From ceb988dc0730f2929fd784b53ec37c9e8d53bb93 Mon Sep 17 00:00:00 2001 From: Chris Beck Date: Sun, 14 Dec 2025 11:22:28 -0700 Subject: [PATCH] abstract the claude module out of the signal-gateway and put interface and ipml in their own crates --- Cargo.lock | 33 +- Cargo.toml | 5 + signal-gateway-app-code/Cargo.toml | 2 +- signal-gateway-app-code/src/lib.rs | 2 +- signal-gateway-assistant/Cargo.toml | 14 + signal-gateway-assistant/claude/Cargo.toml | 26 ++ .../claude/src/lib.rs | 410 ++++++++---------- signal-gateway-assistant/src/assistant.rs | 32 ++ signal-gateway-assistant/src/chat_message.rs | 44 ++ signal-gateway-assistant/src/lib.rs | 14 + signal-gateway-assistant/src/response.rs | 49 +++ .../src}/tools.rs | 8 +- signal-gateway-bin/Cargo.toml | 3 +- signal-gateway-bin/src/main.rs | 16 +- signal-gateway-log-ingest/Cargo.toml | 2 +- signal-gateway/Cargo.toml | 1 + signal-gateway/src/assistant/mod.rs | 118 +++++ signal-gateway/src/assistant/worker.rs | 124 ++++++ signal-gateway/src/claude/mod.rs | 210 --------- signal-gateway/src/gateway/log_handler.rs | 2 +- signal-gateway/src/gateway/mod.rs | 106 +++-- signal-gateway/src/lib.rs | 2 +- signal-gateway/src/prometheus/mod.rs | 2 +- 23 files changed, 736 insertions(+), 489 deletions(-) create mode 100644 signal-gateway-assistant/Cargo.toml create mode 100644 signal-gateway-assistant/claude/Cargo.toml rename signal-gateway/src/claude/worker.rs => signal-gateway-assistant/claude/src/lib.rs (60%) create mode 100644 signal-gateway-assistant/src/assistant.rs create mode 100644 signal-gateway-assistant/src/chat_message.rs create mode 100644 signal-gateway-assistant/src/lib.rs create mode 100644 signal-gateway-assistant/src/response.rs rename {signal-gateway/src/claude => signal-gateway-assistant/src}/tools.rs (88%) create mode 100644 signal-gateway/src/assistant/mod.rs create mode 100644 signal-gateway/src/assistant/worker.rs delete mode 100644 signal-gateway/src/claude/mod.rs diff --git a/Cargo.lock b/Cargo.lock index f12a5d8..6dea536 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1735,6 +1735,7 @@ dependencies = [ "reqwest", "serde", "serde_json", + "signal-gateway-assistant", "thiserror", "tokio", "tokio-util", @@ -1753,12 +1754,41 @@ dependencies = [ "reqwest", "serde", "serde_json", - "signal-gateway", + "signal-gateway-assistant", "tar", "tokio", "tracing", ] +[[package]] +name = "signal-gateway-assistant" +version = "0.1.0" +dependencies = [ + "async-trait", + "chrono", + "serde", + "serde_json", + "tokio-util", +] + +[[package]] +name = "signal-gateway-assistant-claude" +version = "0.1.0" +dependencies = [ + "async-trait", + "chrono", + "conf", + "conf-extra", + "reqwest", + "serde", + "serde_json", + "signal-gateway-assistant", + "thiserror", + "tokio", + "tokio-util", + "tracing", +] + [[package]] name = "signal-gateway-bin" version = "0.1.0" @@ -1774,6 +1804,7 @@ dependencies = [ "serde_json", "signal-gateway", "signal-gateway-app-code", + "signal-gateway-assistant-claude", "signal-gateway-log-ingest", "tokio", "tokio-util", diff --git a/Cargo.toml b/Cargo.toml index 9afc731..26031cb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,6 +4,8 @@ members = [ "prometheus-http-client", "signal-gateway", "signal-gateway-app-code", + "signal-gateway-assistant", + "signal-gateway-assistant/claude", "signal-gateway-bin", "signal-gateway-log-ingest", ] @@ -25,6 +27,9 @@ result_large_err = "allow" [workspace.dependencies] prometheus-http-client = { path = "prometheus-http-client", default-features = false } +signal-gateway = { path = "signal-gateway", default-features = false } +signal-gateway-assistant = { path = "signal-gateway-assistant" } +signal-gateway-assistant-claude = { path = "signal-gateway-assistant/claude" } async-trait = "0.1" chrono = { version = "0.4", default-features = false, features = ["clock", "serde", "std"] } diff --git a/signal-gateway-app-code/Cargo.toml b/signal-gateway-app-code/Cargo.toml index fb37b66..bd778cd 100644 --- a/signal-gateway-app-code/Cargo.toml +++ b/signal-gateway-app-code/Cargo.toml @@ -7,7 +7,7 @@ edition.workspace = true workspace = true [dependencies] -signal-gateway = { path = "../signal-gateway", default-features = false } +signal-gateway-assistant = { workspace = true } async-trait = { workspace = true } flate2 = { workspace = true } diff --git a/signal-gateway-app-code/src/lib.rs b/signal-gateway-app-code/src/lib.rs index cfb1c8a..07a446e 100644 --- a/signal-gateway-app-code/src/lib.rs +++ b/signal-gateway-app-code/src/lib.rs @@ -7,7 +7,7 @@ use async_trait::async_trait; use flate2::read::GzDecoder; use regex::Regex; use serde::Deserialize; -use signal_gateway::claude::{Tool, ToolExecutor, ToolResult}; +use signal_gateway_assistant::{Tool, ToolExecutor, ToolResult}; use std::{ collections::HashMap, error::Error, fmt::Write, future::Future, io::Read, path::PathBuf, pin::Pin, str::FromStr, sync::Arc, diff --git a/signal-gateway-assistant/Cargo.toml b/signal-gateway-assistant/Cargo.toml new file mode 100644 index 0000000..7735fa5 --- /dev/null +++ b/signal-gateway-assistant/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "signal-gateway-assistant" +version = "0.1.0" +edition.workspace = true + +[lints] +workspace = true + +[dependencies] +async-trait = { workspace = true } +chrono = { workspace = true } +serde = { workspace = true } +serde_json = { workspace = true } +tokio-util = { workspace = true } diff --git a/signal-gateway-assistant/claude/Cargo.toml b/signal-gateway-assistant/claude/Cargo.toml new file mode 100644 index 0000000..6c8cd1d --- /dev/null +++ b/signal-gateway-assistant/claude/Cargo.toml @@ -0,0 +1,26 @@ +[package] +name = "signal-gateway-assistant-claude" +version = "0.1.0" +edition.workspace = true + +[lints] +workspace = true + +[features] +default = ["rustls-tls"] +rustls-tls = ["reqwest/rustls-tls"] + +[dependencies] +signal-gateway-assistant = { workspace = true } + +async-trait = { workspace = true } +chrono = { workspace = true } +conf = { workspace = true } +conf-extra = { workspace = true } +reqwest = { workspace = true } +serde = { workspace = true } +serde_json = { workspace = true } +thiserror = { workspace = true } +tokio = { workspace = true } +tokio-util = { workspace = true } +tracing = { workspace = true } diff --git a/signal-gateway/src/claude/worker.rs b/signal-gateway-assistant/claude/src/lib.rs similarity index 60% rename from signal-gateway/src/claude/worker.rs rename to signal-gateway-assistant/claude/src/lib.rs index b852ee9..9311796 100644 --- a/signal-gateway/src/claude/worker.rs +++ b/signal-gateway-assistant/claude/src/lib.rs @@ -1,87 +1,94 @@ -//! Background worker that processes Claude API requests serially. +//! Claude API implementation of the Assistant trait. -use super::{ANTHROPIC_API_VERSION, ClaudeConfig, ClaudeError, Tool, ToolExecutor}; -use crate::message_handler::AdminMessageResponse; -use chrono::{DateTime, Utc}; +use conf::Conf; use serde::{Deserialize, Serialize}; use serde_json::Value; -use std::{path::PathBuf, sync::Weak}; -use tokio::sync::{mpsc, oneshot}; +use signal_gateway_assistant::{Assistant, AssistantResponse, ChatMessage, Tool, ToolExecutor}; +use std::{path::PathBuf, sync::Weak, time::Duration}; +use tokio_util::sync::CancellationToken; use tracing::{error, info, warn}; -/// Sent with inputs to claude that claude is expected to respond to. The sender -/// gives the worker a way to return the results to the caller asynchronously. -pub type ResultSender = oneshot::Sender>; +/// The Anthropic API version header value. +const ANTHROPIC_API_VERSION: &str = "2023-06-01"; -/// Indicates the "role" i.e. the manner in which a particular message was sent -pub enum SentBy { - /// User message directed at the system (commands, etc.) - UserToSystem, - /// User message directed at Claude (prompts) - UserToClaude, - /// Response from Claude - Claude, - /// System-generated message - System, - /// Alert from alertmanager - AlertManager, +/// Configuration for the Claude API integration. +#[derive(Clone, Conf, Debug)] +#[conf(serde)] +pub struct ClaudeConfig { + /// Path to file containing the Claude API key. + #[conf(long, env)] + pub api_key_file: PathBuf, + /// Paths to files containing system prompt components (in order, last one is cached). + #[conf(repeat, long, env)] + pub system_prompt_files: Vec, + /// Claude API URL. + #[conf(long, env, default_value = "https://api.anthropic.com/v1/messages")] + pub claude_api_url: String, + /// Claude model to use. + #[conf(long, env, default_value = "claude-sonnet-4-5-20250929")] + pub claude_model: String, + /// Maximum tokens in the response. + #[conf(long, env, default_value = "1024")] + pub claude_max_tokens: u32, + /// Maximum tool use iterations before giving up. + #[conf(long, env, default_value = "10")] + pub claude_max_iterations: u32, + /// Enable prompt caching (adds cache_control to system prompts). + #[conf(long, env)] + pub prompt_caching: bool, + /// Compaction configuration. + #[conf(flatten, prefix)] + pub compaction: CompactionConfig, } -impl SentBy { - /// Returns the API role: "assistant" for Claude, "user" for everything else. - fn api_role(&self) -> &'static str { - match self { - Self::Claude => "assistant", - _ => "user", - } - } - - /// Returns a prefix to prepend to message text for context. - fn prefix(&self) -> Option<&'static str> { - match self { - Self::UserToSystem => Some("[user to system]"), - Self::UserToClaude => None, // No prefix needed for direct user messages - Self::Claude => None, - Self::System => Some("[system]"), - Self::AlertManager => Some("[alertmanager]"), - } - } +/// Configuration for message buffer compaction. +#[derive(Clone, Conf, Debug)] +#[conf(serde)] +pub struct CompactionConfig { + /// Path to file containing the compaction prompt. + #[conf(long, env)] + pub prompt_file: PathBuf, + /// Model to use for compaction (typically a faster/cheaper model). + #[conf(long, env, default_value = "claude-sonnet-4-5-20250929")] + pub model: String, + /// Maximum tokens for the compaction response. + #[conf(long, env, default_value = "2048")] + pub max_tokens: u32, + /// Trigger compaction when message buffer exceeds this many characters. + #[conf(long, env, default_value = "50000")] + pub trigger_chars: u32, + /// Minimum interval between automatic compactions (not user-requested). + /// If omitted, AI compaction always runs (no rate limiting). + /// If set, compaction is rate-limited and low-priority messages are dropped instead. + #[conf(long, env, value_parser = conf_extra::parse_duration, serde(use_value_parser))] + pub min_interval: Option, } -/// An input to the worker sent through the channel -pub enum Input { - Chat(ChatMessage), - Compact, - Debug, +/// Error type for Claude API operations. +#[derive(Debug, thiserror::Error)] +pub enum ClaudeError { + /// Failed to read API key file. + #[error("failed to read API key file: {0}")] + ApiKeyRead(std::io::Error), + /// Failed to read system prompt file. + #[error("failed to read system prompt file {0:?}: {1}")] + SystemPromptRead(PathBuf, std::io::Error), + /// HTTP request failed. + #[error("HTTP request failed: {0}")] + Request(#[from] reqwest::Error), + /// API returned an error response. + #[error("API error: {0}")] + ApiError(Box), + /// Too many tool use iterations. + #[error("exceeded maximum tool use iterations ({0})")] + TooManyIterations(u32), + /// Tool executor is no longer available. + #[error("tool executor is gone")] + ToolExecutorGone, } -/// A chat message -pub struct ChatMessage { - pub sent_by: SentBy, - pub timestamp: DateTime, - pub text: Box, - /// Present when claude is expected to respond to the message - pub result_sender: Option, -} - -impl ChatMessage { - fn into_content_and_sender(self) -> (MessageContent, Option) { - let ts = self.timestamp.format("%Y-%m-%dT%H:%M:%S%.3fZ"); - let text = match self.sent_by.prefix() { - Some(prefix) => format!("[{}] {} {}", ts, prefix, self.text).into(), - None => format!("[{}] {}", ts, self.text).into(), - }; - let mc = MessageContent { - role: self.sent_by.api_role().into(), - content: vec![ContentBlock::Text { text }], - }; - - (mc, self.result_sender) - } -} - -/// Background worker that processes Claude API requests serially. -pub struct ClaudeWorker { +/// Claude API assistant implementation. +pub struct ClaudeAssistant { config: ClaudeConfig, client: reqwest::Client, api_key: String, @@ -91,41 +98,35 @@ pub struct ClaudeWorker { summary: String, messages: Vec, tool_executor: Weak, - input_rx: mpsc::Receiver, - stop_rx: mpsc::Receiver<()>, /// Timestamp of the last automatic compaction (not user-requested). last_auto_compaction: Option, } -impl ClaudeWorker { - /// Create a new Claude worker. +impl ClaudeAssistant { + /// Create a new Claude assistant from configuration. /// /// Reads the API key, system prompts, and compaction prompt from the configured files. pub fn new( config: ClaudeConfig, tool_executor: Weak, - input_rx: mpsc::Receiver, - stop_rx: mpsc::Receiver<()>, ) -> Result { let api_key = std::fs::read_to_string(&config.api_key_file) .map_err(ClaudeError::ApiKeyRead)? .trim() - .to_owned(); + .to_string(); - let system_prompts: Vec = config + let system_prompts: Result, ClaudeError> = config .system_prompt_files .iter() .map(|path| { std::fs::read_to_string(path) .map_err(|e| ClaudeError::SystemPromptRead(path.clone(), e)) }) - .collect::>()?; + .collect(); + let system_prompts = system_prompts?; - let compaction_prompt = { - let path = &config.compaction.prompt_file; - std::fs::read_to_string(path) - .map_err(|e| ClaudeError::SystemPromptRead(path.clone(), e))? - }; + let compaction_prompt = std::fs::read_to_string(&config.compaction.prompt_file) + .map_err(|e| ClaudeError::SystemPromptRead(config.compaction.prompt_file.clone(), e))?; Ok(Self { config, @@ -134,71 +135,12 @@ impl ClaudeWorker { system_prompts, compaction_prompt, summary: String::new(), - messages: Default::default(), + messages: Vec::new(), tool_executor, - input_rx, - stop_rx, last_auto_compaction: None, }) } - /// Run the worker loop, processing requests serially. - pub async fn run(mut self) { - loop { - tokio::select! { - input = self.input_rx.recv() => { - let Some(input) = input else { - // Channel closed, exit - break; - }; - match input { - Input::Chat(msg) => { - let (mc, maybe_sender) = msg.into_content_and_sender(); - self.messages.push(mc); - - // Check if we need to trigger automatic compaction - let buffer_chars = self.message_buffer_chars(); - if buffer_chars > self.config.compaction.trigger_chars as usize { - // Check if rate limiting is configured - let can_compact = match self.config.compaction.min_interval { - None => true, // No rate limit configured - Some(min_interval) => self.last_auto_compaction - .map(|t| t.elapsed() >= min_interval) - .unwrap_or(true), - }; - if can_compact { - self.handle_compact(true).await; - } else { - // Rate limited - drop oldest messages instead - self.drop_oldest_messages(); - } - } - - if let Some(sender) = maybe_sender { - let result = self.handle_request().await; - // If handle_request was interrupted by stop request, go on to drain the queues - if matches!(result, Err(ClaudeError::StopRequested)) { - self.handle_stop(); - } - // Ignore send errors - the caller may have dropped the receiver - let _ = sender.send(result); - } - }, - Input::Compact => { - self.handle_compact(false).await; - } - Input::Debug => { - self.handle_debug(); - } - } - } - _ = self.stop_rx.recv() => { - self.handle_stop(); - } - } - } - } - /// Calculate total characters in the message buffer. fn message_buffer_chars(&self) -> usize { self.messages @@ -216,6 +158,29 @@ impl ClaudeWorker { .sum() } + /// Check if automatic compaction should be triggered and handle it. + async fn maybe_compact(&mut self) { + let buffer_chars = self.message_buffer_chars(); + if buffer_chars <= self.config.compaction.trigger_chars as usize { + return; + } + + // Check if rate limiting is configured + let can_compact = match self.config.compaction.min_interval { + None => true, + Some(min_interval) => self + .last_auto_compaction + .map(|t| t.elapsed() >= min_interval) + .unwrap_or(true), + }; + + if can_compact { + self.do_compact(true).await; + } else { + self.drop_oldest_messages(); + } + } + /// Drop oldest messages until buffer is under the trigger threshold. fn drop_oldest_messages(&mut self) { let target = self.config.compaction.trigger_chars as usize; @@ -228,7 +193,6 @@ impl ClaudeWorker { let mut current_chars = before_chars; let mut dropped = 0; - // Remove from the front (oldest) until under limit while current_chars > target && !self.messages.is_empty() { let msg_chars = self.messages[0] .content @@ -251,45 +215,12 @@ impl ClaudeWorker { ); } - /// Handle a stop request by draining queues and sending errors to pending requests. - fn handle_stop(&mut self) { - // Drain the stop_rx queue - while self.stop_rx.try_recv().is_ok() {} - - // Drain the input_rx queue and send StopRequested to each prompt request - while let Ok(input) = self.input_rx.try_recv() { - match input { - Input::Chat(msg) => { - let (mc, maybe_sender) = msg.into_content_and_sender(); - self.messages.push(mc); - if let Some(sender) = maybe_sender { - let _ = sender.send(Err(ClaudeError::StopRequested)); - } - } - Input::Compact | Input::Debug => {} - } - } - } - - /// Check if stop has been requested. - fn check_stop(&mut self) -> Result<(), ClaudeError> { - match self.stop_rx.try_recv() { - Ok(()) => Err(ClaudeError::StopRequested), - Err(mpsc::error::TryRecvError::Empty) => Ok(()), - Err(mpsc::error::TryRecvError::Disconnected) => Err(ClaudeError::StopRequested), - } - } - /// Perform compaction by summarizing messages and storing the result. - /// - /// If `is_automatic` is true, this was triggered by buffer size threshold - /// and the last_auto_compaction timestamp will be updated. - async fn handle_compact(&mut self, is_automatic: bool) { + async fn do_compact(&mut self, is_automatic: bool) { if self.messages.is_empty() { return; } - // Update timestamp for automatic compactions (before the API call) if is_automatic { self.last_auto_compaction = Some(std::time::Instant::now()); } @@ -312,7 +243,6 @@ impl ClaudeWorker { if !self.summary.is_empty() { system.push(SystemContent::text(&self.summary)); } - // Mark the last one as cached if prompt caching is enabled if self.config.prompt_caching && let Some(last) = system.last_mut() { @@ -341,13 +271,12 @@ impl ClaudeWorker { Ok(response) if response.status().is_success() => { match response.json::().await { Ok(parsed) => { - // Extract text from response let summary_text: String = parsed .content .into_iter() .filter_map(|block| { if let ContentBlock::Text { text } = block { - Some(text) + Some(String::from(text)) } else { None } @@ -355,7 +284,6 @@ impl ClaudeWorker { .collect::>() .join("\n"); - // Wrap in XML tags and store self.summary = format!("\n{}\n", summary_text); self.messages.clear(); @@ -382,25 +310,18 @@ impl ClaudeWorker { } } - /// Log the message buffer for debugging. - fn handle_debug(&self) { - if self.summary.is_empty() { - info!("Claude summary: (empty)"); - } else { - info!("Claude summary:\n{}", self.summary); - } - info!("Claude message buffer:\n{:#?}", self.messages); - } - - /// Handle a single request to the Claude API. - async fn handle_request(&mut self) -> Result { + /// Handle a single request to the Claude API with tool use loop. + /// + /// Returns `Ok(None)` if the operation was cancelled. + async fn handle_request( + &mut self, + cancel: &CancellationToken, + ) -> Result, ClaudeError> { let max_iterations = self.config.claude_max_iterations; - // Get tools from the executor if still alive let executor = self.tool_executor.upgrade(); let tools = executor.as_ref().map(|te| te.tools()).unwrap_or_default(); - // Collect attachments from tool results across all iterations let mut attachments: Vec = Vec::new(); if let Some(last) = self.messages.last() @@ -409,7 +330,7 @@ impl ClaudeWorker { info!("Claude request: {}", text); } - // Build system content blocks: system prompts + summary (if any), caching only the last one + // Build system content blocks let mut system: Vec = self .system_prompts .iter() @@ -418,7 +339,6 @@ impl ClaudeWorker { if !self.summary.is_empty() { system.push(SystemContent::text(&self.summary)); } - // Mark the last one as cached if prompt caching is enabled if self.config.prompt_caching && let Some(last) = system.last_mut() { @@ -426,8 +346,9 @@ impl ClaudeWorker { } for iteration in 0..max_iterations { - // Check for stop before making API call - self.check_stop()?; + if cancel.is_cancelled() { + return Ok(None); + } let request_body = MessagesRequest { model: &self.config.claude_model, @@ -458,28 +379,24 @@ impl ClaudeWorker { response.stop_reason, response.content ); - // Check if we need to handle tool use if response.stop_reason.as_ref() == "tool_use" { - // Try to get a strong reference to the executor let Some(executor) = self.tool_executor.upgrade() else { return Err(ClaudeError::ToolExecutorGone); }; - // Add assistant's response to messages self.messages .push(MessageContent::assistant(response.content.clone())); - // Execute each tool use and collect results for block in &response.content { if let ContentBlock::ToolUse { id, name, input } = block { - // Check for stop before each tool use - self.check_stop()?; + if cancel.is_cancelled() { + return Ok(None); + } info!("Claude tool use: {}({})", name, input); let (result_text, is_error) = match executor.execute(name, input).await { Ok(tool_result) => { info!("Tool result: {}", tool_result.text); - // Collect any attachments from the tool result attachments.extend(tool_result.attachments); (tool_result.text, false) } @@ -496,7 +413,6 @@ impl ClaudeWorker { } } - // Continue the loop to get Claude's next response info!("Tool use iteration {}, continuing...", iteration + 1); continue; } @@ -507,7 +423,7 @@ impl ClaudeWorker { .into_iter() .filter_map(|block| { if let ContentBlock::Text { text, .. } = block { - Some(text) + Some(text.into_string()) } else { None } @@ -516,34 +432,88 @@ impl ClaudeWorker { .join("\n"); info!("Claude final result: {}", text); - return Ok(AdminMessageResponse::new(text).with_attachments(attachments)); + return Ok(Some(AssistantResponse::with_attachments(text, attachments))); } Err(ClaudeError::TooManyIterations(max_iterations)) } } +#[async_trait::async_trait] +impl Assistant for ClaudeAssistant { + async fn record_message(&mut self, message: ChatMessage) { + let mc = message_to_content(message); + self.messages.push(mc); + self.maybe_compact().await; + } + + async fn prompt( + &mut self, + message: ChatMessage, + cancel: CancellationToken, + ) -> Result, Box> { + let mc = message_to_content(message); + self.messages.push(mc); + self.maybe_compact().await; + + if cancel.is_cancelled() { + return Ok(None); + } + + Ok(self.handle_request(&cancel).await?) + } + + async fn compact(&mut self) { + // Manual compaction always runs (is_automatic = false) + self.do_compact(false).await; + } + + fn debug_log(&mut self) { + if self.summary.is_empty() { + info!("Claude summary: (empty)"); + } else { + info!("Claude summary:\n{}", self.summary); + } + info!("Claude message buffer:\n{:#?}", self.messages); + } +} + +/// Convert a ChatMessage to a MessageContent for the API. +fn message_to_content(msg: ChatMessage) -> MessageContent { + let ts = msg.timestamp.format("%Y-%m-%dT%H:%M:%S%.3fZ"); + let text = match msg.sent_by.prefix() { + Some(prefix) => format!("[{}] {} {}", ts, prefix, msg.text), + None => format!("[{}] {}", ts, msg.text), + }; + MessageContent { + role: msg.sent_by.api_role().into(), + content: vec![ContentBlock::Text { text: text.into() }], + } +} + /// Estimate the serialized size of a JSON Value without allocating. fn estimate_json_size(value: &Value) -> usize { match value { - Value::Null => 4, // "null" - Value::Bool(true) => 4, // "true" - Value::Bool(false) => 5, // "false" - Value::Number(n) => n.to_string().len(), // Numbers are small, ok to alloc - Value::String(s) => s.len() + 2, // quotes + Value::Null => 4, + Value::Bool(true) => 4, + Value::Bool(false) => 5, + Value::Number(n) => n.to_string().len(), + Value::String(s) => s.len() + 2, Value::Array(arr) => { 2 + arr.iter().map(estimate_json_size).sum::() + arr.len().saturating_sub(1) } Value::Object(obj) => { 2 + obj .iter() - .map(|(k, v)| k.len() + 3 + estimate_json_size(v)) // "key":value + .map(|(k, v)| k.len() + 3 + estimate_json_size(v)) .sum::() - + obj.len().saturating_sub(1) // commas + + obj.len().saturating_sub(1) } } } +// ---- API Types ---- + /// Request body for the Claude Messages API. #[derive(Serialize)] struct MessagesRequest<'a> { diff --git a/signal-gateway-assistant/src/assistant.rs b/signal-gateway-assistant/src/assistant.rs new file mode 100644 index 0000000..66e0aca --- /dev/null +++ b/signal-gateway-assistant/src/assistant.rs @@ -0,0 +1,32 @@ +use crate::{AssistantResponse, ChatMessage}; +use std::error::Error; +use tokio_util::sync::CancellationToken; + +/// Assistant trait which handles chat messages and commands. +/// +/// Implementations handle the actual LLM API calls, message history management, +/// and tool execution. +#[async_trait::async_trait] +pub trait Assistant: Send { + /// Record a chat message in the assistant's history without expecting a response. + async fn record_message(&mut self, message: ChatMessage); + + /// Process a chat message and generate a response. + /// + /// The cancellation token can be used to interrupt long-running operations. + /// Returns `Ok(None)` if the operation was cancelled. + async fn prompt( + &mut self, + message: ChatMessage, + cancel: CancellationToken, + ) -> Result, Box>; + + /// Compact the assistant's message history (e.g., by summarizing). + /// + /// This is called when the user explicitly requests compaction. + /// Automatic compaction is an internal implementation detail. + async fn compact(&mut self); + + /// Log the assistant's current state for debugging. + fn debug_log(&mut self); +} diff --git a/signal-gateway-assistant/src/chat_message.rs b/signal-gateway-assistant/src/chat_message.rs new file mode 100644 index 0000000..b485fd5 --- /dev/null +++ b/signal-gateway-assistant/src/chat_message.rs @@ -0,0 +1,44 @@ +use chrono::{DateTime, Utc}; + +/// Indicates the "role" i.e. the manner in which a particular message was sent +#[non_exhaustive] +pub enum SentBy { + /// User message directed at the system (commands, etc.) + UserToSystem, + /// User message directed at assistant (prompts) + UserToAssistant, + /// Response from assistant + Assistant, + /// System-generated message + System, + /// Alert from alertmanager + AlertManager, +} + +impl SentBy { + /// Returns the API role: "assistant" for Assistant, "user" for everything else. + pub fn api_role(&self) -> &'static str { + match self { + Self::Assistant => "assistant", + _ => "user", + } + } + + /// Returns a prefix to prepend to message text for context. + pub fn prefix(&self) -> Option<&'static str> { + match self { + Self::UserToSystem => Some("[user to system]"), + Self::UserToAssistant => None, // No prefix needed for direct user messages + Self::Assistant => None, + Self::System => Some("[system]"), + Self::AlertManager => Some("[alertmanager]"), + } + } +} + +/// A chat message +pub struct ChatMessage { + pub sent_by: SentBy, + pub timestamp: DateTime, + pub text: Box, +} diff --git a/signal-gateway-assistant/src/lib.rs b/signal-gateway-assistant/src/lib.rs new file mode 100644 index 0000000..17d0a92 --- /dev/null +++ b/signal-gateway-assistant/src/lib.rs @@ -0,0 +1,14 @@ +//! Assistant API used by signal-gateway. +//! +//! This crate provides abstract types for LLM assistant interactions, +//! making it easy to swap in different LLM implementations. + +mod assistant; +mod chat_message; +mod response; +mod tools; + +pub use assistant::Assistant; +pub use chat_message::{ChatMessage, SentBy}; +pub use response::AssistantResponse; +pub use tools::{Tool, ToolExecutor, ToolResult}; diff --git a/signal-gateway-assistant/src/response.rs b/signal-gateway-assistant/src/response.rs new file mode 100644 index 0000000..e4cbe4f --- /dev/null +++ b/signal-gateway-assistant/src/response.rs @@ -0,0 +1,49 @@ +//! Response type for assistant interactions. + +use std::path::PathBuf; + +/// Response from an assistant interaction. +#[derive(Clone, Debug, Default)] +#[non_exhaustive] +pub struct AssistantResponse { + /// Text response from the assistant. + pub text: String, + /// Optional file attachments generated during the interaction. + pub attachments: Vec, +} + +impl AssistantResponse { + /// Create a new response with just text. + pub fn new(text: impl Into) -> Self { + Self { + text: text.into(), + attachments: Vec::new(), + } + } + + /// Create a response with text and attachments. + pub fn with_attachments(text: impl Into, attachments: Vec) -> Self { + Self { + text: text.into(), + attachments, + } + } + + /// Add attachments to this response. + pub fn add_attachments(mut self, attachments: Vec) -> Self { + self.attachments.extend(attachments); + self + } +} + +impl From for AssistantResponse { + fn from(text: String) -> Self { + Self::new(text) + } +} + +impl From<&str> for AssistantResponse { + fn from(text: &str) -> Self { + Self::new(text) + } +} diff --git a/signal-gateway/src/claude/tools.rs b/signal-gateway-assistant/src/tools.rs similarity index 88% rename from signal-gateway/src/claude/tools.rs rename to signal-gateway-assistant/src/tools.rs index 3dbcb09..e56f299 100644 --- a/signal-gateway/src/claude/tools.rs +++ b/signal-gateway-assistant/src/tools.rs @@ -1,4 +1,4 @@ -//! Tool definitions and executor trait for the Claude API. +//! Tool definitions and executor trait for the LLM API. use async_trait::async_trait; use serde::Serialize; @@ -8,7 +8,7 @@ use std::path::PathBuf; /// Result of executing a tool. #[derive(Clone, Debug, Default)] pub struct ToolResult { - /// Text result to return to Claude. + /// Text result to return to LLM. pub text: String, /// Optional file attachments generated by the tool. pub attachments: Vec, @@ -56,11 +56,11 @@ pub trait ToolExecutor: Send + Sync { } /// Execute a tool by name with the given input arguments. - /// Returns the result to be sent back to Claude, potentially with attachments. + /// Returns the result to be sent back to LLM, potentially with attachments. async fn execute(&self, name: &str, input: &Value) -> Result; } -/// A tool definition for the Claude API. +/// A tool definition for the LLM API. #[derive(Clone, Debug, Serialize)] pub struct Tool { /// The name of the tool. diff --git a/signal-gateway-bin/Cargo.toml b/signal-gateway-bin/Cargo.toml index 9c6f912..9fbcde4 100644 --- a/signal-gateway-bin/Cargo.toml +++ b/signal-gateway-bin/Cargo.toml @@ -16,7 +16,8 @@ plot = ["signal-gateway/plot"] rustls-tls = ["signal-gateway/rustls-tls"] [dependencies] -signal-gateway = { path = "../signal-gateway" } +signal-gateway = { workspace = true } +signal-gateway-assistant-claude = { workspace = true } signal-gateway-app-code = { path = "../signal-gateway-app-code" } signal-gateway-log-ingest = { path = "../signal-gateway-log-ingest" } diff --git a/signal-gateway-bin/src/main.rs b/signal-gateway-bin/src/main.rs index 9f3310c..3ee8dbb 100644 --- a/signal-gateway-bin/src/main.rs +++ b/signal-gateway-bin/src/main.rs @@ -7,6 +7,7 @@ use hyper::service::service_fn; use hyper_util::{rt::TokioIo, server::conn::auto}; use signal_gateway::{CommandRouter, Gateway, GatewayConfig, Handling}; use signal_gateway_app_code::AppCodeTools; +use signal_gateway_assistant_claude::{ClaudeAssistant, ClaudeConfig}; use std::{env, fs, net::SocketAddr, path::PathBuf, sync::Arc, time::Duration}; use tokio::net::TcpListener; use tokio_util::sync::CancellationToken; @@ -46,6 +47,9 @@ pub struct Config { /// Application source code configurations for Claude tools. #[conf(long, env, value_parser = serde_json::from_str, default, default_help_str = "[]")] app_code: Vec, + /// Claude API configuration for AI-powered responses. + #[conf(flatten, prefix)] + claude: Option, #[conf(flatten, serde(flatten))] gateway: GatewayConfig, } @@ -121,7 +125,7 @@ async fn main() -> Result<(), Box> { router_builder = router_builder.route("/", Handling::GatewayCommand); // Add Claude as default handler if configured - if config.gateway.claude.is_some() { + if config.claude.is_some() { router_builder = router_builder.route("", Handling::Claude); } @@ -153,6 +157,16 @@ async fn main() -> Result<(), Box> { gateway_builder = gateway_builder.with_tools(tools); } + // Add Claude assistant if configured + if let Some(claude_config) = config.claude { + gateway_builder = gateway_builder.with_assistant(move |tool_executor| { + Box::new( + ClaudeAssistant::new(claude_config, tool_executor) + .expect("Failed to initialize Claude assistant"), + ) + }); + } + let gateway = gateway_builder.build().await; let listener = TcpListener::bind(config.http_listen_addr).await.unwrap(); diff --git a/signal-gateway-log-ingest/Cargo.toml b/signal-gateway-log-ingest/Cargo.toml index 0e6184e..abf0a14 100644 --- a/signal-gateway-log-ingest/Cargo.toml +++ b/signal-gateway-log-ingest/Cargo.toml @@ -7,7 +7,7 @@ edition.workspace = true workspace = true [dependencies] -signal-gateway = { path = "../signal-gateway" } +signal-gateway = { workspace = true } chrono = { workspace = true } conf = { workspace = true } diff --git a/signal-gateway/Cargo.toml b/signal-gateway/Cargo.toml index 9e36e9c..9785ab1 100644 --- a/signal-gateway/Cargo.toml +++ b/signal-gateway/Cargo.toml @@ -13,6 +13,7 @@ rustls-tls = ["prometheus-http-client/rustls-tls", "reqwest/rustls-tls"] [dependencies] prometheus-http-client = { workspace = true, default-features = false } +signal-gateway-assistant = { workspace = true } async-trait = { workspace = true } chrono = { workspace = true } diff --git a/signal-gateway/src/assistant/mod.rs b/signal-gateway/src/assistant/mod.rs new file mode 100644 index 0000000..30deab3 --- /dev/null +++ b/signal-gateway/src/assistant/mod.rs @@ -0,0 +1,118 @@ +//! Assistant integration for AI-powered responses. +//! +//! This module provides a channel-based wrapper around an `Assistant` implementation, +//! handling message queuing and background processing. + +mod worker; + +// Re-export the assistant types that callers need +pub use signal_gateway_assistant::{ + Assistant, AssistantResponse, ChatMessage, SentBy, Tool, ToolExecutor, ToolResult, +}; + +use crate::message_handler::AdminMessageResponse; +use chrono::{DateTime, Utc}; +use tokio::sync::{mpsc, oneshot}; +use worker::{AssistantWorker, Input}; + +/// Size of the request queue for the assistant worker. +const REQUEST_QUEUE_SIZE: usize = 16; + +/// Error type for assistant operations at the gateway level. +#[derive(Debug, thiserror::Error)] +pub enum AssistantError { + /// Request queue is full. + #[error("request queue is full")] + QueueFull, + /// Worker has shut down. + #[error("worker has shut down")] + WorkerGone, + /// Request was cancelled. + #[error("request cancelled")] + Cancelled, +} + +/// Assistant agent that processes requests via a background worker. +/// +/// Requests are processed serially by a background worker to prevent +/// concurrent API calls. +pub struct AssistantAgent { + input_tx: mpsc::Sender, + stop_tx: mpsc::Sender<()>, + #[allow(dead_code)] + worker_handle: tokio::task::JoinHandle<()>, +} + +impl AssistantAgent { + /// Create a new assistant agent with the given assistant implementation. + /// + /// Spawns a background worker task that processes requests serially. + pub fn new(assistant: Box) -> Self { + let (input_tx, input_rx) = mpsc::channel(REQUEST_QUEUE_SIZE); + let (stop_tx, stop_rx) = mpsc::channel(REQUEST_QUEUE_SIZE); + + let worker = AssistantWorker::new(assistant, input_rx, stop_rx); + + let worker_handle = tokio::spawn(async move { + worker.run().await; + tracing::info!("Assistant worker task exited"); + }); + + Self { + input_tx, + stop_tx, + worker_handle, + } + } + + /// Send a prompt and wait for a response. + /// + /// Returns `QueueFull` error if the request queue is full. + pub async fn request( + &self, + prompt: &str, + ts_ms: u64, + ) -> Result { + let (result_tx, result_rx) = oneshot::channel(); + + let timestamp = DateTime::from_timestamp_millis(ts_ms as i64).unwrap_or_else(Utc::now); + let msg = ChatMessage { + sent_by: SentBy::UserToAssistant, + text: prompt.into(), + timestamp, + }; + + self.input_tx + .try_send(Input::Prompt(msg, result_tx)) + .map_err(|_| AssistantError::QueueFull)?; + + result_rx.await.map_err(|_| AssistantError::WorkerGone)? + } + + /// Record a message in the assistant's history without expecting a response. + pub fn record_message(&self, sent_by: SentBy, text: &str, ts_ms: u64) { + let timestamp = DateTime::from_timestamp_millis(ts_ms as i64).unwrap_or_else(Utc::now); + let msg = ChatMessage { + sent_by, + text: text.into(), + timestamp, + }; + + let _ = self.input_tx.try_send(Input::Record(msg)); + } + + /// Request the worker to compact the assistant's message history. + pub fn request_compaction(&self) { + let _ = self.input_tx.try_send(Input::Compact); + } + + /// Request the worker to log its state for debugging. + pub fn request_debug(&self) { + let _ = self.input_tx.try_send(Input::Debug); + } + + /// Request the worker to stop processing and cancel pending requests. + pub fn request_stop(&self) { + let _ = self.stop_tx.try_send(()); + } +} diff --git a/signal-gateway/src/assistant/worker.rs b/signal-gateway/src/assistant/worker.rs new file mode 100644 index 0000000..90ee6df --- /dev/null +++ b/signal-gateway/src/assistant/worker.rs @@ -0,0 +1,124 @@ +//! Background worker that processes assistant requests serially. + +use super::AssistantError; +use crate::message_handler::AdminMessageResponse; +use signal_gateway_assistant::{Assistant, AssistantResponse, ChatMessage}; +use tokio::sync::{mpsc, oneshot}; +use tokio_util::sync::CancellationToken; +use tracing::error; + +/// Result sender for prompt requests. +pub type ResultSender = oneshot::Sender>; + +/// Input messages for the worker. +pub enum Input { + /// A prompt that expects a response. + Prompt(ChatMessage, ResultSender), + /// A message to record without expecting a response. + Record(ChatMessage), + /// Request to compact the message history. + Compact, + /// Request to log debug info. + Debug, +} + +/// Background worker that processes assistant requests serially. +pub struct AssistantWorker { + assistant: Box, + input_rx: mpsc::Receiver, + stop_rx: mpsc::Receiver<()>, + cancel_token: CancellationToken, +} + +impl AssistantWorker { + /// Create a new assistant worker. + pub fn new( + assistant: Box, + input_rx: mpsc::Receiver, + stop_rx: mpsc::Receiver<()>, + ) -> Self { + Self { + assistant, + input_rx, + stop_rx, + cancel_token: CancellationToken::new(), + } + } + + /// Run the worker loop, processing requests serially. + pub async fn run(mut self) { + loop { + tokio::select! { + input = self.input_rx.recv() => { + let Some(input) = input else { + break; // Channel closed + }; + self.handle_input(input).await; + } + _ = self.stop_rx.recv() => { + self.handle_stop().await; + } + } + } + } + + async fn handle_input(&mut self, input: Input) { + match input { + Input::Prompt(msg, sender) => { + // Reset the cancel token for each new request + self.cancel_token = CancellationToken::new(); + + let result = self.assistant.prompt(msg, self.cancel_token.clone()).await; + + let response = match result { + Ok(Some(resp)) => Ok(assistant_response_to_admin(resp)), + Ok(None) => Err(AssistantError::Cancelled), + Err(e) => { + error!("Assistant error: {}", e); + // Return the error message as the response text + Ok(AdminMessageResponse::new(format!("Error: {}", e))) + } + }; + + let _ = sender.send(response); + } + Input::Record(msg) => { + self.assistant.record_message(msg).await; + } + Input::Compact => { + self.assistant.compact().await; + } + Input::Debug => { + self.assistant.debug_log(); + } + } + } + + async fn handle_stop(&mut self) { + // Cancel any in-progress request + self.cancel_token.cancel(); + + // Drain remaining stop signals + while self.stop_rx.try_recv().is_ok() {} + + // Drain pending inputs, recording messages but cancelling prompts + while let Ok(input) = self.input_rx.try_recv() { + match input { + Input::Prompt(msg, sender) => { + // Record the message even though we're cancelling + self.assistant.record_message(msg).await; + let _ = sender.send(Err(AssistantError::Cancelled)); + } + Input::Record(msg) => { + self.assistant.record_message(msg).await; + } + Input::Compact | Input::Debug => {} + } + } + } +} + +/// Convert an AssistantResponse to an AdminMessageResponse. +fn assistant_response_to_admin(resp: AssistantResponse) -> AdminMessageResponse { + AdminMessageResponse::new(resp.text).with_attachments(resp.attachments) +} diff --git a/signal-gateway/src/claude/mod.rs b/signal-gateway/src/claude/mod.rs deleted file mode 100644 index 5c0d698..0000000 --- a/signal-gateway/src/claude/mod.rs +++ /dev/null @@ -1,210 +0,0 @@ -//! Claude API integration for AI-powered responses with tool use support. - -mod tools; -mod worker; - -pub use tools::{Tool, ToolExecutor, ToolResult}; -pub use worker::SentBy; - -use crate::message_handler::AdminMessageResponse; -use chrono::{DateTime, Utc}; -use conf::Conf; -use std::{path::PathBuf, sync::Weak, time::Duration}; -use tokio::sync::{mpsc, oneshot}; -use worker::{ChatMessage, ClaudeWorker, Input}; - -/// The Anthropic API version header value. This is a stable version identifier, -/// not a date indicating when the API was released. Anthropic adds new features -/// through separate `anthropic-beta` headers rather than bumping this version. -/// The official Anthropic SDKs use this same value. -const ANTHROPIC_API_VERSION: &str = "2023-06-01"; - -/// Size of the request queue for the Claude worker. -const REQUEST_QUEUE_SIZE: usize = 16; - -/// Configuration for the Claude API integration. -#[derive(Clone, Conf, Debug)] -#[conf(serde)] -pub struct ClaudeConfig { - /// Path to file containing the Claude API key. - #[conf(long, env)] - pub api_key_file: PathBuf, - /// Paths to files containing system prompt components (in order, last one is cached). - #[conf(repeat, long, env)] - pub system_prompt_files: Vec, - /// Claude API URL. - #[conf(long, env, default_value = "https://api.anthropic.com/v1/messages")] - pub claude_api_url: String, - /// Claude model to use. - #[conf(long, env, default_value = "claude-sonnet-4-5-20250929")] - pub claude_model: String, - /// Maximum tokens in the response. - #[conf(long, env, default_value = "1024")] - pub claude_max_tokens: u32, - /// Maximum tool use iterations before giving up. - #[conf(long, env, default_value = "10")] - pub claude_max_iterations: u32, - /// Enable prompt caching (adds cache_control to system prompts). - #[conf(long, env)] - pub prompt_caching: bool, - /// Compaction configuration. - #[conf(flatten, prefix)] - pub compaction: CompactionConfig, -} - -/// Configuration for message buffer compaction. -#[derive(Clone, Conf, Debug)] -#[conf(serde)] -pub struct CompactionConfig { - /// Path to file containing the compaction prompt. - #[conf(long, env)] - pub prompt_file: PathBuf, - /// Model to use for compaction (typically a faster/cheaper model). - #[conf(long, env, default_value = "claude-sonnet-4-5-20250929")] - pub model: String, - /// Maximum tokens for the compaction response. - #[conf(long, env, default_value = "2048")] - pub max_tokens: u32, - /// Trigger compaction when message buffer exceeds this many characters. - #[conf(long, env, default_value = "50000")] - pub trigger_chars: u32, - /// Minimum interval between automatic compactions (not user-requested). - /// If omitted, AI compaction always runs (no rate limiting). - /// If set, compaction is rate-limited and low-priority messages are dropped instead. - #[conf(long, env, value_parser = conf_extra::parse_duration, serde(use_value_parser))] - pub min_interval: Option, -} - -/// Error type for Claude API operations. -#[derive(Debug, thiserror::Error)] -pub enum ClaudeError { - /// Failed to read API key file. - #[error("failed to read API key file: {0}")] - ApiKeyRead(std::io::Error), - /// Failed to read system prompt file. - #[error("failed to read system prompt file {0:?}: {1}")] - SystemPromptRead(PathBuf, std::io::Error), - /// HTTP request failed. - #[error("HTTP request failed: {0}")] - Request(#[from] reqwest::Error), - /// API returned an error response. - #[error("API error: {0}")] - ApiError(Box), - /// Too many tool use iterations. - #[error("exceeded maximum tool use iterations ({0})")] - TooManyIterations(u32), - /// Tool executor is no longer available. - #[error("tool executor is gone")] - ToolExecutorGone, - /// Request queue is full. - #[error("request queue is full")] - QueueFull, - /// Worker has shut down. - #[error("worker has shut down")] - WorkerGone, - /// Stop was requested. - #[error("stop requested")] - StopRequested, -} - -/// Claude API client. -/// -/// Requests are processed serially by a background worker to prevent -/// concurrent API calls. -pub struct ClaudeAgent { - input_tx: mpsc::Sender, - stop_tx: mpsc::Sender<()>, - #[allow(dead_code)] - worker_handle: tokio::task::JoinHandle<()>, -} - -impl ClaudeAgent { - /// Create a new Claude API client from configuration. - /// - /// Reads the API key and system prompt from the configured files. - /// The tool executor is held as a weak reference, so it will not prevent - /// the executor from being dropped. If the executor is dropped during a - /// request, tool use will fail with `ToolExecutorGone`. - /// - /// Spawns a background worker task that processes requests serially. - pub fn new( - config: ClaudeConfig, - tool_executor: Weak, - ) -> Result { - let (input_tx, input_rx) = mpsc::channel(REQUEST_QUEUE_SIZE); - let (stop_tx, stop_rx) = mpsc::channel(REQUEST_QUEUE_SIZE); - - let worker = ClaudeWorker::new(config, tool_executor, input_rx, stop_rx)?; - - let worker_handle = tokio::spawn(async move { - worker.run().await; - tracing::info!("Claude worker task exited"); - }); - - Ok(Self { - input_tx, - stop_tx, - worker_handle, - }) - } - - /// Send a prompt to the Claude API, execute tools etc., and return the response. - /// - /// Returns a future that resolves when the request is complete. - /// Returns `QueueFull` error if the request queue is full. - pub async fn request( - &self, - prompt: &str, - ts_ms: u64, - ) -> Result { - let (result_tx, result_rx) = oneshot::channel(); - - let timestamp = DateTime::from_timestamp_millis(ts_ms as i64).unwrap_or_else(Utc::now); - let msg = ChatMessage { - sent_by: SentBy::UserToClaude, - text: prompt.into(), - timestamp, - result_sender: Some(result_tx), - }; - - self.input_tx - .try_send(Input::Chat(msg)) - .map_err(|_| ClaudeError::QueueFull)?; - - // result_rx.await has type Result, RecvError> - // The outer Result is for channel errors, the inner is the actual response - result_rx.await.map_err(|_| ClaudeError::WorkerGone)? - } - - /// Record a message into claude's chat log that claude is not expected to respond to. - pub fn record_message(&self, sent_by: SentBy, text: &str, ts_ms: u64) { - let timestamp = DateTime::from_timestamp_millis(ts_ms as i64).unwrap_or_else(Utc::now); - let msg = ChatMessage { - sent_by, - text: text.into(), - timestamp, - result_sender: None, - }; - - let _ = self.input_tx.try_send(Input::Chat(msg)); - } - - /// Request the worker to perform compaction (summarizing their recent message history and discarding it) - pub fn request_compaction(&self) { - let _ = self.input_tx.try_send(Input::Compact); - } - - /// Request the worker to log its message buffer for debugging. - pub fn request_debug(&self) { - let _ = self.input_tx.try_send(Input::Debug); - } - - /// Request the worker to interrupt any prompts it is responding to, discard any queued messages, - /// and get ready to accept different prompts. - /// - /// This will cause the current request (if any) to be interrupted at the - /// next opportunity, and all pending requests to receive `StopRequested` errors. - pub fn request_stop(&self) { - let _ = self.stop_tx.try_send(()); - } -} diff --git a/signal-gateway/src/gateway/log_handler.rs b/signal-gateway/src/gateway/log_handler.rs index 59d93f0..5bb3e5a 100644 --- a/signal-gateway/src/gateway/log_handler.rs +++ b/signal-gateway/src/gateway/log_handler.rs @@ -4,7 +4,7 @@ use super::{ route::{Destination, Limit, Route}, }; use crate::{ - claude::{Tool, ToolExecutor, ToolResult}, + assistant::{Tool, ToolExecutor, ToolResult}, concurrent_map::LazyMap, log_format::LogFormatConfig, log_message::{LogFilter, LogMessage, Origin}, diff --git a/signal-gateway/src/gateway/mod.rs b/signal-gateway/src/gateway/mod.rs index a6e1d94..b1cabe3 100644 --- a/signal-gateway/src/gateway/mod.rs +++ b/signal-gateway/src/gateway/mod.rs @@ -4,7 +4,7 @@ use crate::signal_jsonrpc::connect_ipc; use crate::{ alertmanager::AlertPost, - claude::{ClaudeAgent, ClaudeConfig, SentBy, Tool, ToolExecutor, ToolResult}, + assistant::{AssistantAgent, SentBy, Tool, ToolExecutor, ToolResult}, log_message::{LogMessage, Origin}, message_handler::{AdminMessage, AdminMessageResponse, Context, MessageHandlerResult}, prometheus::{Prometheus, PrometheusConfig}, @@ -101,9 +101,6 @@ pub struct GatewayConfig { /// Log handler configuration for processing log messages. #[conf(flatten, prefix)] pub log_handler: LogHandlerConfig, - /// Claude API configuration for AI-powered responses. - #[conf(flatten, prefix)] - pub claude: Option, } /// Wrapper for parsing gateway commands @@ -255,19 +252,23 @@ pub struct Gateway { log_handler: LogHandler, /// Command router for dispatching admin messages. command_router: CommandRouter, - /// Claude API client for AI-powered responses. - /// Initialized after Arc creation so it can hold a weak reference back to Gateway. - claude: OnceLock>, + /// Assistant agent for AI-powered responses. + assistant: OnceLock, /// Additional tool executors added via the builder. extra_tool_executors: Vec>, } +/// Type alias for the assistant factory function. +pub type AssistantFactory = + Box) -> Box + Send>; + /// Builder for creating a [`Gateway`] with optional additional tool executors. pub struct GatewayBuilder { config: GatewayConfig, token: Option, command_router: Option, extra_tool_executors: Vec>, + assistant_factory: Option, } impl GatewayBuilder { @@ -278,6 +279,7 @@ impl GatewayBuilder { token: None, command_router: None, extra_tool_executors: Vec::new(), + assistant_factory: None, } } @@ -295,13 +297,25 @@ impl GatewayBuilder { /// Add an additional tool executor to the gateway. /// - /// Tools from these executors will be available to Claude alongside + /// Tools from these executors will be available to the assistant alongside /// the built-in tools (log handler, prometheus, etc.). pub fn with_tools(mut self, executor: Arc) -> Self { self.extra_tool_executors.push(executor); self } + /// Set the assistant factory for AI-powered responses. + /// + /// The factory receives a weak reference to the gateway (as a ToolExecutor) + /// and returns the assistant implementation. + pub fn with_assistant(mut self, factory: F) -> Self + where + F: FnOnce(Weak) -> Box + Send + 'static, + { + self.assistant_factory = Some(Box::new(factory)); + self + } + /// Build the gateway. pub async fn build(self) -> Arc { Gateway::new_internal( @@ -309,6 +323,7 @@ impl GatewayBuilder { self.token.unwrap_or_default(), self.command_router.unwrap_or_default(), self.extra_tool_executors, + self.assistant_factory, ) .await } @@ -326,6 +341,7 @@ impl Gateway { token: CancellationToken, command_router: CommandRouter, extra_tool_executors: Vec>, + assistant_factory: Option, ) -> Arc { let (signal_alert_mq_tx, signal_alert_mq_rx) = unbounded_channel(); @@ -338,8 +354,6 @@ impl Gateway { let log_handler = LogHandler::new(config.log_handler.clone(), signal_alert_mq_tx.clone()); - let claude_config = config.claude.clone(); - let gateway = Arc::new(Self { config, signal_alert_mq_tx, @@ -348,18 +362,18 @@ impl Gateway { prometheus, log_handler, command_router, - claude: OnceLock::new(), + assistant: OnceLock::new(), extra_tool_executors, }); - // Initialize Claude with a weak reference back to the gateway - if let Some(cc) = claude_config { - let claude = ClaudeAgent::new(cc, Arc::downgrade(&gateway) as Weak) - .expect("Invalid claude config"); + // Initialize the assistant agent using the factory if one was provided + if let Some(factory) = assistant_factory { + let assistant = factory(Arc::downgrade(&gateway) as Weak); + let agent = AssistantAgent::new(assistant); gateway - .claude - .set(Box::new(claude)) - .unwrap_or_else(|_| panic!("claude OnceLock was already set")); + .assistant + .set(agent) + .unwrap_or_else(|_| panic!("assistant OnceLock was already set")); } gateway @@ -475,8 +489,8 @@ impl Gateway { attachments, }.send(signal_cli).await?; - if let Some(claude) = self.claude.get() { - claude.record_message(SentBy::System, &message, Utc::now().timestamp_millis() as u64); + if let Some(assistant) = self.assistant.get() { + assistant.record_message(SentBy::System, &message, Utc::now().timestamp_millis() as u64); } } else { warn!("alert_rx is closed, halting service"); @@ -549,8 +563,8 @@ impl Gateway { attachments, }.send(signal_cli).await?; - if let Some(claude) = self.claude.get() { - claude.record_message(sourced.source, &sourced.resp.text, Utc::now().timestamp_millis() as u64); + if let Some(assistant) = self.assistant.get() { + assistant.record_message(sourced.source, &sourced.resp.text, Utc::now().timestamp_millis() as u64); } } } @@ -570,8 +584,8 @@ impl Gateway { match handling { Handling::Help => { // Record as system command - if let Some(claude) = self.claude.get() { - claude.record_message(SentBy::UserToSystem, &data.message, data.timestamp); + if let Some(assistant) = self.assistant.get() { + assistant.record_message(SentBy::UserToSystem, &data.message, data.timestamp); } Some(Ok(SourcedAdminMessageResponse { resp: AdminMessageResponse::new(self.command_router.help()), @@ -585,9 +599,9 @@ impl Gateway { Err(err) => return Some(Err((400u16, err.into()))), }; - // Record this as a system command in Claude's history - if let Some(claude) = self.claude.get() { - claude.record_message(SentBy::UserToSystem, &data.message, data.timestamp); + // Record this as a system command in the assistant's history + if let Some(assistant) = self.assistant.get() { + assistant.record_message(SentBy::UserToSystem, &data.message, data.timestamp); } let resp = match self.handle_gateway_command(cmd).await { @@ -600,23 +614,23 @@ impl Gateway { })) } Handling::Claude => { - // Send directly to Claude (use stripped message) - let Some(claude) = self.claude.get() else { - return Some(Err((501u16, "Claude is not configured".into()))); + // Send directly to the assistant (use stripped message) + let Some(assistant) = self.assistant.get() else { + return Some(Err((501u16, "Assistant is not configured".into()))); }; - match claude.request(stripped_message, data.timestamp).await { + match assistant.request(stripped_message, data.timestamp).await { Ok(resp) => Some(Ok(SourcedAdminMessageResponse { resp, - source: SentBy::Claude, + source: SentBy::Assistant, })), Err(err) => Some(Err((500, err.to_string().into()))), } } Handling::Custom(handler) => { // Record as system message - if let Some(claude) = self.claude.get() { - claude.record_message(SentBy::UserToSystem, &data.message, data.timestamp); + if let Some(assistant) = self.assistant.get() { + assistant.record_message(SentBy::UserToSystem, &data.message, data.timestamp); } // Pass stripped message to custom handler @@ -843,30 +857,30 @@ impl Gateway { } } GatewayCommand::ClaudeStop => { - let claude = self - .claude + let assistant = self + .assistant .get() - .ok_or_else(|| (501u16, "claude was not configured".into()))?; + .ok_or_else(|| (501u16, "assistant was not configured".into()))?; - claude.request_stop(); + assistant.request_stop(); Ok(AdminMessageResponse::new("stop requested")) } GatewayCommand::ClaudeCompact => { - let claude = self - .claude + let assistant = self + .assistant .get() - .ok_or_else(|| (501u16, "claude was not configured".into()))?; + .ok_or_else(|| (501u16, "assistant was not configured".into()))?; - claude.request_compaction(); + assistant.request_compaction(); Ok(AdminMessageResponse::new("compaction requested")) } GatewayCommand::ClaudeDebug => { - let claude = self - .claude + let assistant = self + .assistant .get() - .ok_or_else(|| (501u16, "claude was not configured".into()))?; + .ok_or_else(|| (501u16, "assistant was not configured".into()))?; - claude.request_debug(); + assistant.request_debug(); Ok(AdminMessageResponse::new("debug logged")) } } diff --git a/signal-gateway/src/lib.rs b/signal-gateway/src/lib.rs index 5db20cc..b350a56 100644 --- a/signal-gateway/src/lib.rs +++ b/signal-gateway/src/lib.rs @@ -6,7 +6,7 @@ #![deny(missing_docs)] pub mod alertmanager; -pub mod claude; +pub mod assistant; pub mod gateway; pub mod message_handler; diff --git a/signal-gateway/src/prometheus/mod.rs b/signal-gateway/src/prometheus/mod.rs index dfeaece..42b0def 100644 --- a/signal-gateway/src/prometheus/mod.rs +++ b/signal-gateway/src/prometheus/mod.rs @@ -1,4 +1,4 @@ -use crate::claude::{Tool, ToolExecutor, ToolResult}; +use crate::assistant::{Tool, ToolExecutor, ToolResult}; use async_trait::async_trait; use conf::Conf; use prometheus_http_client::{