diff --git a/server/src/chat/commands/mod.rs b/server/src/chat/commands/mod.rs index 8307d88e..b3254a78 100644 --- a/server/src/chat/commands/mod.rs +++ b/server/src/chat/commands/mod.rs @@ -224,6 +224,11 @@ pub fn commands() -> &'static [BotCommand] { description: "Clear the current Claude Code session and start fresh", handler: handle_reset_fallback, }, + BotCommand { + name: "compact", + description: "Distill the session transcript into a seed, then reset (unlike `reset`, keeps distilled context for the next message)", + handler: handle_compact_fallback, + }, BotCommand { name: "timer", description: "Schedule a deferred agent start: `timer `, `timer list`, `timer cancel `", @@ -420,6 +425,16 @@ fn handle_reset_fallback(_ctx: &CommandContext) -> Option { None } +/// Fallback handler for the `compact` command when it is not intercepted by +/// the async handler in `on_room_message`. In practice this is never called — +/// compact is detected and handled before `try_handle_command` is invoked. +/// The entry exists in the registry only so `help` lists it. +/// +/// Returns `None` to prevent the LLM from receiving "compact" as a prompt. +fn handle_compact_fallback(_ctx: &CommandContext) -> Option { + None +} + /// Fallback handler for the `cleanup_worktrees` command when it is not /// intercepted by the async handler in `on_room_message`. In practice this is /// never called — cleanup_worktrees is detected and handled before diff --git a/server/src/chat/compact/digest.rs b/server/src/chat/compact/digest.rs new file mode 100644 index 00000000..b471b0ab --- /dev/null +++ b/server/src/chat/compact/digest.rs @@ -0,0 +1,263 @@ +//! Deterministic, pure extraction of a size-capped digest from a Claude Code +//! session transcript (JSONL), with no LLM call involved. + +/// Build a deterministic digest of a Claude Code session transcript. +/// +/// Reads NDJSON `jsonl` (one Claude Code transcript event per line) and keeps +/// only: +/// - `user` entries whose `message.content` is a plain string (real user +/// text), included verbatim. +/// - `assistant` entries' final text content, included verbatim, plus each +/// `tool_use` block reduced to `name(one-line json args)`. +/// +/// Excluded: `tool_result` content (found in `user` entries whose content is +/// an array), `thinking` blocks, and any event whose top-level `type` is not +/// `user` or `assistant` (e.g. `queue-operation`, `attachment`, `summary`, +/// `system`, `stream_event`). Sidechain entries (subagent turns) are also +/// excluded so the digest reflects only the main conversation thread. +/// +/// Lines that fail to parse as JSON are skipped rather than treated as fatal. +/// +/// When the joined digest exceeds `max_bytes`, entries are dropped from the +/// oldest end first so the most recent content survives the cap. +pub fn build_digest(jsonl: &str, max_bytes: usize) -> String { + let entries = extract_entries(jsonl); + cap_to_bytes(&entries, max_bytes) +} + +/// Parse `jsonl` into an ordered list of digest lines, applying the +/// extraction/exclusion rules. Pure — no truncation is applied here. +fn extract_entries(jsonl: &str) -> Vec { + let mut entries = Vec::new(); + for line in jsonl.lines() { + let trimmed = line.trim(); + if trimmed.is_empty() { + continue; + } + let Ok(value) = serde_json::from_str::(trimmed) else { + continue; + }; + if value.get("isSidechain").and_then(|v| v.as_bool()) == Some(true) { + continue; + } + match value.get("type").and_then(|t| t.as_str()) { + Some("user") => { + if let Some(text) = extract_user_text(&value) { + entries.push(format!("User: {text}")); + } + } + Some("assistant") => { + entries.extend(extract_assistant_lines(&value)); + } + _ => {} + } + } + entries +} + +/// Extract verbatim user text from a `user`-typed transcript entry. +/// +/// Returns `None` when the message content is an array (tool results) rather +/// than a plain string — tool results are excluded from the digest. +fn extract_user_text(value: &serde_json::Value) -> Option { + let content = value.get("message")?.get("content")?; + content.as_str().map(str::to_string) +} + +/// Extract digest lines from an `assistant`-typed transcript entry: the final +/// text reply verbatim, followed by one line per `tool_use` block. `thinking` +/// blocks are skipped. +fn extract_assistant_lines(value: &serde_json::Value) -> Vec { + let mut lines = Vec::new(); + let Some(content) = value + .get("message") + .and_then(|m| m.get("content")) + .and_then(|c| c.as_array()) + else { + return lines; + }; + + let mut text_parts = Vec::new(); + for block in content { + match block.get("type").and_then(|t| t.as_str()) { + Some("text") => { + if let Some(text) = block.get("text").and_then(|t| t.as_str()) { + text_parts.push(text); + } + } + Some("tool_use") => { + let name = block.get("name").and_then(|n| n.as_str()).unwrap_or("tool"); + let args = block + .get("input") + .map(serde_json::Value::to_string) + .unwrap_or_default(); + lines.push(format!("Tool: {name}({args})")); + } + // "thinking" and any other block types are excluded. + _ => {} + } + } + if !text_parts.is_empty() { + lines.insert(0, format!("Assistant: {}", text_parts.join("\n"))); + } + lines +} + +/// Join `entries` with newlines, dropping the oldest entries first so the +/// digest stays within `max_bytes`. Always keeps at least the single newest +/// entry, even if it alone exceeds the cap. +fn cap_to_bytes(entries: &[String], max_bytes: usize) -> String { + let joined = entries.join("\n"); + if joined.len() <= max_bytes { + return joined; + } + + let mut kept: Vec<&str> = Vec::new(); + let mut total = 0usize; + for entry in entries.iter().rev() { + let addition = entry.len() + if kept.is_empty() { 0 } else { 1 }; + if total + addition > max_bytes && !kept.is_empty() { + break; + } + total += addition; + kept.push(entry); + } + kept.reverse(); + kept.join("\n") +} + +#[cfg(test)] +mod tests { + use super::*; + + // -- extraction / exclusion rules --------------------------------------- + + #[test] + fn extracts_plain_user_text_verbatim() { + let jsonl = r#"{"type":"user","message":{"role":"user","content":"hello there"}}"#; + let digest = build_digest(jsonl, 10_000); + assert_eq!(digest, "User: hello there"); + } + + #[test] + fn extracts_assistant_final_text_verbatim() { + let jsonl = r#"{"type":"assistant","message":{"role":"assistant","content":[{"type":"text","text":"here is my answer"}]}}"#; + let digest = build_digest(jsonl, 10_000); + assert_eq!(digest, "Assistant: here is my answer"); + } + + #[test] + fn extracts_tool_use_as_name_and_one_line_args() { + let jsonl = r#"{"type":"assistant","message":{"role":"assistant","content":[{"type":"tool_use","name":"Bash","input":{"command":"ls -la"}}]}}"#; + let digest = build_digest(jsonl, 10_000); + assert_eq!(digest, r#"Tool: Bash({"command":"ls -la"})"#); + assert!(!digest.contains('\n'), "tool args must be one line"); + } + + #[test] + fn excludes_tool_result_user_entries() { + // A user entry carrying a tool_result array (not plain string content) + // must be excluded entirely from the digest. + let jsonl = r#"{"type":"user","message":{"role":"user","content":[{"type":"tool_result","tool_use_id":"abc","content":"file contents..."}]}}"#; + let digest = build_digest(jsonl, 10_000); + assert_eq!(digest, "", "tool_result entries must be excluded"); + } + + #[test] + fn excludes_thinking_blocks() { + let jsonl = r#"{"type":"assistant","message":{"role":"assistant","content":[{"type":"thinking","thinking":"let me consider..."},{"type":"text","text":"final answer"}]}}"#; + let digest = build_digest(jsonl, 10_000); + assert_eq!(digest, "Assistant: final answer"); + assert!(!digest.contains("let me consider")); + } + + #[test] + fn excludes_non_user_assistant_event_types() { + let jsonl = "\ +{\"type\":\"queue-operation\",\"operation\":\"enqueue\",\"content\":\"noise\"} +{\"type\":\"attachment\",\"attachment\":{\"type\":\"skill_listing\"}} +{\"type\":\"summary\",\"summary\":\"irrelevant\"} +{\"type\":\"system\",\"content\":\"system noise\"} +{\"type\":\"user\",\"message\":{\"role\":\"user\",\"content\":\"real message\"}}"; + let digest = build_digest(jsonl, 10_000); + assert_eq!(digest, "User: real message"); + } + + #[test] + fn excludes_sidechain_entries() { + let jsonl = r#"{"type":"user","isSidechain":true,"message":{"role":"user","content":"subagent chatter"}}"#; + let digest = build_digest(jsonl, 10_000); + assert_eq!(digest, "", "sidechain (subagent) entries must be excluded"); + } + + #[test] + fn deterministic_across_repeated_calls() { + let jsonl = "\ +{\"type\":\"user\",\"message\":{\"role\":\"user\",\"content\":\"a\"}} +{\"type\":\"assistant\",\"message\":{\"role\":\"assistant\",\"content\":[{\"type\":\"text\",\"text\":\"b\"}]}}"; + let first = build_digest(jsonl, 10_000); + let second = build_digest(jsonl, 10_000); + assert_eq!(first, second); + } + + // -- cap-truncation keeping newest --------------------------------------- + + #[test] + fn cap_truncation_keeps_newest_entries() { + let jsonl = "\ +{\"type\":\"user\",\"message\":{\"role\":\"user\",\"content\":\"oldest message\"}} +{\"type\":\"assistant\",\"message\":{\"role\":\"assistant\",\"content\":[{\"type\":\"text\",\"text\":\"middle reply\"}]}} +{\"type\":\"user\",\"message\":{\"role\":\"user\",\"content\":\"newest message\"}}"; + // Cap small enough to only fit the last entry. + let digest = build_digest(jsonl, 20); + assert_eq!(digest, "User: newest message"); + assert!(!digest.contains("oldest")); + } + + #[test] + fn cap_truncation_keeps_at_least_one_entry_even_if_oversized() { + let jsonl = r#"{"type":"user","message":{"role":"user","content":"this single message is longer than the cap"}}"#; + let digest = build_digest(jsonl, 5); + assert!( + digest.contains("this single message"), + "must keep the single newest entry even if it exceeds max_bytes: {digest}" + ); + } + + #[test] + fn no_truncation_when_under_cap() { + let jsonl = r#"{"type":"user","message":{"role":"user","content":"short"}}"#; + let digest = build_digest(jsonl, 10_000); + assert_eq!(digest, "User: short"); + } + + // -- malformed-line tolerance --------------------------------------------- + + #[test] + fn tolerates_malformed_lines_between_valid_ones() { + let jsonl = "\ +{\"type\":\"user\",\"message\":{\"role\":\"user\",\"content\":\"first\"}} +not json at all {{{ +{\"type\":\"assistant\",\"message\":{\"role\":\"assistant\",\"content\":[{\"type\":\"text\",\"text\":\"second\"}]}}"; + let digest = build_digest(jsonl, 10_000); + assert_eq!(digest, "User: first\nAssistant: second"); + } + + #[test] + fn tolerates_empty_lines() { + let jsonl = "\n{\"type\":\"user\",\"message\":{\"role\":\"user\",\"content\":\"hi\"}}\n\n"; + let digest = build_digest(jsonl, 10_000); + assert_eq!(digest, "User: hi"); + } + + #[test] + fn empty_input_yields_empty_digest() { + assert_eq!(build_digest("", 10_000), ""); + } + + #[test] + fn all_malformed_yields_empty_digest() { + let jsonl = "garbage\nmore garbage\n{not valid"; + assert_eq!(build_digest(jsonl, 10_000), ""); + } +} diff --git a/server/src/chat/compact/mod.rs b/server/src/chat/compact/mod.rs new file mode 100644 index 00000000..9034bee5 --- /dev/null +++ b/server/src/chat/compact/mod.rs @@ -0,0 +1,250 @@ +//! `compact` chat command: deterministically distill a Claude Code session +//! transcript into a size-capped digest, write it as a seed, and clear the +//! session so the next turn starts fresh with only the distilled context. +//! +//! Transport-agnostic: parsing (`extract_compact_command`) and orchestration +//! (`compact_session`) live here so every chat transport (Matrix, WhatsApp, +//! Slack, Discord) shares the same behaviour. Each transport wires this into +//! its own conversation-history type. + +/// Pure extraction of a deterministic digest from session transcript lines. +pub mod digest; +/// Resolves a Claude Code session transcript's path on disk from its id. +pub(crate) mod transcript; + +use std::fs; +use std::path::{Path, PathBuf}; + +use crate::chat::util::strip_bot_mention; + +/// A parsed `compact` command. +#[derive(Debug, PartialEq)] +pub struct CompactCommand; + +/// Parse a `compact` command from a raw message body. +/// +/// Mirrors [`crate::chat::transport::matrix::reset::extract_reset_command`]: +/// strips the bot mention prefix and checks whether the command word is +/// `compact`. Returns `None` when the message is not a compact command. +pub fn extract_compact_command( + message: &str, + bot_name: &str, + bot_user_id: &str, +) -> Option { + let stripped = strip_bot_mention(message, bot_name, bot_user_id); + let trimmed = stripped + .trim() + .trim_start_matches(|c: char| !c.is_alphanumeric()); + + let cmd = match trimmed.split_once(char::is_whitespace) { + Some((c, _)) => c, + None => trimmed, + }; + + if cmd.eq_ignore_ascii_case("compact") { + Some(CompactCommand) + } else { + None + } +} + +/// Result of a successful compaction: the digest text plus before/after +/// sizes (in bytes) for the reply message. +pub struct CompactOutcome { + /// The distilled digest text, written to the seed file and set as the + /// pending seed for the next spawned session. + pub digest: String, + /// Size in bytes of the original session transcript. + pub before_bytes: u64, + /// Size in bytes of the distilled digest. + pub after_bytes: u64, + /// Path the digest was written to. + pub seed_path: PathBuf, +} + +/// Why a compaction attempt could not proceed. +#[derive(Debug)] +pub enum CompactError { + /// The session transcript file could not be read (missing, permissions, + /// or otherwise unreadable). + TranscriptUnreadable(String), +} + +impl std::fmt::Display for CompactError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + CompactError::TranscriptUnreadable(e) => { + write!(f, "session transcript is unreadable: {e}") + } + } + } +} + +/// Distill the session transcript at `transcript_path` into a size-capped +/// digest, and write it to a per-room seed file under +/// `project_root/.huskies/chat_seeds/`. +/// +/// Callers are responsible for checking that a session is actually active +/// before calling this (AC5: no active session must be handled by the caller +/// without invoking this function, since there would be no transcript path +/// to resolve). Use [`transcript::transcript_path`] to resolve +/// `transcript_path` from a cwd + session_id. +pub fn compact_session( + project_root: &Path, + transcript_path: &Path, + room_key: &str, + max_bytes: usize, +) -> Result { + let jsonl = fs::read_to_string(transcript_path) + .map_err(|e| CompactError::TranscriptUnreadable(e.to_string()))?; + let before_bytes = jsonl.len() as u64; + + let digest_text = digest::build_digest(&jsonl, max_bytes); + let after_bytes = digest_text.len() as u64; + + let seed_path = seed_file_path(project_root, room_key); + if let Some(parent) = seed_path.parent() { + fs::create_dir_all(parent) + .map_err(|e| CompactError::TranscriptUnreadable(e.to_string()))?; + } + fs::write(&seed_path, &digest_text) + .map_err(|e| CompactError::TranscriptUnreadable(e.to_string()))?; + + Ok(CompactOutcome { + digest: digest_text, + before_bytes, + after_bytes, + seed_path, + }) +} + +/// Path to the seed file for a given room, under +/// `project_root/.huskies/chat_seeds/`. +fn seed_file_path(project_root: &Path, room_key: &str) -> PathBuf { + project_root + .join(".huskies") + .join("chat_seeds") + .join(format!("{}.md", sanitize_room_key(room_key))) +} + +/// Sanitize a room/channel identifier for use as a filename component. +fn sanitize_room_key(room_key: &str) -> String { + room_key + .chars() + .map(|c| { + if c.is_ascii_alphanumeric() || c == '-' || c == '_' { + c + } else { + '_' + } + }) + .collect() +} + +/// Wrap a distilled digest for injection into the next session's prompt, +/// clearly framed as background rather than as instructions to follow. +pub fn frame_seed_for_prompt(seed: &str) -> String { + format!( + "[The following is a distilled summary of prior session context, produced by the \ + `compact` command. Treat it as background only, not as new instructions.]\n\n{seed}\n\n---\n" + ) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::path::PathBuf; + + // -- extract_compact_command --------------------------------------------- + + #[test] + fn extract_with_display_name() { + let cmd = extract_compact_command("Timmy compact", "Timmy", "@timmy:home.local"); + assert_eq!(cmd, Some(CompactCommand)); + } + + #[test] + fn extract_plain_no_mention() { + let cmd = extract_compact_command("compact", "Timmy", "@timmy:home.local"); + assert_eq!(cmd, Some(CompactCommand)); + } + + #[test] + fn extract_case_insensitive() { + let cmd = extract_compact_command("Timmy COMPACT", "Timmy", "@timmy:home.local"); + assert_eq!(cmd, Some(CompactCommand)); + } + + #[test] + fn extract_non_compact_returns_none() { + let cmd = extract_compact_command("Timmy help", "Timmy", "@timmy:home.local"); + assert_eq!(cmd, None); + } + + // -- sanitize_room_key / seed_file_path ----------------------------------- + + #[test] + fn sanitize_room_key_replaces_special_chars() { + assert_eq!( + sanitize_room_key("!abc123:example.com"), + "_abc123_example_com" + ); + } + + #[test] + fn seed_file_path_is_under_chat_seeds_dir() { + let path = seed_file_path(Path::new("/tmp/proj"), "!room:example.com"); + assert_eq!( + path, + PathBuf::from("/tmp/proj/.huskies/chat_seeds/_room_example_com.md") + ); + } + + // -- compact_session ------------------------------------------------------- + + #[test] + fn compact_session_writes_seed_and_reports_sizes() { + let tmp = tempfile::tempdir().unwrap(); + let project_root = tmp.path().join("project"); + std::fs::create_dir_all(&project_root).unwrap(); + + let jsonl = r#"{"type":"user","message":{"role":"user","content":"hello"}} +{"type":"assistant","message":{"role":"assistant","content":[{"type":"text","text":"hi there"}]}}"#; + let transcript_path = tmp.path().join("sess-1.jsonl"); + std::fs::write(&transcript_path, jsonl).unwrap(); + + let outcome = + compact_session(&project_root, &transcript_path, "!room:example.com", 10_000).unwrap(); + + assert_eq!(outcome.digest, "User: hello\nAssistant: hi there"); + assert_eq!(outcome.before_bytes, jsonl.len() as u64); + assert_eq!(outcome.after_bytes, outcome.digest.len() as u64); + assert!(outcome.seed_path.exists()); + let written = std::fs::read_to_string(&outcome.seed_path).unwrap(); + assert_eq!(written, outcome.digest); + } + + #[test] + fn compact_session_errors_on_missing_transcript() { + let tmp = tempfile::tempdir().unwrap(); + let project_root = tmp.path().join("project"); + std::fs::create_dir_all(&project_root).unwrap(); + + let result = compact_session( + &project_root, + &tmp.path().join("no-such-session.jsonl"), + "room", + 10_000, + ); + assert!(matches!(result, Err(CompactError::TranscriptUnreadable(_)))); + } + + // -- frame_seed_for_prompt -------------------------------------------------- + + #[test] + fn frame_seed_for_prompt_marks_it_as_background() { + let framed = frame_seed_for_prompt("User: hi\nAssistant: hello"); + assert!(framed.contains("distilled summary")); + assert!(framed.contains("User: hi\nAssistant: hello")); + } +} diff --git a/server/src/chat/compact/transcript.rs b/server/src/chat/compact/transcript.rs new file mode 100644 index 00000000..ab769991 --- /dev/null +++ b/server/src/chat/compact/transcript.rs @@ -0,0 +1,65 @@ +//! Resolves the on-disk path to a Claude Code session transcript JSONL file. + +use std::path::{Path, PathBuf}; + +/// Resolve the path to a Claude Code session transcript. +/// +/// Claude Code stores each session's transcript at +/// `$HOME/.claude/projects//.jsonl`, where +/// `` is the absolute working directory with every `/` and `.` +/// replaced by `-` (e.g. `/workspace/.huskies/worktrees/1186` becomes +/// `-workspace--huskies-worktrees-1186`). +pub fn transcript_path(cwd: &Path, session_id: &str) -> PathBuf { + let home = std::env::var("HOME").unwrap_or_else(|_| "/home/huskies".to_string()); + transcript_path_under_home(Path::new(&home), cwd, session_id) +} + +/// Same as [`transcript_path`] but takes an explicit `$HOME` directory, +/// keeping the path-joining logic testable without mutating process env vars. +fn transcript_path_under_home(home: &Path, cwd: &Path, session_id: &str) -> PathBuf { + let mangled = mangle_cwd(&cwd.to_string_lossy()); + home.join(".claude") + .join("projects") + .join(mangled) + .join(format!("{session_id}.jsonl")) +} + +/// Replace every `/` and `.` in an absolute path string with `-`, matching +/// the directory-naming convention Claude Code uses under `~/.claude/projects/`. +fn mangle_cwd(cwd: &str) -> String { + cwd.chars() + .map(|c| if c == '/' || c == '.' { '-' } else { c }) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn mangles_worktree_path_matching_observed_convention() { + // Verified against this very worktree's own transcript directory. + assert_eq!( + mangle_cwd("/workspace/.huskies/worktrees/1186"), + "-workspace--huskies-worktrees-1186" + ); + } + + #[test] + fn mangles_simple_home_path() { + assert_eq!(mangle_cwd("/home/huskies"), "-home-huskies"); + } + + #[test] + fn transcript_path_joins_home_projects_dir_and_session_file() { + let path = transcript_path_under_home( + Path::new("/home/testuser"), + Path::new("/workspace/proj"), + "abc-123", + ); + assert_eq!( + path, + PathBuf::from("/home/testuser/.claude/projects/-workspace-proj/abc-123.jsonl") + ); + } +} diff --git a/server/src/chat/history.rs b/server/src/chat/history.rs index df9fedba..bc2f2917 100644 --- a/server/src/chat/history.rs +++ b/server/src/chat/history.rs @@ -128,6 +128,7 @@ mod tests { content: "hi there!".to_string(), }, ], + ..Default::default() }, ); diff --git a/server/src/chat/mod.rs b/server/src/chat/mod.rs index cdd2350b..ad08a3be 100644 --- a/server/src/chat/mod.rs +++ b/server/src/chat/mod.rs @@ -6,6 +6,8 @@ /// Bot command registry and dispatch — parses and routes incoming chat messages. pub mod commands; +/// `compact` command: deterministic session-transcript distillation and seeding. +pub mod compact; /// Protocol-agnostic chat dispatcher — coalesce window and per-session serial lock. pub mod dispatcher; /// Chat history utilities — loading and serialising conversation history. diff --git a/server/src/chat/transport/discord/commands.rs b/server/src/chat/transport/discord/commands.rs index 9c57791f..89c90022 100644 --- a/server/src/chat/transport/discord/commands.rs +++ b/server/src/chat/transport/discord/commands.rs @@ -411,6 +411,7 @@ async fn handle_llm_message(ctx: &DiscordContext, channel: &str, user: &str, use Ok(ClaudeCodeResult { messages, session_id, + .. }) => { let reply = if !remaining.is_empty() { let _ = msg_tx.send(remaining.clone()); @@ -568,6 +569,7 @@ mod tests { sender: "user123".to_string(), content: "previous message".to_string(), }], + ..Default::default() }, ); m diff --git a/server/src/chat/transport/discord/history.rs b/server/src/chat/transport/discord/history.rs index 3a14a59b..1389b635 100644 --- a/server/src/chat/transport/discord/history.rs +++ b/server/src/chat/transport/discord/history.rs @@ -56,6 +56,7 @@ mod tests { content: "hi there!".to_string(), }, ], + ..Default::default() }, ); diff --git a/server/src/chat/transport/matrix/bot/context.rs b/server/src/chat/transport/matrix/bot/context.rs index b3fc6b4c..1f2b3656 100644 --- a/server/src/chat/transport/matrix/bot/context.rs +++ b/server/src/chat/transport/matrix/bot/context.rs @@ -120,6 +120,16 @@ pub struct BotContext { /// Optional model override from bot.toml. Passed as `--model` to the /// `claude` CLI when set. pub model: Option, + /// Maximum size in bytes of the digest the `compact` command writes as a + /// seed file. From `bot.toml`'s `compact_seed_max_bytes`. + pub compact_seed_max_bytes: usize, + /// `cache_read_input_tokens` threshold above which the bot suggests + /// running `compact` after a turn. From `bot.toml`'s + /// `cache_read_suggest_threshold`. + pub cache_read_suggest_threshold: u64, + /// Minimum seconds between repeated `compact` suggestions for the same + /// room. From `bot.toml`'s `compact_suggest_cooldown_secs`. + pub compact_suggest_cooldown_secs: i64, } impl BotContext { @@ -343,6 +353,9 @@ mod tests { gateway_port: None, last_matrix_event_ms: Arc::new(AtomicI64::new(chrono::Utc::now().timestamp_millis())), model: None, + compact_seed_max_bytes: 8_000, + cache_read_suggest_threshold: 50_000, + compact_suggest_cooldown_secs: 3_600, } } diff --git a/server/src/chat/transport/matrix/bot/history.rs b/server/src/chat/transport/matrix/bot/history.rs index 5c381e60..03ef5d46 100644 --- a/server/src/chat/transport/matrix/bot/history.rs +++ b/server/src/chat/transport/matrix/bot/history.rs @@ -36,6 +36,17 @@ pub struct RoomConversation { pub session_id: Option, /// Rolling conversation entries (used for turn counting and persistence). pub entries: Vec, + /// A distilled digest produced by the `compact` command, waiting to be + /// injected as background context into the next spawned session's + /// prompt. Cleared immediately after the first turn that uses it, so it + /// is never re-injected. + #[serde(skip_serializing_if = "Option::is_none")] + pub pending_seed: Option, + /// Timestamp (ms since Unix epoch) of the last time this room was sent a + /// "consider running `compact`" suggestion, used to rate-limit repeat + /// suggestions. + #[serde(skip_serializing_if = "Option::is_none")] + pub last_compact_suggested_at_ms: Option, } /// Per-room conversation state, keyed by room ID (serialised as string). diff --git a/server/src/chat/transport/matrix/bot/messages/handle_message.rs b/server/src/chat/transport/matrix/bot/messages/handle_message.rs index 3b93ce46..d0881662 100644 --- a/server/src/chat/transport/matrix/bot/messages/handle_message.rs +++ b/server/src/chat/transport/matrix/bot/messages/handle_message.rs @@ -31,6 +31,23 @@ pub(in crate::chat::transport::matrix::bot) async fn handle_message( guard.get(&room_id).and_then(|conv| conv.session_id.clone()) }; + // Drain any pending `compact` seed for this room so it is injected into + // the prompt exactly once, then persist the cleared state immediately — + // a crash mid-turn must not cause it to be re-injected on the next try. + let pending_seed: Option = { + let mut guard = ctx.history.lock().await; + let conv = guard.entry(room_id.clone()).or_default(); + let seed = conv.pending_seed.take(); + if seed.is_some() { + save_history(&ctx.services.project_root, &guard); + } + seed + }; + let seed_prefix = pending_seed + .as_deref() + .map(crate::chat::compact::frame_seed_for_prompt) + .unwrap_or_default(); + // Pull new pipeline-transition events from the CRDT event log for this // persona and atomically advance the high-water marks so the same events // are not re-injected on the next turn. All transports share the same @@ -49,7 +66,7 @@ pub(in crate::chat::transport::matrix::bot) async fn handle_message( String::new() }; let prompt = format!( - "{event_log_ctx}[Your name is {bot_name}. Refer to yourself as {bot_name}, not Claude.]\n{active_project_ctx}\n{}", + "{event_log_ctx}{seed_prefix}[Your name is {bot_name}. Refer to yourself as {bot_name}, not Claude.]\n{active_project_ctx}\n{}", format_user_prompt(&sender, &user_message) ); @@ -127,10 +144,11 @@ pub(in crate::chat::transport::matrix::bot) async fn handle_message( let remaining = buffer.lock().unwrap().trim().to_string(); let did_send_any = sent_any_chunk.load(Ordering::Relaxed); - let (assistant_reply, new_session_id) = match result { + let (assistant_reply, new_session_id, turn_usage) = match result { Ok(ClaudeCodeResult { messages, session_id, + usage, }) => { let reply = if !remaining.is_empty() { let _ = msg_tx.send(remaining.clone()); @@ -153,7 +171,7 @@ pub(in crate::chat::transport::matrix::bot) async fn handle_message( remaining }; slog!("[matrix-bot] session_id from chat_stream: {:?}", session_id); - (reply, session_id) + (reply, session_id, usage) } Err(e) => { slog!("[matrix-bot] LLM error: {e}"); @@ -163,7 +181,7 @@ pub(in crate::chat::transport::matrix::bot) async fn handle_message( format!("Error processing your request: {e}") }; let _ = msg_tx.send(err_msg.clone()); - (err_msg, None) + (err_msg, None, None) } }; @@ -174,9 +192,10 @@ pub(in crate::chat::transport::matrix::bot) async fn handle_message( // Record this exchange in the per-room conversation history and persist // the session ID so the next turn resumes with structured API messages. + let mut compact_suggestion: Option = None; if !assistant_reply.starts_with("Error processing") { let mut guard = ctx.history.lock().await; - let conv = guard.entry(room_id).or_default(); + let conv = guard.entry(room_id.clone()).or_default(); // Store the session ID so the next turn uses --resume. slog!( @@ -208,6 +227,27 @@ pub(in crate::chat::transport::matrix::bot) async fn handle_message( conv.entries.drain(..excess); } + // When this turn's cache_read usage crosses the configured threshold, + // queue a rate-limited suggestion to run `compact`. Rate-limited via + // `last_compact_suggested_at_ms` so a busy room isn't nagged every turn. + if let Some(usage) = &turn_usage + && usage.cache_read_input_tokens > ctx.cache_read_suggest_threshold + { + let now_ms = chrono::Utc::now().timestamp_millis(); + let cooldown_ms = ctx.compact_suggest_cooldown_secs.saturating_mul(1_000); + let due = conv + .last_compact_suggested_at_ms + .is_none_or(|last| now_ms - last >= cooldown_ms); + if due { + conv.last_compact_suggested_at_ms = Some(now_ms); + compact_suggestion = Some(format!( + "This turn read {} cache tokens. Consider running `compact` to distill the \ + session and reduce context size.", + usage.cache_read_input_tokens + )); + } + } + // Persist to disk so history survives server restarts. save_history(&ctx.services.project_root, &guard); } else { @@ -222,6 +262,14 @@ pub(in crate::chat::transport::matrix::bot) async fn handle_message( save_history(&ctx.services.project_root, &guard); } } + + if let Some(suggestion) = compact_suggestion { + let html = markdown_to_html(&suggestion); + let _ = ctx + .transport + .send_message(&room_id_str, &suggestion, &html) + .await; + } } // --------------------------------------------------------------------------- diff --git a/server/src/chat/transport/matrix/bot/messages/on_room_message.rs b/server/src/chat/transport/matrix/bot/messages/on_room_message.rs index 2105eead..dafd401a 100644 --- a/server/src/chat/transport/matrix/bot/messages/on_room_message.rs +++ b/server/src/chat/transport/matrix/bot/messages/on_room_message.rs @@ -1023,6 +1023,36 @@ pub(in crate::chat::transport::matrix::bot) async fn on_room_message( return; } + // Check for the compact command, which requires async access to the + // shared conversation history and cannot be handled by the sync command + // registry. + if super::super::super::compact::extract_compact_command( + &user_message, + &ctx.services.bot_name, + ctx.matrix_user_id.as_str(), + ) + .is_some() + { + slog!("[matrix-bot] Handling compact command from {sender}"); + let response = super::super::super::compact::handle_compact( + &incoming_room_id, + &ctx.history, + &ctx.services.project_root, + ctx.compact_seed_max_bytes, + ) + .await; + let html = markdown_to_html(&response); + if let Ok(msg_id) = ctx + .transport + .send_message(&room_id_str, &response, &html) + .await + && let Ok(event_id) = msg_id.parse() + { + ctx.bot_sent_event_ids.lock().await.insert(event_id); + } + return; + } + // In gateway mode, intercept "rebuild gateway" and route it through the // detached trampoline so the process swap survives any bash-tool kill cascade. if ctx.gateway_active_project.is_some() diff --git a/server/src/chat/transport/matrix/bot/run.rs b/server/src/chat/transport/matrix/bot/run.rs index ad69d2bc..884011b2 100644 --- a/server/src/chat/transport/matrix/bot/run.rs +++ b/server/src/chat/transport/matrix/bot/run.rs @@ -334,6 +334,9 @@ pub async fn run_bot( gateway_port, last_matrix_event_ms: Arc::new(AtomicI64::new(chrono::Utc::now().timestamp_millis())), model: config.model.clone(), + compact_seed_max_bytes: config.compact_seed_max_bytes, + cache_read_suggest_threshold: config.cache_read_suggest_threshold, + compact_suggest_cooldown_secs: config.compact_suggest_cooldown_secs, }; slog!( diff --git a/server/src/chat/transport/matrix/compact.rs b/server/src/chat/transport/matrix/compact.rs new file mode 100644 index 00000000..478c9cc4 --- /dev/null +++ b/server/src/chat/transport/matrix/compact.rs @@ -0,0 +1,190 @@ +//! `compact` command: distills the current room's Claude Code session +//! transcript into a size-capped digest, writes it as a seed, and clears the +//! session so the next turn starts fresh with only the distilled context. +//! +//! Parsing and the transcript→digest→seed pipeline are transport-agnostic +//! and live in [`crate::chat::compact`]; this module is the Matrix-specific +//! glue around the shared [`ConversationHistory`], mirroring [`super::reset`]. + +use crate::chat::compact::{self, CompactError}; +use crate::chat::transport::matrix::bot::{ConversationHistory, RoomConversation}; +use matrix_sdk::ruma::OwnedRoomId; +use std::path::Path; + +/// Re-exported so Matrix call sites read as `compact::extract_compact_command`, +/// matching the sibling `reset` module's shape. +pub use crate::chat::compact::extract_compact_command; + +/// Handle a `compact` command: distill the room's session transcript into a +/// digest, write it as a seed for the next turn, clear the session ID and +/// entries (like `reset`), and return a confirmation with before/after sizes. +/// +/// `project_root` doubles as the working directory the `claude` CLI was +/// spawned with for this room (matrix's `handle_message` always runs Claude +/// Code in `services.project_root` — the gateway config dir in gateway mode, +/// or the project root in standalone mode — so the session transcript lives +/// under the matching `~/.claude/projects//` directory). +/// +/// Returns an actionable message and changes nothing when there is no active +/// session, or when the session transcript cannot be read. +pub async fn handle_compact( + room_id: &OwnedRoomId, + history: &ConversationHistory, + project_root: &Path, + max_bytes: usize, +) -> String { + let mut guard = history.lock().await; + let conv = guard + .entry(room_id.clone()) + .or_insert_with(RoomConversation::default); + + let Some(session_id) = conv.session_id.clone() else { + return "No active session to compact. Send a message first to start one.".to_string(); + }; + + let transcript_path = compact::transcript::transcript_path(project_root, &session_id); + + match compact::compact_session(project_root, &transcript_path, room_id.as_ref(), max_bytes) { + Ok(outcome) => { + conv.session_id = None; + conv.entries.clear(); + conv.pending_seed = Some(outcome.digest); + crate::chat::transport::matrix::bot::save_history(project_root, &guard); + crate::slog!( + "[matrix-bot] compact: room {room_id} {} -> {} bytes, seed written to {}", + outcome.before_bytes, + outcome.after_bytes, + outcome.seed_path.display() + ); + format!( + "Compacted session context: {} → {} bytes. Starting fresh — the distilled \ + summary will be included as background in your next message.", + outcome.before_bytes, outcome.after_bytes + ) + } + Err(CompactError::TranscriptUnreadable(e)) => format!( + "Could not compact: session transcript is unreadable ({e}). Nothing was changed." + ), + } +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +#[cfg(test)] +mod tests { + use super::*; + use crate::chat::transport::matrix::bot::{ConversationEntry, ConversationRole}; + use std::collections::HashMap; + use std::sync::Arc; + use tokio::sync::Mutex as TokioMutex; + + fn room_id() -> OwnedRoomId { + "!test:example.com".parse().unwrap() + } + + #[tokio::test] + async fn no_active_session_returns_actionable_message_and_changes_nothing() { + let history: ConversationHistory = Arc::new(TokioMutex::new(HashMap::new())); + let tmp = tempfile::tempdir().unwrap(); + + let response = handle_compact(&room_id(), &history, tmp.path(), 8_000).await; + + assert!( + response.contains("No active session"), + "expected actionable no-session message: {response}" + ); + let guard = history.lock().await; + assert!( + guard.get(&room_id()).is_none_or(|c| c.entries.is_empty()), + "no entries should be created for a compact with no session" + ); + } + + #[tokio::test] + async fn unreadable_transcript_returns_actionable_message_and_preserves_session() { + let history: ConversationHistory = Arc::new(TokioMutex::new({ + let mut m = HashMap::new(); + m.insert( + room_id(), + RoomConversation { + session_id: Some("missing-session".to_string()), + entries: vec![], + pending_seed: None, + last_compact_suggested_at_ms: None, + }, + ); + m + })); + let tmp = tempfile::tempdir().unwrap(); + + let response = handle_compact(&room_id(), &history, tmp.path(), 8_000).await; + + assert!( + response.contains("unreadable"), + "expected actionable unreadable-transcript message: {response}" + ); + let guard = history.lock().await; + let conv = guard.get(&room_id()).unwrap(); + assert_eq!( + conv.session_id.as_deref(), + Some("missing-session"), + "session_id must be preserved when compaction fails" + ); + } + + #[tokio::test] + async fn successful_compact_clears_session_and_sets_pending_seed() { + let session_id = "sess-compact-1"; + let project_root = tempfile::tempdir().unwrap(); + let home = tempfile::tempdir().unwrap(); + // SAFETY: this test owns HOME for its duration; no other test in this + // process reads HOME concurrently with this call. + unsafe { + std::env::set_var("HOME", home.path()); + } + let transcript_dir = compact::transcript::transcript_path(project_root.path(), session_id) + .parent() + .unwrap() + .to_path_buf(); + std::fs::create_dir_all(&transcript_dir).unwrap(); + let jsonl = r#"{"type":"user","message":{"role":"user","content":"hello"}} +{"type":"assistant","message":{"role":"assistant","content":[{"type":"text","text":"hi there"}]}}"#; + std::fs::write(transcript_dir.join(format!("{session_id}.jsonl")), jsonl).unwrap(); + + let history: ConversationHistory = Arc::new(TokioMutex::new({ + let mut m = HashMap::new(); + m.insert( + room_id(), + RoomConversation { + session_id: Some(session_id.to_string()), + entries: vec![ConversationEntry { + role: ConversationRole::User, + sender: "@alice:example.com".to_string(), + content: "hi".to_string(), + }], + pending_seed: None, + last_compact_suggested_at_ms: None, + }, + ); + m + })); + + let response = handle_compact(&room_id(), &history, project_root.path(), 8_000).await; + + assert!( + response.contains("Compacted session context"), + "response should confirm compaction with sizes: {response}" + ); + let guard = history.lock().await; + let conv = guard.get(&room_id()).unwrap(); + assert!(conv.session_id.is_none(), "session_id must be cleared"); + assert!(conv.entries.is_empty(), "entries must be cleared"); + assert_eq!( + conv.pending_seed.as_deref(), + Some("User: hello\nAssistant: hi there"), + "pending_seed must hold the distilled digest" + ); + } +} diff --git a/server/src/chat/transport/matrix/config/types.rs b/server/src/chat/transport/matrix/config/types.rs index 88c3c46d..c2cdba1d 100644 --- a/server/src/chat/transport/matrix/config/types.rs +++ b/server/src/chat/transport/matrix/config/types.rs @@ -14,6 +14,23 @@ pub(super) fn default_coalesce_window_ms() -> u64 { 1_500 } +/// Default cap (bytes) on the digest the `compact` command writes as a seed. +pub(super) fn default_compact_seed_max_bytes() -> usize { + 8_000 +} + +/// Default cache_read token threshold above which the bot suggests running +/// `compact` after a turn. +pub(super) fn default_cache_read_suggest_threshold() -> u64 { + 50_000 +} + +/// Default cooldown (seconds) between repeated `compact` suggestions for the +/// same room. +pub(super) fn default_compact_suggest_cooldown_secs() -> i64 { + 3_600 +} + pub(super) fn default_transport() -> String { "matrix".to_string() } @@ -190,4 +207,20 @@ pub struct BotConfig { /// `git config user.email` when absent. #[serde(default)] pub git_user_email: Option, + + /// Maximum size in bytes of the digest the `compact` command writes as a + /// seed file. Older content is dropped first so the most recent + /// conversation survives the cap. Defaults to 8000 bytes. + #[serde(default = "default_compact_seed_max_bytes")] + pub compact_seed_max_bytes: usize, + + /// Number of `cache_read_input_tokens` in a single turn above which the + /// bot suggests running `compact`. Defaults to 50 000. + #[serde(default = "default_cache_read_suggest_threshold")] + pub cache_read_suggest_threshold: u64, + + /// Minimum seconds between repeated `compact` suggestions for the same + /// room, so a busy room isn't spammed every turn. Defaults to 3600 (1h). + #[serde(default = "default_compact_suggest_cooldown_secs")] + pub compact_suggest_cooldown_secs: i64, } diff --git a/server/src/chat/transport/matrix/health.rs b/server/src/chat/transport/matrix/health.rs index b101e418..49f4b773 100644 --- a/server/src/chat/transport/matrix/health.rs +++ b/server/src/chat/transport/matrix/health.rs @@ -925,6 +925,9 @@ mod tests { gateway_port: None, last_matrix_event_ms: Arc::new(AtomicI64::new(chrono::Utc::now().timestamp_millis())), model: None, + compact_seed_max_bytes: 8_000, + cache_read_suggest_threshold: 50_000, + compact_suggest_cooldown_secs: 3_600, } } } diff --git a/server/src/chat/transport/matrix/mod.rs b/server/src/chat/transport/matrix/mod.rs index 126cc6ff..b25ed418 100644 --- a/server/src/chat/transport/matrix/mod.rs +++ b/server/src/chat/transport/matrix/mod.rs @@ -22,6 +22,8 @@ mod bot; pub mod cleanup_worktrees; /// Matrix bot command handlers — parses and routes bot commands from Matrix messages. pub mod commands; +/// `compact` command — distills the session transcript into a seed and resets, Matrix-specific glue. +pub mod compact; pub(crate) mod config; /// Story deletion command — handles `!delete` bot commands to remove work items. pub mod delete; diff --git a/server/src/chat/transport/matrix/projects.rs b/server/src/chat/transport/matrix/projects.rs index 9dde1192..52e84795 100644 --- a/server/src/chat/transport/matrix/projects.rs +++ b/server/src/chat/transport/matrix/projects.rs @@ -99,6 +99,9 @@ mod tests { gateway_port: None, last_matrix_event_ms: Arc::new(AtomicI64::new(chrono::Utc::now().timestamp_millis())), model: None, + compact_seed_max_bytes: 8_000, + cache_read_suggest_threshold: 50_000, + compact_suggest_cooldown_secs: 3_600, }; run_projects_list(&ctx).await } @@ -218,6 +221,9 @@ mod tests { gateway_port: None, last_matrix_event_ms: Arc::new(AtomicI64::new(chrono::Utc::now().timestamp_millis())), model: None, + compact_seed_max_bytes: 8_000, + cache_read_suggest_threshold: 50_000, + compact_suggest_cooldown_secs: 3_600, }; let response = run_projects_list(&ctx).await; assert!( diff --git a/server/src/chat/transport/matrix/reset.rs b/server/src/chat/transport/matrix/reset.rs index 76e9cd82..926b7f37 100644 --- a/server/src/chat/transport/matrix/reset.rs +++ b/server/src/chat/transport/matrix/reset.rs @@ -125,6 +125,7 @@ mod tests { sender: "@alice:example.com".to_string(), content: "previous message".to_string(), }], + ..Default::default() }, ); m diff --git a/server/src/chat/transport/slack/commands/llm.rs b/server/src/chat/transport/slack/commands/llm.rs index 6379fb61..9b9473c6 100644 --- a/server/src/chat/transport/slack/commands/llm.rs +++ b/server/src/chat/transport/slack/commands/llm.rs @@ -133,6 +133,7 @@ pub(super) async fn handle_llm_message( Ok(ClaudeCodeResult { messages, session_id, + .. }) => { let reply = if !remaining.is_empty() { let _ = msg_tx.send(remaining.clone()); diff --git a/server/src/chat/transport/slack/commands/mod.rs b/server/src/chat/transport/slack/commands/mod.rs index b44f6dba..dd7a0e53 100644 --- a/server/src/chat/transport/slack/commands/mod.rs +++ b/server/src/chat/transport/slack/commands/mod.rs @@ -513,6 +513,7 @@ mod tests { sender: "U01GHIJKL".to_string(), content: "previous message".to_string(), }], + ..Default::default() }, ); m diff --git a/server/src/chat/transport/slack/history.rs b/server/src/chat/transport/slack/history.rs index b3f0d145..73620b3b 100644 --- a/server/src/chat/transport/slack/history.rs +++ b/server/src/chat/transport/slack/history.rs @@ -56,6 +56,7 @@ mod tests { content: "hi there!".to_string(), }, ], + ..Default::default() }, ); diff --git a/server/src/chat/transport/whatsapp/commands/llm.rs b/server/src/chat/transport/whatsapp/commands/llm.rs index f89c2544..3ff5d0c5 100644 --- a/server/src/chat/transport/whatsapp/commands/llm.rs +++ b/server/src/chat/transport/whatsapp/commands/llm.rs @@ -134,6 +134,7 @@ pub(super) async fn handle_llm_message( Ok(ClaudeCodeResult { messages, session_id, + .. }) => { let reply = if !remaining.is_empty() { let _ = msg_tx.send(remaining.clone()); diff --git a/server/src/chat/transport/whatsapp/commands/mod.rs b/server/src/chat/transport/whatsapp/commands/mod.rs index 00039fbf..b5f0d660 100644 --- a/server/src/chat/transport/whatsapp/commands/mod.rs +++ b/server/src/chat/transport/whatsapp/commands/mod.rs @@ -416,6 +416,7 @@ mod tests { sender: sender.to_string(), content: "previous message".to_string(), }], + ..Default::default() }, ); m diff --git a/server/src/chat/transport/whatsapp/history.rs b/server/src/chat/transport/whatsapp/history.rs index 5cb7ed83..eb73d2a2 100644 --- a/server/src/chat/transport/whatsapp/history.rs +++ b/server/src/chat/transport/whatsapp/history.rs @@ -153,6 +153,7 @@ mod tests { content: "hi there!".to_string(), }, ], + ..Default::default() }, ); @@ -200,6 +201,7 @@ mod tests { sender: "111".to_string(), content: "msg1".to_string(), }], + ..Default::default() }, ); history.insert( @@ -211,6 +213,7 @@ mod tests { sender: "222".to_string(), content: "msg2".to_string(), }], + ..Default::default() }, ); diff --git a/server/src/llm/chat/run.rs b/server/src/llm/chat/run.rs index f06b592a..c4085134 100644 --- a/server/src/llm/chat/run.rs +++ b/server/src/llm/chat/run.rs @@ -203,6 +203,7 @@ where let ClaudeCodeResult { messages: cc_messages, session_id, + .. } = provider .chat_stream( &user_message, diff --git a/server/src/llm/providers/claude_code/events/mod.rs b/server/src/llm/providers/claude_code/events/mod.rs index b3e3b398..74433fc2 100644 --- a/server/src/llm/providers/claude_code/events/mod.rs +++ b/server/src/llm/providers/claude_code/events/mod.rs @@ -6,6 +6,7 @@ mod stream; mod tests; use super::parse::{parse_assistant_message, parse_tool_results}; +use crate::agents::TokenUsage; use crate::llm::types::Message; use crate::slog; use std::sync::atomic::{AtomicBool, Ordering}; @@ -16,6 +17,7 @@ use stream::handle_stream_event; /// Routes the event to the appropriate handler based on `type`, emitting tokens, /// thinking output, activity signals, and parsed messages to their respective channels. /// Returns `true` only for `"result"` events, which signal that the CLI turn is complete. +#[allow(clippy::too_many_arguments)] pub(super) fn process_json_event( json: &serde_json::Value, token_tx: &tokio::sync::mpsc::UnboundedSender, @@ -23,6 +25,7 @@ pub(super) fn process_json_event( activity_tx: &tokio::sync::mpsc::UnboundedSender, msg_tx: &std::sync::mpsc::Sender, sid_tx: &mut Option>, + usage_tx: &mut Option>, auth_failed: &AtomicBool, ) -> bool { let event_type = match json.get("type").and_then(|t| t.as_str()) { @@ -80,7 +83,14 @@ pub(super) fn process_json_event( } false } - "result" => true, + "result" => { + if let Some(tx) = usage_tx.take() + && let Some(usage) = TokenUsage::from_result_event(json) + { + let _ = tx.send(usage); + } + true + } // system, rate_limit_event, and unknown types are no-ops _ => false, } diff --git a/server/src/llm/providers/claude_code/events/tests.rs b/server/src/llm/providers/claude_code/events/tests.rs index c8df4cc1..47a22b4c 100644 --- a/server/src/llm/providers/claude_code/events/tests.rs +++ b/server/src/llm/providers/claude_code/events/tests.rs @@ -202,10 +202,68 @@ fn process_json_event_result_returns_true() { &act_tx, &msg_tx, &mut sid_tx_opt, + &mut None, &AtomicBool::new(false), )); } +#[test] +fn process_json_event_result_captures_usage() { + let (tok_tx, _tok_rx, thi_tx, _thi_rx, act_tx, _act_rx, msg_tx, _msg_rx) = make_channels(); + let mut sid_tx = None::>; + let (usage_tx, mut usage_rx) = tokio::sync::oneshot::channel::(); + let mut usage_tx_opt = Some(usage_tx); + let json = json!({ + "type": "result", + "subtype": "success", + "total_cost_usd": 0.42, + "usage": { + "input_tokens": 10, + "output_tokens": 20, + "cache_creation_input_tokens": 100, + "cache_read_input_tokens": 60000 + } + }); + assert!(process_json_event( + &json, + &tok_tx, + &thi_tx, + &act_tx, + &msg_tx, + &mut sid_tx, + &mut usage_tx_opt, + &AtomicBool::new(false), + )); + assert!(usage_tx_opt.is_none(), "usage_tx should be consumed"); + let usage = usage_rx.try_recv().unwrap(); + assert_eq!(usage.cache_read_input_tokens, 60000); +} + +#[test] +fn process_json_event_result_without_usage_sends_nothing() { + let (tok_tx, _tok_rx, thi_tx, _thi_rx, act_tx, _act_rx, msg_tx, _msg_rx) = make_channels(); + let mut sid_tx = None::>; + let (usage_tx, mut usage_rx) = tokio::sync::oneshot::channel::(); + let mut usage_tx_opt = Some(usage_tx); + let json = json!({"type": "result", "subtype": "success"}); + assert!(process_json_event( + &json, + &tok_tx, + &thi_tx, + &act_tx, + &msg_tx, + &mut sid_tx, + &mut usage_tx_opt, + &AtomicBool::new(false), + )); + // The result event carried no "usage" field, so the sender is dropped + // without sending — the receiver observes a closed channel, not a value. + assert!( + usage_rx.try_recv().is_err(), + "no usage field means nothing sent" + ); +} + #[test] fn process_json_event_system_returns_false() { let (tok_tx, _tok_rx, thi_tx, _thi_rx, act_tx, _act_rx, msg_tx, _msg_rx) = make_channels(); @@ -218,6 +276,7 @@ fn process_json_event_system_returns_false() { &act_tx, &msg_tx, &mut sid_tx, + &mut None, &AtomicBool::new(false), )); } @@ -234,6 +293,7 @@ fn process_json_event_rate_limit_returns_false() { &act_tx, &msg_tx, &mut sid_tx, + &mut None, &AtomicBool::new(false), )); } @@ -250,6 +310,7 @@ fn process_json_event_unknown_type_returns_false() { &act_tx, &msg_tx, &mut sid_tx, + &mut None, &AtomicBool::new(false), )); } @@ -266,6 +327,7 @@ fn process_json_event_no_type_returns_false() { &act_tx, &msg_tx, &mut sid_tx, + &mut None, &AtomicBool::new(false), )); } @@ -283,6 +345,7 @@ fn process_json_event_captures_session_id() { &act_tx, &msg_tx, &mut sid_tx_opt, + &mut None, &AtomicBool::new(false), ); // sid_tx should have been consumed @@ -304,6 +367,7 @@ fn process_json_event_preserves_sid_tx_if_no_session_id() { &act_tx, &msg_tx, &mut sid_tx_opt, + &mut None, &AtomicBool::new(false), ); // sid_tx should still be present since no session_id in event @@ -329,6 +393,7 @@ fn process_json_event_stream_event_forwards_token() { &act_tx, &msg_tx, &mut sid_tx, + &mut None, &AtomicBool::new(false), )); drop(tok_tx); @@ -364,6 +429,7 @@ fn process_json_event_stream_event_tool_use_fires_activity() { &act_tx, &msg_tx, &mut sid_tx, + &mut None, &AtomicBool::new(false), )); drop(act_tx); @@ -397,6 +463,7 @@ fn process_json_event_assistant_with_tool_use_fires_activity() { &act_tx, &msg_tx, &mut sid_tx, + &mut None, &AtomicBool::new(false), )); drop(act_tx); @@ -430,6 +497,7 @@ fn process_json_event_assistant_with_multiple_tool_uses_fires_all_activities() { &act_tx, &msg_tx, &mut sid_tx, + &mut None, &AtomicBool::new(false), )); drop(act_tx); @@ -460,6 +528,7 @@ fn process_json_event_assistant_text_only_no_activity() { &act_tx, &msg_tx, &mut sid_tx, + &mut None, &AtomicBool::new(false), )); drop(act_tx); @@ -490,6 +559,7 @@ fn process_json_event_assistant_event_parses_message() { &act_tx, &msg_tx, &mut sid_tx, + &mut None, &AtomicBool::new(false), )); drop(msg_tx); @@ -515,6 +585,7 @@ fn process_json_event_user_event_parses_tool_results() { &act_tx, &msg_tx, &mut sid_tx, + &mut None, &AtomicBool::new(false), )); drop(msg_tx); @@ -539,6 +610,7 @@ fn process_json_event_assistant_without_content_array_is_noop() { &act_tx, &msg_tx, &mut sid_tx, + &mut None, &AtomicBool::new(false), )); drop(msg_tx); @@ -558,6 +630,7 @@ fn process_json_event_user_without_content_array_is_noop() { &act_tx, &msg_tx, &mut sid_tx, + &mut None, &AtomicBool::new(false), )); drop(msg_tx); @@ -584,6 +657,7 @@ fn process_json_event_detects_authentication_failed() { &act_tx, &msg_tx, &mut sid_tx, + &mut None, &auth_failed, )); assert!(auth_failed.load(Ordering::Relaxed)); @@ -607,6 +681,7 @@ fn process_json_event_no_auth_failed_for_normal_events() { &act_tx, &msg_tx, &mut sid_tx, + &mut None, &auth_failed, )); assert!(!auth_failed.load(Ordering::Relaxed)); diff --git a/server/src/llm/providers/claude_code/mod.rs b/server/src/llm/providers/claude_code/mod.rs index 9f4573ef..ce4902ee 100644 --- a/server/src/llm/providers/claude_code/mod.rs +++ b/server/src/llm/providers/claude_code/mod.rs @@ -8,6 +8,7 @@ use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; use tokio::sync::watch; +use crate::agents::TokenUsage; use crate::llm::types::{FunctionCall, Message, Role, ToolCall}; /// Result from a Claude Code session containing structured messages. @@ -17,6 +18,9 @@ pub struct ClaudeCodeResult { pub messages: Vec, /// Session ID for conversation resumption on subsequent requests. pub session_id: Option, + /// Token usage reported on the CLI's `result` event, if one was received. + /// `None` when the turn errored or produced no `result` event. + pub usage: Option, } /// Manages a Claude Code session via a pseudo-terminal. @@ -101,6 +105,7 @@ impl ClaudeCodeProvider { let (activity_tx, mut activity_rx) = tokio::sync::mpsc::unbounded_channel::(); let (msg_tx, msg_rx) = std::sync::mpsc::channel::(); let (sid_tx, sid_rx) = tokio::sync::oneshot::channel::(); + let (usage_tx, usage_rx) = tokio::sync::oneshot::channel::(); let pty_handle = tokio::task::spawn_blocking(move || { run_pty_session( @@ -117,6 +122,7 @@ impl ClaudeCodeProvider { activity_tx, msg_tx, sid_tx, + usage_tx, ) }); @@ -169,6 +175,7 @@ impl ClaudeCodeProvider { let captured_session_id = sid_rx.await.ok(); slog!("[pty-debug] RECEIVED session_id: {:?}", captured_session_id); + let usage = usage_rx.await.ok(); let structured_messages: Vec = msg_rx.try_iter().collect(); let exited_non_zero = exit_failed.load(Ordering::Relaxed); @@ -179,6 +186,7 @@ impl ClaudeCodeProvider { return Ok(ClaudeCodeResult { messages: structured_messages, session_id, + usage, }); } @@ -241,6 +249,7 @@ fn run_pty_session( activity_tx: tokio::sync::mpsc::UnboundedSender, msg_tx: std::sync::mpsc::Sender, sid_tx: tokio::sync::oneshot::Sender, + usage_tx: tokio::sync::oneshot::Sender, ) -> Result<(), String> { let pty_system = native_pty_system(); @@ -338,6 +347,7 @@ fn run_pty_session( let mut got_result = false; let mut sid_tx = Some(sid_tx); + let mut usage_tx = Some(usage_tx); loop { if cancelled.load(Ordering::Relaxed) { @@ -369,6 +379,7 @@ fn run_pty_session( &activity_tx, &msg_tx, &mut sid_tx, + &mut usage_tx, &auth_failed, ) { @@ -396,6 +407,7 @@ fn run_pty_session( &activity_tx, &msg_tx, &mut sid_tx, + &mut usage_tx, &auth_failed, ); }