From 51657dbd385e177208549a667c1d12f5b2f7680e Mon Sep 17 00:00:00 2001 From: Chris Beck Date: Sat, 13 Dec 2025 21:02:17 -0700 Subject: [PATCH] cargo fmt, allow rate-limiting requests to claude --- signal-gateway-app-code/src/lib.rs | 56 ++++++-------- signal-gateway-bin/src/app_code.rs | 13 ++-- signal-gateway/src/claude/mod.rs | 10 ++- signal-gateway/src/claude/worker.rs | 77 +++++++++++++++++-- signal-gateway/src/gateway/mod.rs | 5 +- .../src/gateway/signal_trust_set.rs | 12 ++- signal-gateway/src/prometheus/mod.rs | 3 +- 7 files changed, 117 insertions(+), 59 deletions(-) diff --git a/signal-gateway-app-code/src/lib.rs b/signal-gateway-app-code/src/lib.rs index 8477506..cfb1c8a 100644 --- a/signal-gateway-app-code/src/lib.rs +++ b/signal-gateway-app-code/src/lib.rs @@ -267,9 +267,7 @@ impl AppCode { /// If `path` is None or empty, lists the root directory. pub async fn ls(&self, path: Option<&str>) -> Result { let cache = self.get_current_tarball().await; - let cached = cache - .as_ref() - .ok_or("source code not available")?; + let cached = cache.as_ref().ok_or("source code not available")?; let prefix = path.unwrap_or("").trim_start_matches('/'); let prefix = if prefix.is_empty() { @@ -293,16 +291,17 @@ impl AppCode { // Get just the first component (file or directory name) if let Some(first) = remainder.split('/').next() - && !first.is_empty() { - // Check if it's a directory (has more components) - let is_dir = remainder.contains('/'); - let entry = if is_dir { - format!("{}/", first) - } else { - first.to_string() - }; - entries.insert(entry); - } + && !first.is_empty() + { + // Check if it's a directory (has more components) + let is_dir = remainder.contains('/'); + let entry = if is_dir { + format!("{}/", first) + } else { + first.to_string() + }; + entries.insert(entry); + } } } @@ -321,9 +320,7 @@ impl AppCode { /// Supports simple glob patterns with `*` wildcards. pub async fn find(&self, pattern: Option<&str>) -> Result { let cache = self.get_current_tarball().await; - let cached = cache - .as_ref() - .ok_or("source code not available")?; + let cached = cache.as_ref().ok_or("source code not available")?; let pattern = pattern.unwrap_or("*"); @@ -357,9 +354,7 @@ impl AppCode { line_end: Option, ) -> Result { let cache = self.get_current_tarball().await; - let cached = cache - .as_ref() - .ok_or("source code not available")?; + let cached = cache.as_ref().ok_or("source code not available")?; let path = path.trim_start_matches('/'); let file = cached @@ -404,9 +399,7 @@ impl AppCode { let regex = Regex::new(pattern).map_err(|e| format!("Invalid regex: {e}"))?; let cache = self.get_current_tarball().await; - let cached = cache - .as_ref() - .ok_or("source code not available")?; + let cached = cache.as_ref().ok_or("source code not available")?; let prefix = path_prefix.map(|p| p.trim_start_matches('/')); @@ -421,9 +414,10 @@ impl AppCode { 'outer: for (path, file) in sorted_files { // Skip if path doesn't match prefix if let Some(prefix) = prefix - && !path.starts_with(prefix) { - continue; - } + && !path.starts_with(prefix) + { + continue; + } // Skip binary-looking files if looks_binary(&file.content) { @@ -472,10 +466,11 @@ impl AppCode { // Add separator if there's a gap if let Some(&last) = printed.iter().next_back() - && start > last + 1 { - writeln!(&mut output, "---") - .map_err(|e| format!("Format error: {e}"))?; - } + && start > last + 1 + { + writeln!(&mut output, "---") + .map_err(|e| format!("Format error: {e}"))?; + } for (i, line) in lines[start..end].iter().enumerate() { let line_idx = start + i; @@ -662,8 +657,7 @@ impl ToolExecutor for AppCodeTools { }, Tool { name: "code_search", - description: - "Search for a regex pattern in an application's source code (like grep).", + description: "Search for a regex pattern in an application's source code (like grep).", input_schema: serde_json::json!({ "type": "object", "properties": { diff --git a/signal-gateway-bin/src/app_code.rs b/signal-gateway-bin/src/app_code.rs index 477e3bd..f27a12d 100644 --- a/signal-gateway-bin/src/app_code.rs +++ b/signal-gateway-bin/src/app_code.rs @@ -26,21 +26,20 @@ impl AppCodeConfigExt { let url = url.clone(); let client = client.clone(); Box::pin(async move { - let response = client - .get(url.as_str()) - .send() - .await - .map_err(|e| -> Box { + let response = client.get(url.as_str()).send().await.map_err( + |e| -> Box { Box::new(std::io::Error::other(format!( "HTTP request to {url} failed: {e}" ))) - })?; + }, + )?; if !response.status().is_success() { return Err(Box::new(std::io::Error::other(format!( "HTTP request to {url} returned {}", response.status() - ))) as Box); + ))) + as Box); } let mut sha = response diff --git a/signal-gateway/src/claude/mod.rs b/signal-gateway/src/claude/mod.rs index dd9a090..ad135f4 100644 --- a/signal-gateway/src/claude/mod.rs +++ b/signal-gateway/src/claude/mod.rs @@ -9,8 +9,7 @@ pub use worker::SentBy; use crate::message_handler::AdminMessageResponse; use chrono::{DateTime, Utc}; use conf::Conf; -use std::path::PathBuf; -use std::sync::Weak; +use std::{path::PathBuf, sync::Weak, time::Duration}; use tokio::sync::{mpsc, oneshot}; use worker::{ChatMessage, ClaudeWorker, Input}; @@ -69,6 +68,11 @@ pub struct CompactionConfig { /// 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. @@ -172,7 +176,7 @@ impl ClaudeAgent { result_rx.await.map_err(|_| ClaudeError::WorkerGone)? } - /// Record a message into claude's chat log that claude is not expected to respond to + /// 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 { diff --git a/signal-gateway/src/claude/worker.rs b/signal-gateway/src/claude/worker.rs index d2931ea..00d01f2 100644 --- a/signal-gateway/src/claude/worker.rs +++ b/signal-gateway/src/claude/worker.rs @@ -5,8 +5,7 @@ use crate::message_handler::AdminMessageResponse; use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; use serde_json::Value; -use std::path::PathBuf; -use std::sync::Weak; +use std::{path::PathBuf, sync::Weak}; use tokio::sync::{mpsc, oneshot}; use tracing::{error, info, warn}; @@ -94,6 +93,8 @@ pub struct ClaudeWorker { 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 { @@ -137,6 +138,7 @@ impl ClaudeWorker { tool_executor, input_rx, stop_rx, + last_auto_compaction: None, }) } @@ -154,10 +156,22 @@ impl ClaudeWorker { let (mc, maybe_sender) = msg.into_content_and_sender(); self.messages.push(mc); - // Check if we need to trigger compaction + // 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 { - self.handle_compact().await; + // 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 { @@ -171,7 +185,7 @@ impl ClaudeWorker { } }, Input::Compact => { - self.handle_compact().await; + self.handle_compact(false).await; } Input::Debug => { self.handle_debug(); @@ -202,6 +216,43 @@ impl ClaudeWorker { .sum() } + /// 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; + let before_chars = self.message_buffer_chars(); + + if before_chars <= target { + return; + } + + 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 + .iter() + .map(|block| match block { + ContentBlock::Text { text } => text.len(), + ContentBlock::ToolUse { input, .. } => estimate_json_size(input), + ContentBlock::ToolResult { content, .. } => content.len(), + }) + .sum::(); + current_chars = current_chars.saturating_sub(msg_chars); + self.messages.remove(0); + dropped += 1; + } + + let after_chars = self.message_buffer_chars(); + warn!( + "Compaction rate-limited: dropped {} oldest messages ({} -> {} chars)", + dropped, + before_chars, + after_chars + ); + } + /// Handle a stop request by draining queues and sending errors to pending requests. fn handle_stop(&mut self) { // Drain the stop_rx queue @@ -232,16 +283,26 @@ impl ClaudeWorker { } /// Perform compaction by summarizing messages and storing the result. - async fn handle_compact(&mut self) { + /// + /// 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() { return; } + // Update timestamp for automatic compactions (before the API call) + if is_automatic { + self.last_auto_compaction = Some(std::time::Instant::now()); + } + let num_messages = self.messages.len(); let buffer_chars = self.message_buffer_chars(); warn!( - "Starting compaction: {} messages, {} chars", - num_messages, buffer_chars + "Starting {} compaction: {} messages, {} chars", + if is_automatic { "automatic" } else { "manual" }, + num_messages, + buffer_chars ); // Build system content: compaction prompt first, then other system prompts, then existing summary diff --git a/signal-gateway/src/gateway/mod.rs b/signal-gateway/src/gateway/mod.rs index 0c396ff..a6e1d94 100644 --- a/signal-gateway/src/gateway/mod.rs +++ b/signal-gateway/src/gateway/mod.rs @@ -21,7 +21,10 @@ use http_body::Body; use http_body_util::BodyExt; use prometheus_http_client::{AlertStatus, ExtractLabels}; use std::{ - fmt::Write, net::SocketAddr, path::PathBuf, sync::Arc, sync::Mutex, sync::OnceLock, sync::Weak, + fmt::Write, + net::SocketAddr, + path::PathBuf, + sync::{Arc, Mutex, OnceLock, Weak}, time::Duration, }; use tokio::{ diff --git a/signal-gateway/src/gateway/signal_trust_set.rs b/signal-gateway/src/gateway/signal_trust_set.rs index f566c56..70921b1 100644 --- a/signal-gateway/src/gateway/signal_trust_set.rs +++ b/signal-gateway/src/gateway/signal_trust_set.rs @@ -5,13 +5,11 @@ //! - Sequence: `["uuid1", "uuid2"]` - UUIDs with no safety numbers (simpler) use crate::signal_jsonrpc::{Envelope, Identity, RpcClient}; -use serde::de::{MapAccess, SeqAccess, Visitor}; -use serde::{Deserialize, Deserializer}; -use std::borrow::Borrow; -use std::collections::HashMap; -use std::fmt; -use std::ops::Deref; -use std::str::FromStr; +use serde::{ + Deserialize, Deserializer, + de::{MapAccess, SeqAccess, Visitor}, +}; +use std::{borrow::Borrow, collections::HashMap, fmt, ops::Deref, str::FromStr}; use tracing::{debug, info, warn}; /// A validated Signal UUID in the format `xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx`. diff --git a/signal-gateway/src/prometheus/mod.rs b/signal-gateway/src/prometheus/mod.rs index 145cdf1..dfeaece 100644 --- a/signal-gateway/src/prometheus/mod.rs +++ b/signal-gateway/src/prometheus/mod.rs @@ -5,8 +5,7 @@ use prometheus_http_client::{ AlertInfo, AlertsRequest, ExtractLabels, Labels, LabelsRequest, MetricVal, MetricValue, PromRequest, QueryRequest, ReqwestClient, SeriesRequest, }; -use std::error::Error; -use std::fmt::Write; +use std::{error::Error, fmt::Write}; use tracing::info; type BoxError = Box;