abstract the claude module out of the signal-gateway and put interface and ipml in their own crates

This commit is contained in:
Chris Beck
2025-12-14 11:22:28 -07:00
parent 2682ee8ee6
commit ceb988dc07
23 changed files with 736 additions and 489 deletions
Generated
+32 -1
View File
@@ -1735,6 +1735,7 @@ dependencies = [
"reqwest", "reqwest",
"serde", "serde",
"serde_json", "serde_json",
"signal-gateway-assistant",
"thiserror", "thiserror",
"tokio", "tokio",
"tokio-util", "tokio-util",
@@ -1753,12 +1754,41 @@ dependencies = [
"reqwest", "reqwest",
"serde", "serde",
"serde_json", "serde_json",
"signal-gateway", "signal-gateway-assistant",
"tar", "tar",
"tokio", "tokio",
"tracing", "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]] [[package]]
name = "signal-gateway-bin" name = "signal-gateway-bin"
version = "0.1.0" version = "0.1.0"
@@ -1774,6 +1804,7 @@ dependencies = [
"serde_json", "serde_json",
"signal-gateway", "signal-gateway",
"signal-gateway-app-code", "signal-gateway-app-code",
"signal-gateway-assistant-claude",
"signal-gateway-log-ingest", "signal-gateway-log-ingest",
"tokio", "tokio",
"tokio-util", "tokio-util",
+5
View File
@@ -4,6 +4,8 @@ members = [
"prometheus-http-client", "prometheus-http-client",
"signal-gateway", "signal-gateway",
"signal-gateway-app-code", "signal-gateway-app-code",
"signal-gateway-assistant",
"signal-gateway-assistant/claude",
"signal-gateway-bin", "signal-gateway-bin",
"signal-gateway-log-ingest", "signal-gateway-log-ingest",
] ]
@@ -25,6 +27,9 @@ result_large_err = "allow"
[workspace.dependencies] [workspace.dependencies]
prometheus-http-client = { path = "prometheus-http-client", default-features = false } 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" async-trait = "0.1"
chrono = { version = "0.4", default-features = false, features = ["clock", "serde", "std"] } chrono = { version = "0.4", default-features = false, features = ["clock", "serde", "std"] }
+1 -1
View File
@@ -7,7 +7,7 @@ edition.workspace = true
workspace = true workspace = true
[dependencies] [dependencies]
signal-gateway = { path = "../signal-gateway", default-features = false } signal-gateway-assistant = { workspace = true }
async-trait = { workspace = true } async-trait = { workspace = true }
flate2 = { workspace = true } flate2 = { workspace = true }
+1 -1
View File
@@ -7,7 +7,7 @@ use async_trait::async_trait;
use flate2::read::GzDecoder; use flate2::read::GzDecoder;
use regex::Regex; use regex::Regex;
use serde::Deserialize; use serde::Deserialize;
use signal_gateway::claude::{Tool, ToolExecutor, ToolResult}; use signal_gateway_assistant::{Tool, ToolExecutor, ToolResult};
use std::{ use std::{
collections::HashMap, error::Error, fmt::Write, future::Future, io::Read, path::PathBuf, collections::HashMap, error::Error, fmt::Write, future::Future, io::Read, path::PathBuf,
pin::Pin, str::FromStr, sync::Arc, pin::Pin, str::FromStr, sync::Arc,
+14
View File
@@ -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 }
@@ -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 }
@@ -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 conf::Conf;
use crate::message_handler::AdminMessageResponse;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use serde_json::Value; use serde_json::Value;
use std::{path::PathBuf, sync::Weak}; use signal_gateway_assistant::{Assistant, AssistantResponse, ChatMessage, Tool, ToolExecutor};
use tokio::sync::{mpsc, oneshot}; use std::{path::PathBuf, sync::Weak, time::Duration};
use tokio_util::sync::CancellationToken;
use tracing::{error, info, warn}; use tracing::{error, info, warn};
/// Sent with inputs to claude that claude is expected to respond to. The sender /// The Anthropic API version header value.
/// gives the worker a way to return the results to the caller asynchronously. const ANTHROPIC_API_VERSION: &str = "2023-06-01";
pub type ResultSender = oneshot::Sender<Result<AdminMessageResponse, ClaudeError>>;
/// Indicates the "role" i.e. the manner in which a particular message was sent /// Configuration for the Claude API integration.
pub enum SentBy { #[derive(Clone, Conf, Debug)]
/// User message directed at the system (commands, etc.) #[conf(serde)]
UserToSystem, pub struct ClaudeConfig {
/// User message directed at Claude (prompts) /// Path to file containing the Claude API key.
UserToClaude, #[conf(long, env)]
/// Response from Claude pub api_key_file: PathBuf,
Claude, /// Paths to files containing system prompt components (in order, last one is cached).
/// System-generated message #[conf(repeat, long, env)]
System, pub system_prompt_files: Vec<PathBuf>,
/// Alert from alertmanager /// Claude API URL.
AlertManager, #[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 { /// Configuration for message buffer compaction.
/// Returns the API role: "assistant" for Claude, "user" for everything else. #[derive(Clone, Conf, Debug)]
fn api_role(&self) -> &'static str { #[conf(serde)]
match self { pub struct CompactionConfig {
Self::Claude => "assistant", /// Path to file containing the compaction prompt.
_ => "user", #[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<Duration>,
} }
/// Returns a prefix to prepend to message text for context. /// Error type for Claude API operations.
fn prefix(&self) -> Option<&'static str> { #[derive(Debug, thiserror::Error)]
match self { pub enum ClaudeError {
Self::UserToSystem => Some("[user to system]"), /// Failed to read API key file.
Self::UserToClaude => None, // No prefix needed for direct user messages #[error("failed to read API key file: {0}")]
Self::Claude => None, ApiKeyRead(std::io::Error),
Self::System => Some("[system]"), /// Failed to read system prompt file.
Self::AlertManager => Some("[alertmanager]"), #[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<str>),
/// 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,
} }
/// An input to the worker sent through the channel /// Claude API assistant implementation.
pub enum Input { pub struct ClaudeAssistant {
Chat(ChatMessage),
Compact,
Debug,
}
/// A chat message
pub struct ChatMessage {
pub sent_by: SentBy,
pub timestamp: DateTime<Utc>,
pub text: Box<str>,
/// Present when claude is expected to respond to the message
pub result_sender: Option<ResultSender>,
}
impl ChatMessage {
fn into_content_and_sender(self) -> (MessageContent, Option<ResultSender>) {
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 {
config: ClaudeConfig, config: ClaudeConfig,
client: reqwest::Client, client: reqwest::Client,
api_key: String, api_key: String,
@@ -91,41 +98,35 @@ pub struct ClaudeWorker {
summary: String, summary: String,
messages: Vec<MessageContent>, messages: Vec<MessageContent>,
tool_executor: Weak<dyn ToolExecutor>, tool_executor: Weak<dyn ToolExecutor>,
input_rx: mpsc::Receiver<Input>,
stop_rx: mpsc::Receiver<()>,
/// Timestamp of the last automatic compaction (not user-requested). /// Timestamp of the last automatic compaction (not user-requested).
last_auto_compaction: Option<std::time::Instant>, last_auto_compaction: Option<std::time::Instant>,
} }
impl ClaudeWorker { impl ClaudeAssistant {
/// Create a new Claude worker. /// Create a new Claude assistant from configuration.
/// ///
/// Reads the API key, system prompts, and compaction prompt from the configured files. /// Reads the API key, system prompts, and compaction prompt from the configured files.
pub fn new( pub fn new(
config: ClaudeConfig, config: ClaudeConfig,
tool_executor: Weak<dyn ToolExecutor>, tool_executor: Weak<dyn ToolExecutor>,
input_rx: mpsc::Receiver<Input>,
stop_rx: mpsc::Receiver<()>,
) -> Result<Self, ClaudeError> { ) -> Result<Self, ClaudeError> {
let api_key = std::fs::read_to_string(&config.api_key_file) let api_key = std::fs::read_to_string(&config.api_key_file)
.map_err(ClaudeError::ApiKeyRead)? .map_err(ClaudeError::ApiKeyRead)?
.trim() .trim()
.to_owned(); .to_string();
let system_prompts: Vec<String> = config let system_prompts: Result<Vec<String>, ClaudeError> = config
.system_prompt_files .system_prompt_files
.iter() .iter()
.map(|path| { .map(|path| {
std::fs::read_to_string(path) std::fs::read_to_string(path)
.map_err(|e| ClaudeError::SystemPromptRead(path.clone(), e)) .map_err(|e| ClaudeError::SystemPromptRead(path.clone(), e))
}) })
.collect::<Result<_, _>>()?; .collect();
let system_prompts = system_prompts?;
let compaction_prompt = { let compaction_prompt = std::fs::read_to_string(&config.compaction.prompt_file)
let path = &config.compaction.prompt_file; .map_err(|e| ClaudeError::SystemPromptRead(config.compaction.prompt_file.clone(), e))?;
std::fs::read_to_string(path)
.map_err(|e| ClaudeError::SystemPromptRead(path.clone(), e))?
};
Ok(Self { Ok(Self {
config, config,
@@ -134,71 +135,12 @@ impl ClaudeWorker {
system_prompts, system_prompts,
compaction_prompt, compaction_prompt,
summary: String::new(), summary: String::new(),
messages: Default::default(), messages: Vec::new(),
tool_executor, tool_executor,
input_rx,
stop_rx,
last_auto_compaction: None, 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. /// Calculate total characters in the message buffer.
fn message_buffer_chars(&self) -> usize { fn message_buffer_chars(&self) -> usize {
self.messages self.messages
@@ -216,6 +158,29 @@ impl ClaudeWorker {
.sum() .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. /// Drop oldest messages until buffer is under the trigger threshold.
fn drop_oldest_messages(&mut self) { fn drop_oldest_messages(&mut self) {
let target = self.config.compaction.trigger_chars as usize; let target = self.config.compaction.trigger_chars as usize;
@@ -228,7 +193,6 @@ impl ClaudeWorker {
let mut current_chars = before_chars; let mut current_chars = before_chars;
let mut dropped = 0; let mut dropped = 0;
// Remove from the front (oldest) until under limit
while current_chars > target && !self.messages.is_empty() { while current_chars > target && !self.messages.is_empty() {
let msg_chars = self.messages[0] let msg_chars = self.messages[0]
.content .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. /// Perform compaction by summarizing messages and storing the result.
/// async fn do_compact(&mut self, is_automatic: bool) {
/// 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) {
if self.messages.is_empty() { if self.messages.is_empty() {
return; return;
} }
// Update timestamp for automatic compactions (before the API call)
if is_automatic { if is_automatic {
self.last_auto_compaction = Some(std::time::Instant::now()); self.last_auto_compaction = Some(std::time::Instant::now());
} }
@@ -312,7 +243,6 @@ impl ClaudeWorker {
if !self.summary.is_empty() { if !self.summary.is_empty() {
system.push(SystemContent::text(&self.summary)); system.push(SystemContent::text(&self.summary));
} }
// Mark the last one as cached if prompt caching is enabled
if self.config.prompt_caching if self.config.prompt_caching
&& let Some(last) = system.last_mut() && let Some(last) = system.last_mut()
{ {
@@ -341,13 +271,12 @@ impl ClaudeWorker {
Ok(response) if response.status().is_success() => { Ok(response) if response.status().is_success() => {
match response.json::<MessagesResponse>().await { match response.json::<MessagesResponse>().await {
Ok(parsed) => { Ok(parsed) => {
// Extract text from response
let summary_text: String = parsed let summary_text: String = parsed
.content .content
.into_iter() .into_iter()
.filter_map(|block| { .filter_map(|block| {
if let ContentBlock::Text { text } = block { if let ContentBlock::Text { text } = block {
Some(text) Some(String::from(text))
} else { } else {
None None
} }
@@ -355,7 +284,6 @@ impl ClaudeWorker {
.collect::<Vec<_>>() .collect::<Vec<_>>()
.join("\n"); .join("\n");
// Wrap in XML tags and store
self.summary = self.summary =
format!("<summary type=\"activity\">\n{}\n</summary>", summary_text); format!("<summary type=\"activity\">\n{}\n</summary>", summary_text);
self.messages.clear(); self.messages.clear();
@@ -382,25 +310,18 @@ impl ClaudeWorker {
} }
} }
/// Log the message buffer for debugging. /// Handle a single request to the Claude API with tool use loop.
fn handle_debug(&self) { ///
if self.summary.is_empty() { /// Returns `Ok(None)` if the operation was cancelled.
info!("Claude summary: (empty)"); async fn handle_request(
} else { &mut self,
info!("Claude summary:\n{}", self.summary); cancel: &CancellationToken,
} ) -> Result<Option<AssistantResponse>, ClaudeError> {
info!("Claude message buffer:\n{:#?}", self.messages);
}
/// Handle a single request to the Claude API.
async fn handle_request(&mut self) -> Result<AdminMessageResponse, ClaudeError> {
let max_iterations = self.config.claude_max_iterations; let max_iterations = self.config.claude_max_iterations;
// Get tools from the executor if still alive
let executor = self.tool_executor.upgrade(); let executor = self.tool_executor.upgrade();
let tools = executor.as_ref().map(|te| te.tools()).unwrap_or_default(); let tools = executor.as_ref().map(|te| te.tools()).unwrap_or_default();
// Collect attachments from tool results across all iterations
let mut attachments: Vec<PathBuf> = Vec::new(); let mut attachments: Vec<PathBuf> = Vec::new();
if let Some(last) = self.messages.last() if let Some(last) = self.messages.last()
@@ -409,7 +330,7 @@ impl ClaudeWorker {
info!("Claude request: {}", text); 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<SystemContent> = self let mut system: Vec<SystemContent> = self
.system_prompts .system_prompts
.iter() .iter()
@@ -418,7 +339,6 @@ impl ClaudeWorker {
if !self.summary.is_empty() { if !self.summary.is_empty() {
system.push(SystemContent::text(&self.summary)); system.push(SystemContent::text(&self.summary));
} }
// Mark the last one as cached if prompt caching is enabled
if self.config.prompt_caching if self.config.prompt_caching
&& let Some(last) = system.last_mut() && let Some(last) = system.last_mut()
{ {
@@ -426,8 +346,9 @@ impl ClaudeWorker {
} }
for iteration in 0..max_iterations { for iteration in 0..max_iterations {
// Check for stop before making API call if cancel.is_cancelled() {
self.check_stop()?; return Ok(None);
}
let request_body = MessagesRequest { let request_body = MessagesRequest {
model: &self.config.claude_model, model: &self.config.claude_model,
@@ -458,28 +379,24 @@ impl ClaudeWorker {
response.stop_reason, response.content response.stop_reason, response.content
); );
// Check if we need to handle tool use
if response.stop_reason.as_ref() == "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 { let Some(executor) = self.tool_executor.upgrade() else {
return Err(ClaudeError::ToolExecutorGone); return Err(ClaudeError::ToolExecutorGone);
}; };
// Add assistant's response to messages
self.messages self.messages
.push(MessageContent::assistant(response.content.clone())); .push(MessageContent::assistant(response.content.clone()));
// Execute each tool use and collect results
for block in &response.content { for block in &response.content {
if let ContentBlock::ToolUse { id, name, input } = block { if let ContentBlock::ToolUse { id, name, input } = block {
// Check for stop before each tool use if cancel.is_cancelled() {
self.check_stop()?; return Ok(None);
}
info!("Claude tool use: {}({})", name, input); info!("Claude tool use: {}({})", name, input);
let (result_text, is_error) = match executor.execute(name, input).await { let (result_text, is_error) = match executor.execute(name, input).await {
Ok(tool_result) => { Ok(tool_result) => {
info!("Tool result: {}", tool_result.text); info!("Tool result: {}", tool_result.text);
// Collect any attachments from the tool result
attachments.extend(tool_result.attachments); attachments.extend(tool_result.attachments);
(tool_result.text, false) (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); info!("Tool use iteration {}, continuing...", iteration + 1);
continue; continue;
} }
@@ -507,7 +423,7 @@ impl ClaudeWorker {
.into_iter() .into_iter()
.filter_map(|block| { .filter_map(|block| {
if let ContentBlock::Text { text, .. } = block { if let ContentBlock::Text { text, .. } = block {
Some(text) Some(text.into_string())
} else { } else {
None None
} }
@@ -516,34 +432,88 @@ impl ClaudeWorker {
.join("\n"); .join("\n");
info!("Claude final result: {}", text); 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)) 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<Option<AssistantResponse>, Box<dyn std::error::Error + Send + Sync>> {
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. /// Estimate the serialized size of a JSON Value without allocating.
fn estimate_json_size(value: &Value) -> usize { fn estimate_json_size(value: &Value) -> usize {
match value { match value {
Value::Null => 4, // "null" Value::Null => 4,
Value::Bool(true) => 4, // "true" Value::Bool(true) => 4,
Value::Bool(false) => 5, // "false" Value::Bool(false) => 5,
Value::Number(n) => n.to_string().len(), // Numbers are small, ok to alloc Value::Number(n) => n.to_string().len(),
Value::String(s) => s.len() + 2, // quotes Value::String(s) => s.len() + 2,
Value::Array(arr) => { Value::Array(arr) => {
2 + arr.iter().map(estimate_json_size).sum::<usize>() + arr.len().saturating_sub(1) 2 + arr.iter().map(estimate_json_size).sum::<usize>() + arr.len().saturating_sub(1)
} }
Value::Object(obj) => { Value::Object(obj) => {
2 + obj 2 + obj
.iter() .iter()
.map(|(k, v)| k.len() + 3 + estimate_json_size(v)) // "key":value .map(|(k, v)| k.len() + 3 + estimate_json_size(v))
.sum::<usize>() .sum::<usize>()
+ obj.len().saturating_sub(1) // commas + obj.len().saturating_sub(1)
} }
} }
} }
// ---- API Types ----
/// Request body for the Claude Messages API. /// Request body for the Claude Messages API.
#[derive(Serialize)] #[derive(Serialize)]
struct MessagesRequest<'a> { struct MessagesRequest<'a> {
+32
View File
@@ -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<Option<AssistantResponse>, Box<dyn Error + Send + Sync>>;
/// 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);
}
@@ -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<Utc>,
pub text: Box<str>,
}
+14
View File
@@ -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};
+49
View File
@@ -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<PathBuf>,
}
impl AssistantResponse {
/// Create a new response with just text.
pub fn new(text: impl Into<String>) -> Self {
Self {
text: text.into(),
attachments: Vec::new(),
}
}
/// Create a response with text and attachments.
pub fn with_attachments(text: impl Into<String>, attachments: Vec<PathBuf>) -> Self {
Self {
text: text.into(),
attachments,
}
}
/// Add attachments to this response.
pub fn add_attachments(mut self, attachments: Vec<PathBuf>) -> Self {
self.attachments.extend(attachments);
self
}
}
impl From<String> for AssistantResponse {
fn from(text: String) -> Self {
Self::new(text)
}
}
impl From<&str> for AssistantResponse {
fn from(text: &str) -> Self {
Self::new(text)
}
}
@@ -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 async_trait::async_trait;
use serde::Serialize; use serde::Serialize;
@@ -8,7 +8,7 @@ use std::path::PathBuf;
/// Result of executing a tool. /// Result of executing a tool.
#[derive(Clone, Debug, Default)] #[derive(Clone, Debug, Default)]
pub struct ToolResult { pub struct ToolResult {
/// Text result to return to Claude. /// Text result to return to LLM.
pub text: String, pub text: String,
/// Optional file attachments generated by the tool. /// Optional file attachments generated by the tool.
pub attachments: Vec<PathBuf>, pub attachments: Vec<PathBuf>,
@@ -56,11 +56,11 @@ pub trait ToolExecutor: Send + Sync {
} }
/// Execute a tool by name with the given input arguments. /// 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<ToolResult, String>; async fn execute(&self, name: &str, input: &Value) -> Result<ToolResult, String>;
} }
/// A tool definition for the Claude API. /// A tool definition for the LLM API.
#[derive(Clone, Debug, Serialize)] #[derive(Clone, Debug, Serialize)]
pub struct Tool { pub struct Tool {
/// The name of the tool. /// The name of the tool.
+2 -1
View File
@@ -16,7 +16,8 @@ plot = ["signal-gateway/plot"]
rustls-tls = ["signal-gateway/rustls-tls"] rustls-tls = ["signal-gateway/rustls-tls"]
[dependencies] [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-app-code = { path = "../signal-gateway-app-code" }
signal-gateway-log-ingest = { path = "../signal-gateway-log-ingest" } signal-gateway-log-ingest = { path = "../signal-gateway-log-ingest" }
+15 -1
View File
@@ -7,6 +7,7 @@ use hyper::service::service_fn;
use hyper_util::{rt::TokioIo, server::conn::auto}; use hyper_util::{rt::TokioIo, server::conn::auto};
use signal_gateway::{CommandRouter, Gateway, GatewayConfig, Handling}; use signal_gateway::{CommandRouter, Gateway, GatewayConfig, Handling};
use signal_gateway_app_code::AppCodeTools; 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 std::{env, fs, net::SocketAddr, path::PathBuf, sync::Arc, time::Duration};
use tokio::net::TcpListener; use tokio::net::TcpListener;
use tokio_util::sync::CancellationToken; use tokio_util::sync::CancellationToken;
@@ -46,6 +47,9 @@ pub struct Config {
/// Application source code configurations for Claude tools. /// Application source code configurations for Claude tools.
#[conf(long, env, value_parser = serde_json::from_str, default, default_help_str = "[]")] #[conf(long, env, value_parser = serde_json::from_str, default, default_help_str = "[]")]
app_code: Vec<AppCodeConfigExt>, app_code: Vec<AppCodeConfigExt>,
/// Claude API configuration for AI-powered responses.
#[conf(flatten, prefix)]
claude: Option<ClaudeConfig>,
#[conf(flatten, serde(flatten))] #[conf(flatten, serde(flatten))]
gateway: GatewayConfig, gateway: GatewayConfig,
} }
@@ -121,7 +125,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
router_builder = router_builder.route("/", Handling::GatewayCommand); router_builder = router_builder.route("/", Handling::GatewayCommand);
// Add Claude as default handler if configured // 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); router_builder = router_builder.route("", Handling::Claude);
} }
@@ -153,6 +157,16 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
gateway_builder = gateway_builder.with_tools(tools); 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 gateway = gateway_builder.build().await;
let listener = TcpListener::bind(config.http_listen_addr).await.unwrap(); let listener = TcpListener::bind(config.http_listen_addr).await.unwrap();
+1 -1
View File
@@ -7,7 +7,7 @@ edition.workspace = true
workspace = true workspace = true
[dependencies] [dependencies]
signal-gateway = { path = "../signal-gateway" } signal-gateway = { workspace = true }
chrono = { workspace = true } chrono = { workspace = true }
conf = { workspace = true } conf = { workspace = true }
+1
View File
@@ -13,6 +13,7 @@ rustls-tls = ["prometheus-http-client/rustls-tls", "reqwest/rustls-tls"]
[dependencies] [dependencies]
prometheus-http-client = { workspace = true, default-features = false } prometheus-http-client = { workspace = true, default-features = false }
signal-gateway-assistant = { workspace = true }
async-trait = { workspace = true } async-trait = { workspace = true }
chrono = { workspace = true } chrono = { workspace = true }
+118
View File
@@ -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<Input>,
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<dyn Assistant>) -> 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<AdminMessageResponse, AssistantError> {
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(());
}
}
+124
View File
@@ -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<Result<AdminMessageResponse, AssistantError>>;
/// 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<dyn Assistant>,
input_rx: mpsc::Receiver<Input>,
stop_rx: mpsc::Receiver<()>,
cancel_token: CancellationToken,
}
impl AssistantWorker {
/// Create a new assistant worker.
pub fn new(
assistant: Box<dyn Assistant>,
input_rx: mpsc::Receiver<Input>,
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)
}
-210
View File
@@ -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<PathBuf>,
/// 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<Duration>,
}
/// 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<str>),
/// 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<Input>,
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<dyn ToolExecutor>,
) -> Result<Self, ClaudeError> {
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<AdminMessageResponse, ClaudeError> {
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<Result<_, ClaudeError>, 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(());
}
}
+1 -1
View File
@@ -4,7 +4,7 @@ use super::{
route::{Destination, Limit, Route}, route::{Destination, Limit, Route},
}; };
use crate::{ use crate::{
claude::{Tool, ToolExecutor, ToolResult}, assistant::{Tool, ToolExecutor, ToolResult},
concurrent_map::LazyMap, concurrent_map::LazyMap,
log_format::LogFormatConfig, log_format::LogFormatConfig,
log_message::{LogFilter, LogMessage, Origin}, log_message::{LogFilter, LogMessage, Origin},
+60 -46
View File
@@ -4,7 +4,7 @@
use crate::signal_jsonrpc::connect_ipc; use crate::signal_jsonrpc::connect_ipc;
use crate::{ use crate::{
alertmanager::AlertPost, alertmanager::AlertPost,
claude::{ClaudeAgent, ClaudeConfig, SentBy, Tool, ToolExecutor, ToolResult}, assistant::{AssistantAgent, SentBy, Tool, ToolExecutor, ToolResult},
log_message::{LogMessage, Origin}, log_message::{LogMessage, Origin},
message_handler::{AdminMessage, AdminMessageResponse, Context, MessageHandlerResult}, message_handler::{AdminMessage, AdminMessageResponse, Context, MessageHandlerResult},
prometheus::{Prometheus, PrometheusConfig}, prometheus::{Prometheus, PrometheusConfig},
@@ -101,9 +101,6 @@ pub struct GatewayConfig {
/// Log handler configuration for processing log messages. /// Log handler configuration for processing log messages.
#[conf(flatten, prefix)] #[conf(flatten, prefix)]
pub log_handler: LogHandlerConfig, pub log_handler: LogHandlerConfig,
/// Claude API configuration for AI-powered responses.
#[conf(flatten, prefix)]
pub claude: Option<ClaudeConfig>,
} }
/// Wrapper for parsing gateway commands /// Wrapper for parsing gateway commands
@@ -255,19 +252,23 @@ pub struct Gateway {
log_handler: LogHandler, log_handler: LogHandler,
/// Command router for dispatching admin messages. /// Command router for dispatching admin messages.
command_router: CommandRouter, command_router: CommandRouter,
/// Claude API client for AI-powered responses. /// Assistant agent for AI-powered responses.
/// Initialized after Arc creation so it can hold a weak reference back to Gateway. assistant: OnceLock<AssistantAgent>,
claude: OnceLock<Box<ClaudeAgent>>,
/// Additional tool executors added via the builder. /// Additional tool executors added via the builder.
extra_tool_executors: Vec<Arc<dyn ToolExecutor>>, extra_tool_executors: Vec<Arc<dyn ToolExecutor>>,
} }
/// Type alias for the assistant factory function.
pub type AssistantFactory =
Box<dyn FnOnce(Weak<dyn ToolExecutor>) -> Box<dyn crate::assistant::Assistant> + Send>;
/// Builder for creating a [`Gateway`] with optional additional tool executors. /// Builder for creating a [`Gateway`] with optional additional tool executors.
pub struct GatewayBuilder { pub struct GatewayBuilder {
config: GatewayConfig, config: GatewayConfig,
token: Option<CancellationToken>, token: Option<CancellationToken>,
command_router: Option<CommandRouter>, command_router: Option<CommandRouter>,
extra_tool_executors: Vec<Arc<dyn ToolExecutor>>, extra_tool_executors: Vec<Arc<dyn ToolExecutor>>,
assistant_factory: Option<AssistantFactory>,
} }
impl GatewayBuilder { impl GatewayBuilder {
@@ -278,6 +279,7 @@ impl GatewayBuilder {
token: None, token: None,
command_router: None, command_router: None,
extra_tool_executors: Vec::new(), extra_tool_executors: Vec::new(),
assistant_factory: None,
} }
} }
@@ -295,13 +297,25 @@ impl GatewayBuilder {
/// Add an additional tool executor to the gateway. /// 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.). /// the built-in tools (log handler, prometheus, etc.).
pub fn with_tools(mut self, executor: Arc<dyn ToolExecutor>) -> Self { pub fn with_tools(mut self, executor: Arc<dyn ToolExecutor>) -> Self {
self.extra_tool_executors.push(executor); self.extra_tool_executors.push(executor);
self 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<F>(mut self, factory: F) -> Self
where
F: FnOnce(Weak<dyn ToolExecutor>) -> Box<dyn crate::assistant::Assistant> + Send + 'static,
{
self.assistant_factory = Some(Box::new(factory));
self
}
/// Build the gateway. /// Build the gateway.
pub async fn build(self) -> Arc<Gateway> { pub async fn build(self) -> Arc<Gateway> {
Gateway::new_internal( Gateway::new_internal(
@@ -309,6 +323,7 @@ impl GatewayBuilder {
self.token.unwrap_or_default(), self.token.unwrap_or_default(),
self.command_router.unwrap_or_default(), self.command_router.unwrap_or_default(),
self.extra_tool_executors, self.extra_tool_executors,
self.assistant_factory,
) )
.await .await
} }
@@ -326,6 +341,7 @@ impl Gateway {
token: CancellationToken, token: CancellationToken,
command_router: CommandRouter, command_router: CommandRouter,
extra_tool_executors: Vec<Arc<dyn ToolExecutor>>, extra_tool_executors: Vec<Arc<dyn ToolExecutor>>,
assistant_factory: Option<AssistantFactory>,
) -> Arc<Self> { ) -> Arc<Self> {
let (signal_alert_mq_tx, signal_alert_mq_rx) = unbounded_channel(); 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 log_handler = LogHandler::new(config.log_handler.clone(), signal_alert_mq_tx.clone());
let claude_config = config.claude.clone();
let gateway = Arc::new(Self { let gateway = Arc::new(Self {
config, config,
signal_alert_mq_tx, signal_alert_mq_tx,
@@ -348,18 +362,18 @@ impl Gateway {
prometheus, prometheus,
log_handler, log_handler,
command_router, command_router,
claude: OnceLock::new(), assistant: OnceLock::new(),
extra_tool_executors, extra_tool_executors,
}); });
// Initialize Claude with a weak reference back to the gateway // Initialize the assistant agent using the factory if one was provided
if let Some(cc) = claude_config { if let Some(factory) = assistant_factory {
let claude = ClaudeAgent::new(cc, Arc::downgrade(&gateway) as Weak<dyn ToolExecutor>) let assistant = factory(Arc::downgrade(&gateway) as Weak<dyn ToolExecutor>);
.expect("Invalid claude config"); let agent = AssistantAgent::new(assistant);
gateway gateway
.claude .assistant
.set(Box::new(claude)) .set(agent)
.unwrap_or_else(|_| panic!("claude OnceLock was already set")); .unwrap_or_else(|_| panic!("assistant OnceLock was already set"));
} }
gateway gateway
@@ -475,8 +489,8 @@ impl Gateway {
attachments, attachments,
}.send(signal_cli).await?; }.send(signal_cli).await?;
if let Some(claude) = self.claude.get() { if let Some(assistant) = self.assistant.get() {
claude.record_message(SentBy::System, &message, Utc::now().timestamp_millis() as u64); assistant.record_message(SentBy::System, &message, Utc::now().timestamp_millis() as u64);
} }
} else { } else {
warn!("alert_rx is closed, halting service"); warn!("alert_rx is closed, halting service");
@@ -549,8 +563,8 @@ impl Gateway {
attachments, attachments,
}.send(signal_cli).await?; }.send(signal_cli).await?;
if let Some(claude) = self.claude.get() { if let Some(assistant) = self.assistant.get() {
claude.record_message(sourced.source, &sourced.resp.text, Utc::now().timestamp_millis() as u64); assistant.record_message(sourced.source, &sourced.resp.text, Utc::now().timestamp_millis() as u64);
} }
} }
} }
@@ -570,8 +584,8 @@ impl Gateway {
match handling { match handling {
Handling::Help => { Handling::Help => {
// Record as system command // Record as system command
if let Some(claude) = self.claude.get() { if let Some(assistant) = self.assistant.get() {
claude.record_message(SentBy::UserToSystem, &data.message, data.timestamp); assistant.record_message(SentBy::UserToSystem, &data.message, data.timestamp);
} }
Some(Ok(SourcedAdminMessageResponse { Some(Ok(SourcedAdminMessageResponse {
resp: AdminMessageResponse::new(self.command_router.help()), resp: AdminMessageResponse::new(self.command_router.help()),
@@ -585,9 +599,9 @@ impl Gateway {
Err(err) => return Some(Err((400u16, err.into()))), Err(err) => return Some(Err((400u16, err.into()))),
}; };
// Record this as a system command in Claude's history // Record this as a system command in the assistant's history
if let Some(claude) = self.claude.get() { if let Some(assistant) = self.assistant.get() {
claude.record_message(SentBy::UserToSystem, &data.message, data.timestamp); assistant.record_message(SentBy::UserToSystem, &data.message, data.timestamp);
} }
let resp = match self.handle_gateway_command(cmd).await { let resp = match self.handle_gateway_command(cmd).await {
@@ -600,23 +614,23 @@ impl Gateway {
})) }))
} }
Handling::Claude => { Handling::Claude => {
// Send directly to Claude (use stripped message) // Send directly to the assistant (use stripped message)
let Some(claude) = self.claude.get() else { let Some(assistant) = self.assistant.get() else {
return Some(Err((501u16, "Claude is not configured".into()))); 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 { Ok(resp) => Some(Ok(SourcedAdminMessageResponse {
resp, resp,
source: SentBy::Claude, source: SentBy::Assistant,
})), })),
Err(err) => Some(Err((500, err.to_string().into()))), Err(err) => Some(Err((500, err.to_string().into()))),
} }
} }
Handling::Custom(handler) => { Handling::Custom(handler) => {
// Record as system message // Record as system message
if let Some(claude) = self.claude.get() { if let Some(assistant) = self.assistant.get() {
claude.record_message(SentBy::UserToSystem, &data.message, data.timestamp); assistant.record_message(SentBy::UserToSystem, &data.message, data.timestamp);
} }
// Pass stripped message to custom handler // Pass stripped message to custom handler
@@ -843,30 +857,30 @@ impl Gateway {
} }
} }
GatewayCommand::ClaudeStop => { GatewayCommand::ClaudeStop => {
let claude = self let assistant = self
.claude .assistant
.get() .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")) Ok(AdminMessageResponse::new("stop requested"))
} }
GatewayCommand::ClaudeCompact => { GatewayCommand::ClaudeCompact => {
let claude = self let assistant = self
.claude .assistant
.get() .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")) Ok(AdminMessageResponse::new("compaction requested"))
} }
GatewayCommand::ClaudeDebug => { GatewayCommand::ClaudeDebug => {
let claude = self let assistant = self
.claude .assistant
.get() .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")) Ok(AdminMessageResponse::new("debug logged"))
} }
} }
+1 -1
View File
@@ -6,7 +6,7 @@
#![deny(missing_docs)] #![deny(missing_docs)]
pub mod alertmanager; pub mod alertmanager;
pub mod claude; pub mod assistant;
pub mod gateway; pub mod gateway;
pub mod message_handler; pub mod message_handler;
+1 -1
View File
@@ -1,4 +1,4 @@
use crate::claude::{Tool, ToolExecutor, ToolResult}; use crate::assistant::{Tool, ToolExecutor, ToolResult};
use async_trait::async_trait; use async_trait::async_trait;
use conf::Conf; use conf::Conf;
use prometheus_http_client::{ use prometheus_http_client::{