diff --git a/signal-gateway-assistant/claude/src/lib.rs b/signal-gateway-assistant/claude/src/lib.rs index f92ec3f..ecd556f 100644 --- a/signal-gateway-assistant/claude/src/lib.rs +++ b/signal-gateway-assistant/claude/src/lib.rs @@ -175,7 +175,7 @@ impl ClaudeAssistant { }; if can_compact { - self.do_compact(true).await; + self.do_compact(true, None).await; } else { self.drop_oldest_messages(); } @@ -216,7 +216,9 @@ impl ClaudeAssistant { } /// Perform compaction by summarizing messages and storing the result. - async fn do_compact(&mut self, is_automatic: bool) { + /// + /// If `cancel` is provided, the operation can be interrupted. + async fn do_compact(&mut self, is_automatic: bool, cancel: Option<&CancellationToken>) { if self.messages.is_empty() { return; } @@ -257,15 +259,26 @@ impl ClaudeAssistant { tools: Vec::new(), }; - let result = self + let http_fut = self .client .post(&self.config.claude_api_url) .header("x-api-key", &self.api_key) .header("anthropic-version", ANTHROPIC_API_VERSION) .header("content-type", "application/json") .json(&request_body) - .send() - .await; + .send(); + + let result = if let Some(cancel) = cancel { + tokio::select! { + result = http_fut => result, + _ = cancel.cancelled() => { + info!("Compaction cancelled"); + return; + } + } + } else { + http_fut.await + }; match result { Ok(response) if response.status().is_success() => { @@ -469,9 +482,9 @@ impl Assistant for ClaudeAssistant { Ok(self.handle_request(&cancel).await?) } - async fn compact(&mut self) { + async fn compact(&mut self, cancel: CancellationToken) { // Manual compaction always runs (is_automatic = false) - self.do_compact(false).await; + self.do_compact(false, Some(&cancel)).await; } fn debug_log(&mut self) { diff --git a/signal-gateway-assistant/src/assistant.rs b/signal-gateway-assistant/src/assistant.rs index 66e0aca..9dadcb2 100644 --- a/signal-gateway-assistant/src/assistant.rs +++ b/signal-gateway-assistant/src/assistant.rs @@ -25,7 +25,9 @@ pub trait Assistant: Send { /// /// This is called when the user explicitly requests compaction. /// Automatic compaction is an internal implementation detail. - async fn compact(&mut self); + /// + /// The cancellation token can be used to interrupt the operation. + async fn compact(&mut self, cancel: CancellationToken); /// Log the assistant's current state for debugging. fn debug_log(&mut self); diff --git a/signal-gateway/src/assistant/worker.rs b/signal-gateway/src/assistant/worker.rs index 1cb237a..fcaf1b5 100644 --- a/signal-gateway/src/assistant/worker.rs +++ b/signal-gateway/src/assistant/worker.rs @@ -117,7 +117,17 @@ impl AssistantWorker { false } Input::Compact => { - self.assistant.compact().await; + let cancel_token = self.cancellation_token.child_token(); + let mut compact_fut = self.assistant.compact(cancel_token.clone()); + + tokio::select! { + _ = &mut compact_fut => {}, + _ = self.stop_rx.recv() => { + cancel_token.cancel(); + compact_fut.await; + return true; + } + }; false } Input::Debug => {