From 79279a6bf34a82ff17e7b8d112b0b721d1913d97 Mon Sep 17 00:00:00 2001 From: jackie-cqz <2557911191@qq.com> Date: Wed, 21 Oct 2026 10:28:18 +0800 Subject: [PATCH 1/2] fix(zcode): retain per-turn usage and native wall time in history --- docs/agents/zcode.md | 2 + src-tauri/src/acp/connection.rs | 127 ++++++++++++++- src-tauri/src/acp/mod.rs | 1 + src-tauri/src/acp/registry.rs | 3 +- src-tauri/src/acp/zcode_turn.rs | 242 ++++++++++++++++++++++++++++ src-tauri/src/parsers/acp_native.rs | 60 ++++++- 6 files changed, 424 insertions(+), 11 deletions(-) create mode 100644 src-tauri/src/acp/zcode_turn.rs diff --git a/docs/agents/zcode.md b/docs/agents/zcode.md index 36a7334a22..8ad2e5c6f1 100644 --- a/docs/agents/zcode.md +++ b/docs/agents/zcode.md @@ -25,6 +25,8 @@ The bridge writes a disabled sandbox template to `.zcode/acp/sandbox.json` in ea Codeg records its own ACP transcripts for ZCode, including messages, thoughts, tools and usage. These transcripts support history and reconnecting to sessions created in Codeg. Existing custom ZCode entries keep their registration, launch configuration and history; the built-in option records its history separately. This integration does not scan ZCode's native SQLite sessions. +Completed replies retain the bridge's per-turn token usage (including cache reads and writes) and wall time, so their footer and session totals survive reopening. Context-window occupancy remains a separate reading, not additional billable tokens. When the bridge supplies no turn measurements, Codeg keeps its measured duration without inventing token counts. Older transcripts without those measurements are not retroactively filled. + Manage stdio MCP servers through Codeg's MCP settings and assign them to ZCode. The configuration is stored in Codeg's data directory and forwarded when creating or restoring a session; it is not written into official ZCode's provider or credential files. When multi-agent delegation is enabled, Codeg injects its per-launch `codeg-mcp` companion using the same path. ## Current scope diff --git a/src-tauri/src/acp/connection.rs b/src-tauri/src/acp/connection.rs index 044401c52b..6f4bdf958a 100644 --- a/src-tauri/src/acp/connection.rs +++ b/src-tauri/src/acp/connection.rs @@ -1811,6 +1811,7 @@ async fn record_turn_end( stop_reason: &str, started_at_ms: u64, model: Option, + metrics: Option<&crate::acp::zcode_turn::TurnMetrics>, ) { let Some(dir) = transcript_dir_for(agent_type) else { return; @@ -1820,6 +1821,9 @@ async fn record_turn_end( "stopReason": stop_reason, "durationMs": now.saturating_sub(started_at_ms), }); + if let Some(metrics) = metrics { + metrics.apply_to(&mut payload); + } // ACP puts no model on the prompt response, so the session's model selector // is the only honest answer at turn end — and it is the same value the // composer showed while the turn ran. Recorded per turn rather than once in @@ -13388,6 +13392,7 @@ async fn run_conversation_loop( // of a confusing `PendingReview` on a blank conversation — and // the probe's other fields say *why* it looked empty. let mut probe = TurnOutputProbe::new(stderr_mark); + let mut zcode_turn = crate::acp::zcode_turn::TurnTracker::default(); // A CodeBuddy native sub-agent's full lifecycle (Agent tool call // open → completed) happens within one turn, so reset the // suppression window at each turn start. This bounds the tracking @@ -13477,6 +13482,24 @@ async fn run_conversation_loop( let runtime = terminal_runtime.clone(); let session_id = sid.clone(); let cwd_opt = Some(cwd); + // Billing snapshots are per prompt, unlike the + // standard context-occupancy usage_update channel. + // Observe only this prompt's drain: load replay and + // idle notifications cannot seed the next turn. + if agent_type == AgentType::ZCode { + if let Dispatch::Notification(notification) = &dispatch { + if crate::acp::zcode_turn::is_turn_state_method( + notification.method(), + ) { + zcode_turn.observe( + notification.method(), + notification.params(), + &sid.0, + ); + continue; + } + } + } // grok reports `/compact` results on ext methods // that bypass the typed pipeline below and emit a // compaction card/error from `.otherwise`. Count @@ -13729,6 +13752,7 @@ async fn run_conversation_loop( reason_str, turn_started_at_ms, current_session_model_id(state).await, + zcode_turn.metrics(), ) .await; // Same wedge guard as the normal turn exit @@ -13844,6 +13868,7 @@ async fn run_conversation_loop( reason_str, turn_started_at_ms, current_session_model_id(state).await, + zcode_turn.metrics(), ) .await; // The turn is over, so any card still parked here @@ -14152,6 +14177,23 @@ async fn run_conversation_loop( .release_all_for_session(sid.0.as_ref()) .await; tracked_terminal_tool_calls.clear(); + // This path deliberately does not await the + // agent's eventual response. Still close + // ZCode's recorded turn before the refetch + // trigger, with only measurements already + // received for THIS prompt. Late telemetry + // must not become the next prompt's bill. + if agent_type == AgentType::ZCode { + record_turn_end( + agent_type, + &sid.0, + "cancelled", + turn_started_at_ms, + current_session_model_id(state).await, + zcode_turn.metrics(), + ) + .await; + } // Also cancel any pending permission requests // (queued ones included), clearing the card // that was on screen, and immediately emit @@ -40495,9 +40537,12 @@ mod tests { impl CustomLoop { async fn start(turns: Vec>) -> Self { - use agent_client_protocol::schema::v1::PromptResponse; - let agent_type = AgentType::custom(IDLE_PROBE_AGENT).expect("a valid custom id"); + Self::start_for(agent_type, turns).await + } + + async fn start_for(agent_type: AgentType, turns: Vec>) -> Self { + use agent_client_protocol::schema::v1::PromptResponse; let turns = Arc::new(std::sync::Mutex::new(VecDeque::from(turns))); let respond = Arc::new(tokio::sync::Notify::new()); let agent_respond = Arc::clone(&respond); @@ -40711,6 +40756,84 @@ mod tests { } } + #[tokio::test] + async fn zcode_terminal_metrics_reach_history_once_before_turn_complete() { + use crate::parsers::AgentParser as _; + let home = tempfile::tempdir().unwrap(); + let home_str = home.path().to_string_lossy().into_owned(); + temp_env::async_with_vars([("CODEG_HOME", Some(home_str.as_str()))], async { + let start = || ("_zcode/turnState", serde_json::json!({"sessionId":"s1","running":true})); + let end = |input: u64, ms: u64| serde_json::json!({ + "sessionId":"s1","running":false,"turnMs":ms, + "usage":{"inputTokens":input,"outputTokens":350,"cachedReadTokens":900,"cachedWriteTokens":120,"totalTokens":input+1370} + }); + let mut lp = CustomLoop::start_for(AgentType::ZCode, vec![ + vec![start(), + custom_update(serde_json::json!({"sessionUpdate":"usage_update","used":9999,"size":131072})), + ("_zcode/turnState",end(180,4200)), + ("$/zcode/turnState",end(9999,9999)), + ("_zcode/turnState",end(9999,9999)), custom_text("first")], + vec![start(), ("_zcode/turnState",end(12,1700)), custom_text("second")], + // A late prior-turn end and idle/replay data have no opening + // marker in this prompt, so they cannot seed its metadata. + vec![("_zcode/turnState",end(9999,9999)),custom_text("third")], + ]).await; + lp.agent_sends(vec![start(),("_zcode/turnState",end(9999,9999)),grok_barrier()]); + lp.until_barrier().await; + let parser = || crate::parsers::acp_native::AcpNativeParser::new_in( + AgentType::ZCode, crate::paths::codeg_acp_transcripts_root(), + ); + lp.run_turn("one", "first").await; + // TurnComplete is the refetch trigger: its footer must already be + // durable and readable, without a test-only flush after the event. + let first = parser().get_conversation("s1").unwrap(); + let usage = first.turns[1].usage.as_ref().unwrap(); + assert_eq!((usage.input_tokens,usage.output_tokens,usage.cache_read_input_tokens,usage.cache_creation_input_tokens),(180,350,900,120)); + assert_eq!(first.turns[1].duration_ms,Some(4200)); + let stats = first.session_stats.unwrap(); + assert_eq!(stats.total_tokens,Some(1550)); + assert_eq!(stats.context_window_used_tokens,Some(9999)); + assert_eq!(stats.total_duration_ms,4200); + lp.run_turn("two", "second").await; + lp.run_turn("three", "third").await; + let cold = parser().get_conversation("s1").unwrap(); + assert_eq!(cold.turns.len(),6); + assert_eq!(cold.turns[3].usage.as_ref().unwrap().input_tokens,12); + assert_eq!(cold.turns[3].duration_ms,Some(1700)); + assert!(cold.turns[5].usage.is_none()); + let raw = crate::acp_transcript::read_transcript("codeg-zcode-acp","s1"); + assert_eq!(raw.entries.iter().filter(|e| e.k==crate::acp_transcript::EntryKind::TurnEnd).count(),3); + assert!(raw.entries.iter().rev().find(|e| e.k==crate::acp_transcript::EntryKind::TurnEnd).unwrap().p.get("usage").is_none()); + lp.shutdown().await; + }).await; + } + + #[tokio::test] + async fn zcode_cancel_records_a_boundary_without_inheriting_the_previous_bill() { + let home = tempfile::tempdir().unwrap(); + let home_str = home.path().to_string_lossy().into_owned(); + temp_env::async_with_vars([("CODEG_HOME", Some(home_str.as_str()))], async { + let mut lp = CustomLoop::start_for(AgentType::ZCode, vec![ + vec![("_zcode/turnState",serde_json::json!({"sessionId":"s1","running":true})), + ("_zcode/turnState",serde_json::json!({"sessionId":"s1","running":false,"usage":{"inputTokens":30,"outputTokens":12},"turnMs":42})),custom_text("first")], + vec![("_zcode/turnState",serde_json::json!({"sessionId":"s1","running":true})),custom_text("working")], + ]).await; + lp.run_turn("one","first").await; + lp.open_turn("two","working").await; + lp.cmd_tx.send(ConnectionCommand::Cancel).await.unwrap(); + lp.until("cancelled",|e|matches!(e,AcpEvent::TurnComplete{stop_reason,..} if stop_reason=="cancelled")).await; + let raw = crate::acp_transcript::read_transcript("codeg-zcode-acp","s1"); + let ends: Vec<_> = raw.entries.iter().filter(|e|e.k==crate::acp_transcript::EntryKind::TurnEnd).collect(); + assert_eq!(ends.len(),2); + assert_eq!(ends[0].p["usage"]["inputTokens"],30); + assert_eq!(ends[1].p["stopReason"],"cancelled"); + assert!(ends[1].p.get("usage").is_none()); + assert!(ends[1].p.get("durationMs").and_then(|v|v.as_u64()).is_some()); + lp.respond.notify_one(); + lp.shutdown().await; + }).await; + } + /// A setter sent with a reply is answered from the agent's own response to /// that request. Here the agent first re-sends its option list showing the /// requested value — the unprompted echo that makes the broadcasts useless diff --git a/src-tauri/src/acp/mod.rs b/src-tauri/src/acp/mod.rs index 5ec96f8bc9..dd62129955 100644 --- a/src-tauri/src/acp/mod.rs +++ b/src-tauri/src/acp/mod.rs @@ -53,6 +53,7 @@ pub mod trae; pub mod types; pub mod work_task_tools; pub mod zcode; +pub(crate) mod zcode_turn; pub use idle_sweep::{idle_sweep_task, idle_timeout_from_env, SWEEP_INTERVAL_SECS}; pub use internal_bus::{EventBusMetrics, EventBusMetricsSnapshot, InternalEventBus}; diff --git a/src-tauri/src/acp/registry.rs b/src-tauri/src/acp/registry.rs index 7ced7f81e8..539acf9cad 100644 --- a/src-tauri/src/acp/registry.rs +++ b/src-tauri/src/acp/registry.rs @@ -4745,7 +4745,8 @@ pub fn get_agent_meta(agent_type: AgentType) -> AcpAgentMeta { // `ZCODE_ACP_MODE`, else `yolo`) where it reported `yolo` // regardless; codeg's gate confirms its own choice either way. // 0.72.0 adds `usage` and `turnMs` to the `_zcode/turnState` - // notification that ends a turn, which codeg does not read. The + // notification that ends a turn; codeg records these as per-turn + // metadata without adding context-occupancy usage updates. The // rest (quota cards, CLIProxy, sandbox pacing when a workspace arms // the sandbox) is opt-in and outside what codeg drives. distribution: AgentDistribution::Npx { diff --git a/src-tauri/src/acp/zcode_turn.rs b/src-tauri/src/acp/zcode_turn.rs new file mode 100644 index 0000000000..aa86af55c5 --- /dev/null +++ b/src-tauri/src/acp/zcode_turn.rs @@ -0,0 +1,242 @@ +//! Per-prompt billing and wall time supplied by the ZCode ACP bridge. +//! +//! `usage_update` describes context occupancy, not a billable turn. These +//! counters are already cache-exclusive ACP usage; never subtract the cache +//! again or add the deprecated alias/standard context updates to them. + +pub(crate) const TURN_STATE_METHOD: &str = "_zcode/turnState"; + +pub(crate) fn is_turn_state_method(method: &str) -> bool { + matches!(method, TURN_STATE_METHOD | "$/zcode/turnState") +} + +#[derive(Debug, Default, PartialEq)] +pub(crate) struct TurnMetrics { + pub usage: Option, + pub duration_ms: Option, +} + +impl TurnMetrics { + /// Replace measurements on the ONE normal turn-end record, never append a + /// second end or merge token snapshots additively. Missing extras keep the + /// ordinary Codeg-measured duration and do not fabricate billing usage. + pub fn apply_to(&self, payload: &mut serde_json::Value) { + if let Some(usage) = &self.usage { + payload["usage"] = usage.clone(); + } + if let Some(duration_ms) = self.duration_ms { + payload["durationMs"] = duration_ms.into(); + } + } +} + +/// Constructed for each prompt, observed only inside its update drain, and +/// discarded when Codeg ends that prompt. Load replay and idle notifications never enter +/// this tracker. The protocol has no turn ID: its ordered running/ended pair +/// and the sequential prompt boundary are the available correlation markers. +#[derive(Debug, Default)] +pub(crate) struct TurnTracker { + running: bool, + ended: bool, + metrics: Option, +} + +impl TurnTracker { + pub fn observe(&mut self, method: &str, params: &serde_json::Value, session_id: &str) { + if method != TURN_STATE_METHOD + || params.get("sessionId").and_then(|v| v.as_str()) != Some(session_id) + || self.ended + { + return; + } + match params.get("running").and_then(|v| v.as_bool()) { + Some(true) => self.running = true, + Some(false) if self.running => { + self.running = false; + self.ended = true; + self.metrics = Some(TurnMetrics { + usage: params.get("usage").and_then(valid_usage), + duration_ms: params.get("turnMs").and_then(|v| v.as_u64()), + }); + } + _ => {} + } + } + + pub fn metrics(&self) -> Option<&TurnMetrics> { + self.metrics.as_ref() + } +} + +/// Retain only unsigned ACP counters, without inventing missing values. A bad +/// counter drops the usage report as a whole; the independent clock can still +/// be used. Guard the four-bucket sum the history parser will compute as well. +fn valid_usage(value: &serde_json::Value) -> Option { + value.as_object()?; + let mut counters = serde_json::Map::new(); + let mut total = 0_u64; + let mut billing_counter = false; + for key in [ + "inputTokens", + "outputTokens", + "cachedReadTokens", + "cachedWriteTokens", + "totalTokens", + "thoughtTokens", + ] { + let Some(value) = value.get(key).filter(|v| !v.is_null()) else { + continue; + }; + let count = value.as_u64()?; + if matches!( + key, + "inputTokens" | "outputTokens" | "cachedReadTokens" | "cachedWriteTokens" + ) { + total = total.checked_add(count)?; + billing_counter = true; + } + counters.insert(key.to_string(), count.into()); + } + billing_counter.then_some(serde_json::Value::Object(counters)) +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + fn start(tracker: &mut TurnTracker) { + tracker.observe( + TURN_STATE_METHOD, + &json!({"sessionId":"s", "running":true}), + "s", + ); + } + + fn end() -> serde_json::Value { + json!({"sessionId":"s", "running":false, "turnMs":4200, + "usage":{"inputTokens":180,"outputTokens":350,"cachedReadTokens":900, + "cachedWriteTokens":120,"thoughtTokens":200,"totalTokens":1550}}) + } + + #[test] + fn an_ordered_terminal_report_keeps_cache_exclusive_usage_and_wall_time() { + let mut tracker = TurnTracker::default(); + start(&mut tracker); + tracker.observe(TURN_STATE_METHOD, &end(), "s"); + assert_eq!( + tracker.metrics().unwrap().usage.as_ref().unwrap(), + &end()["usage"] + ); + assert_eq!(tracker.metrics().unwrap().duration_ms, Some(4200)); + } + + #[test] + fn aliases_other_sessions_and_unopened_ends_do_not_count() { + let mut tracker = TurnTracker::default(); + tracker.observe(TURN_STATE_METHOD, &end(), "s"); + tracker.observe( + TURN_STATE_METHOD, + &json!({"sessionId":"other", "running":true}), + "s", + ); + tracker.observe( + "$/zcode/turnState", + &json!({"sessionId":"s", "running":true}), + "s", + ); + tracker.observe(TURN_STATE_METHOD, &end(), "s"); + assert!(tracker.metrics().is_none()); + start(&mut tracker); + tracker.observe("$/zcode/turnState", &end(), "s"); + assert!(tracker.metrics().is_none()); + tracker.observe(TURN_STATE_METHOD, &end(), "other"); + assert!(tracker.metrics().is_none()); + } + + #[test] + fn a_terminal_report_is_a_snapshot_and_is_consumed_once() { + let mut tracker = TurnTracker::default(); + start(&mut tracker); + tracker.observe(TURN_STATE_METHOD, &end(), "s"); + let mut duplicate = end(); + duplicate["usage"]["inputTokens"] = json!(9999); + tracker.observe("$/zcode/turnState", &duplicate, "s"); + tracker.observe(TURN_STATE_METHOD, &duplicate, "s"); + start(&mut tracker); + tracker.observe(TURN_STATE_METHOD, &duplicate, "s"); + assert_eq!( + tracker.metrics().unwrap().usage.as_ref().unwrap()["inputTokens"], + 180 + ); + let next = TurnTracker::default(); + assert!( + next.metrics().is_none(), + "a cancelled or rejected next prompt inherits nothing" + ); + } + + #[test] + fn absent_usage_does_not_fabricate_zero_counters() { + let mut tracker = TurnTracker::default(); + start(&mut tracker); + tracker.observe( + TURN_STATE_METHOD, + &json!({"sessionId":"s","running":false}), + "s", + ); + assert_eq!(tracker.metrics(), Some(&TurnMetrics::default())); + } + + #[test] + fn malformed_counters_do_not_discard_a_valid_independent_clock() { + for usage in [ + json!({"inputTokens":-1}), + json!({"outputTokens":"3"}), + json!({"inputTokens":1.5}), + json!({"inputTokens":u64::MAX,"outputTokens":1}), + json!({"totalTokens":10}), + json!(null), + ] { + let mut tracker = TurnTracker::default(); + start(&mut tracker); + tracker.observe( + TURN_STATE_METHOD, + &json!({"sessionId":"s","running":false,"usage":usage,"turnMs":0}), + "s", + ); + assert_eq!(tracker.metrics().unwrap().usage, None); + assert_eq!(tracker.metrics().unwrap().duration_ms, Some(0)); + } + } + + #[test] + fn partial_usage_and_invalid_duration_remain_independent() { + let mut tracker = TurnTracker::default(); + start(&mut tracker); + tracker.observe(TURN_STATE_METHOD, &json!({"sessionId":"s","running":false,"usage":{"cachedReadTokens":42,"cachedWriteTokens":null,"other":"ignored"},"turnMs":-5}), "s"); + assert_eq!( + tracker.metrics().unwrap().usage, + Some(json!({"cachedReadTokens":42})) + ); + assert_eq!(tracker.metrics().unwrap().duration_ms, None); + } + + #[test] + fn one_turn_end_retains_fallbacks_and_replaces_counters_instead_of_adding() { + let mut payload = json!({"stopReason":"cancelled","model":"glm-test","durationMs":9999}); + TurnMetrics::default().apply_to(&mut payload); + assert_eq!(payload["durationMs"], 9999); + assert!(payload.get("usage").is_none()); + let metrics = TurnMetrics { + usage: Some(end()["usage"].clone()), + duration_ms: Some(4200), + }; + metrics.apply_to(&mut payload); + metrics.apply_to(&mut payload); + assert_eq!(payload["usage"], end()["usage"]); + assert_eq!(payload["durationMs"], 4200); + assert_eq!(payload["stopReason"], "cancelled"); + assert_eq!(payload["model"], "glm-test"); + } +} diff --git a/src-tauri/src/parsers/acp_native.rs b/src-tauri/src/parsers/acp_native.rs index 6e12d134b0..e74de21cfe 100644 --- a/src-tauri/src/parsers/acp_native.rs +++ b/src-tauri/src/parsers/acp_native.rs @@ -597,17 +597,27 @@ fn turn_end_model_and_usage(payload: &serde_json::Value) -> (Option, Opt /// Read an ACP usage object. Field names follow the `unstable_session_usage` /// shape, with snake_case accepted as well since agents differ. fn parse_usage(value: &serde_json::Value) -> Option { - let get = |camel: &str, snake: &str| -> u64 { - value - .get(camel) - .or_else(|| value.get(snake)) + let get = |keys: &[&str]| -> u64 { + keys.iter() + .find_map(|key| value.get(*key)) .and_then(|v| v.as_u64()) .unwrap_or(0) }; - let input = get("inputTokens", "input_tokens"); - let output = get("outputTokens", "output_tokens"); - let cache_create = get("cacheCreationInputTokens", "cache_creation_input_tokens"); - let cache_read = get("cacheReadInputTokens", "cache_read_input_tokens"); + let input = get(&["inputTokens", "input_tokens"]); + let output = get(&["outputTokens", "output_tokens"]); + // ACP usage is cache-exclusive. Accept its canonical cache names first, + // retaining the older spellings for transcripts already on disk. Aliases + // name the same bucket; adding them would double-count a cache hit. + let cache_create = get(&[ + "cachedWriteTokens", + "cacheCreationInputTokens", + "cache_creation_input_tokens", + ]); + let cache_read = get(&[ + "cachedReadTokens", + "cacheReadInputTokens", + "cache_read_input_tokens", + ]); if input == 0 && output == 0 && cache_create == 0 && cache_read == 0 { return None; } @@ -2176,6 +2186,40 @@ mod tests { assert_eq!(stats.total_duration_ms, 42); } + #[test] + fn acp_cache_counters_have_precedence_and_are_not_subtracted_twice() { + let usage = parse_usage(&serde_json::json!({ + "inputTokens":180,"outputTokens":350,"cachedReadTokens":900,"cachedWriteTokens":120, + "cacheReadInputTokens":9999,"cacheCreationInputTokens":9999,"thoughtTokens":200 + })) + .unwrap(); + assert_eq!(usage.input_tokens, 180); + assert_eq!(usage.output_tokens, 350); + assert_eq!(usage.cache_read_input_tokens, 900); + assert_eq!(usage.cache_creation_input_tokens, 120); + let cache_only = parse_usage(&serde_json::json!({"cachedReadTokens":42})).unwrap(); + assert_eq!(cache_only.cache_read_input_tokens, 42); + assert_eq!(cache_only.input_tokens, 0); + let canonical_zero = parse_usage( + &serde_json::json!({"outputTokens":1,"cachedReadTokens":0,"cacheReadInputTokens":9999}), + ) + .unwrap(); + assert_eq!(canonical_zero.cache_read_input_tokens, 0); + for legacy in [ + serde_json::json!({"cacheReadInputTokens":2,"cacheCreationInputTokens":3}), + serde_json::json!({"cache_read_input_tokens":2,"cache_creation_input_tokens":3}), + ] { + let usage = parse_usage(&legacy).unwrap(); + assert_eq!( + ( + usage.cache_read_input_tokens, + usage.cache_creation_input_tokens + ), + (2, 3) + ); + } + } + /// The write-side whitelist and the read-side match arms are two /// hand-maintained views of one set. This asserts they agree by RUNNING the /// projection over each variant rather than by restating either list, so From 899a77dfdb79041a766aa2bd6f8919aa78c7b765 Mon Sep 17 00:00:00 2001 From: xintaofei Date: Wed, 21 Oct 2026 16:04:19 +0800 Subject: [PATCH 2/2] fix(zcode): total every busy window a prompt holds and keep cancelled replies whole One prompt can hold several of the bridge's busy windows: the rounds of an /auto or /goal loop, a sandbox continuation round, a compaction the prompt waited behind. Each window's terminal report covers one backend turn, so the turn end now records their sum instead of the first window alone. An end with no running report before it still closes nothing, and the bridge's turnMs is used only when a single window ran; otherwise Codeg's own measurement of the prompt stays. A local cancel no longer records a ZCode turn end. The bridge's last frames arrive after it, and the history parser read them as a reply of their own. This matches every other agent whose history is Codeg's transcript, and the per-prompt tracker still keeps a late report off the next prompt. --- docs/agents/zcode.md | 2 +- src-tauri/src/acp/connection.rs | 183 ++++++++++++++++++++++------ src-tauri/src/acp/zcode_turn.rs | 137 +++++++++++++++++---- src-tauri/src/parsers/acp_native.rs | 5 +- 4 files changed, 265 insertions(+), 62 deletions(-) diff --git a/docs/agents/zcode.md b/docs/agents/zcode.md index 8ad2e5c6f1..3b6729df16 100644 --- a/docs/agents/zcode.md +++ b/docs/agents/zcode.md @@ -25,7 +25,7 @@ The bridge writes a disabled sandbox template to `.zcode/acp/sandbox.json` in ea Codeg records its own ACP transcripts for ZCode, including messages, thoughts, tools and usage. These transcripts support history and reconnecting to sessions created in Codeg. Existing custom ZCode entries keep their registration, launch configuration and history; the built-in option records its history separately. This integration does not scan ZCode's native SQLite sessions. -Completed replies retain the bridge's per-turn token usage (including cache reads and writes) and wall time, so their footer and session totals survive reopening. Context-window occupancy remains a separate reading, not additional billable tokens. When the bridge supplies no turn measurements, Codeg keeps its measured duration without inventing token counts. Older transcripts without those measurements are not retroactively filled. +Completed replies retain the bridge's per-turn token usage (including cache reads and writes) and wall time, so their footer and session totals survive reopening. A prompt the bridge works through in several rounds, such as an `/auto` loop or a sandbox continuation, totals every round's usage and keeps Codeg's measured duration. Context-window occupancy remains a separate reading, not additional billable tokens. When the bridge supplies no turn measurements, Codeg keeps its measured duration without inventing token counts. A cancelled reply records no usage, because the bridge reports it only after Codeg has stopped the turn. Older transcripts without those measurements are not retroactively filled. Manage stdio MCP servers through Codeg's MCP settings and assign them to ZCode. The configuration is stored in Codeg's data directory and forwarded when creating or restoring a session; it is not written into official ZCode's provider or credential files. When multi-agent delegation is enabled, Codeg injects its per-launch `codeg-mcp` companion using the same path. diff --git a/src-tauri/src/acp/connection.rs b/src-tauri/src/acp/connection.rs index 6f4bdf958a..d4a3d8d945 100644 --- a/src-tauri/src/acp/connection.rs +++ b/src-tauri/src/acp/connection.rs @@ -13482,10 +13482,12 @@ async fn run_conversation_loop( let runtime = terminal_runtime.clone(); let session_id = sid.clone(); let cwd_opt = Some(cwd); - // Billing snapshots are per prompt, unlike the - // standard context-occupancy usage_update channel. - // Observe only this prompt's drain: load replay and - // idle notifications cannot seed the next turn. + // ZCode's billing reports are per busy window, + // unlike the standard context-occupancy + // usage_update channel, and this prompt's windows + // add up (see `zcode_turn`). Observe only this + // prompt's drain: load replay and idle + // notifications cannot seed the next turn. if agent_type == AgentType::ZCode { if let Dispatch::Notification(notification) = &dispatch { if crate::acp::zcode_turn::is_turn_state_method( @@ -14177,23 +14179,16 @@ async fn run_conversation_loop( .release_all_for_session(sid.0.as_ref()) .await; tracked_terminal_tool_calls.clear(); - // This path deliberately does not await the - // agent's eventual response. Still close - // ZCode's recorded turn before the refetch - // trigger, with only measurements already - // received for THIS prompt. Late telemetry - // must not become the next prompt's bill. - if agent_type == AgentType::ZCode { - record_turn_end( - agent_type, - &sid.0, - "cancelled", - turn_started_at_ms, - current_session_model_id(state).await, - zcode_turn.metrics(), - ) - .await; - } + // No turn end is recorded here, for ZCode + // as for every agent whose history is + // codeg's transcript: the agent's last + // frames are still on their way and land + // on the idle loop, and a turn end ahead of + // them would make the history parser read + // them as a reply of their own. This + // prompt's `zcode_turn` goes with the + // drain, so a late report never bills the + // next prompt either. // Also cancel any pending permission requests // (queued ones included), clearing the card // that was on screen, and immediately emit @@ -40808,30 +40803,142 @@ mod tests { }).await; } + /// A local cancel lets go of the prompt without waiting for the agent, so + /// the cancelled reply's last frames land on the idle loop. No turn end is + /// recorded ahead of them (as for every agent whose history is codeg's + /// transcript): the late chunk stays part of the reply it ends, and the + /// bridge's late report for the cancelled work bills no other prompt. #[tokio::test] - async fn zcode_cancel_records_a_boundary_without_inheriting_the_previous_bill() { + async fn zcode_cancel_keeps_the_late_tail_in_its_reply_and_bills_no_other_prompt() { + use crate::parsers::AgentParser as _; let home = tempfile::tempdir().unwrap(); let home_str = home.path().to_string_lossy().into_owned(); temp_env::async_with_vars([("CODEG_HOME", Some(home_str.as_str()))], async { - let mut lp = CustomLoop::start_for(AgentType::ZCode, vec![ - vec![("_zcode/turnState",serde_json::json!({"sessionId":"s1","running":true})), - ("_zcode/turnState",serde_json::json!({"sessionId":"s1","running":false,"usage":{"inputTokens":30,"outputTokens":12},"turnMs":42})),custom_text("first")], - vec![("_zcode/turnState",serde_json::json!({"sessionId":"s1","running":true})),custom_text("working")], - ]).await; - lp.run_turn("one","first").await; - lp.open_turn("two","working").await; + let turn_state = |body: serde_json::Value| ("_zcode/turnState", body); + let start = || turn_state(serde_json::json!({"sessionId": "s1", "running": true})); + let end = |input: u64| { + turn_state(serde_json::json!({ + "sessionId": "s1", "running": false, "turnMs": 42, + "usage": {"inputTokens": input, "outputTokens": 12} + })) + }; + let mut lp = CustomLoop::start_for( + AgentType::ZCode, + vec![ + vec![start(), end(30), custom_text("first")], + vec![start(), custom_text("working, ")], + vec![start(), end(5), custom_text("third")], + ], + ) + .await; + lp.run_turn("one", "first").await; + lp.open_turn("two", "working, ").await; lp.cmd_tx.send(ConnectionCommand::Cancel).await.unwrap(); - lp.until("cancelled",|e|matches!(e,AcpEvent::TurnComplete{stop_reason,..} if stop_reason=="cancelled")).await; - let raw = crate::acp_transcript::read_transcript("codeg-zcode-acp","s1"); - let ends: Vec<_> = raw.entries.iter().filter(|e|e.k==crate::acp_transcript::EntryKind::TurnEnd).collect(); - assert_eq!(ends.len(),2); - assert_eq!(ends[0].p["usage"]["inputTokens"],30); - assert_eq!(ends[1].p["stopReason"],"cancelled"); - assert!(ends[1].p.get("usage").is_none()); - assert!(ends[1].p.get("durationMs").and_then(|v|v.as_u64()).is_some()); + lp.until("cancelled", |e| { + matches!(e, AcpEvent::TurnComplete { stop_reason, .. } if stop_reason == "cancelled") + }) + .await; + lp.agent_sends(vec![custom_text("late tail"), end(777), grok_barrier()]); + lp.until_barrier().await; + // The cancelled prompt's own response, which nothing waits on. lp.respond.notify_one(); + lp.run_turn("three", "third").await; + + let raw = crate::acp_transcript::read_transcript("codeg-zcode-acp", "s1"); + let ends: Vec<_> = raw + .entries + .iter() + .filter(|e| e.k == crate::acp_transcript::EntryKind::TurnEnd) + .collect(); + assert_eq!(ends.len(), 2, "a local cancel records no turn end"); + assert_eq!(ends[0].p["usage"]["inputTokens"], 30); + assert_eq!(ends[1].p["usage"]["inputTokens"], 5); + let history = crate::parsers::acp_native::AcpNativeParser::new_in( + AgentType::ZCode, + crate::paths::codeg_acp_transcripts_root(), + ) + .get_conversation("s1") + .unwrap(); + assert_eq!( + turn_shape(&history.turns), + [ + "User: text \"one\"", + "Assistant: text \"first\"", + "User: text \"two\"", + "Assistant: text \"working, late tail\"", + "User: text \"three\"", + "Assistant: text \"third\"", + ] + ); + assert!(history.turns[3].usage.is_none()); lp.shutdown().await; - }).await; + }) + .await; + } + + /// One prompt can hold several busy windows: a compaction it waited + /// behind (reported bare), then the rounds of an `/auto` loop. Each round + /// reports its own usage, so the reply carries their total; its duration + /// is codeg's own measurement of the prompt, not any round's clock. + #[tokio::test] + async fn zcode_a_prompt_held_over_several_windows_totals_them_on_codegs_clock() { + use crate::parsers::AgentParser as _; + let home = tempfile::tempdir().unwrap(); + let home_str = home.path().to_string_lossy().into_owned(); + temp_env::async_with_vars([("CODEG_HOME", Some(home_str.as_str()))], async { + let turn_state = |body: serde_json::Value| ("_zcode/turnState", body); + let start = || turn_state(serde_json::json!({"sessionId": "s1", "running": true})); + let end = |input: u64, ms: u64| { + turn_state(serde_json::json!({ + "sessionId": "s1", "running": false, "turnMs": ms, + "usage": {"inputTokens": input, "outputTokens": 10, "cachedReadTokens": 100} + })) + }; + let mut lp = CustomLoop::start_for( + AgentType::ZCode, + vec![vec![ + start(), + turn_state(serde_json::json!({"sessionId": "s1", "running": false})), + start(), + custom_text("round one, "), + end(40, 3_600_000), + start(), + end(2, 1_800_000), + custom_text("round two"), + ]], + ) + .await; + lp.run_turn("/auto ship it", "round two").await; + + let history = crate::parsers::acp_native::AcpNativeParser::new_in( + AgentType::ZCode, + crate::paths::codeg_acp_transcripts_root(), + ) + .get_conversation("s1") + .unwrap(); + assert_eq!( + turn_shape(&history.turns), + [ + "User: text \"/auto ship it\"", + "Assistant: text \"round one, round two\"", + ] + ); + let reply = &history.turns[1]; + let usage = reply.usage.as_ref().unwrap(); + assert_eq!( + ( + usage.input_tokens, + usage.output_tokens, + usage.cache_read_input_tokens, + usage.cache_creation_input_tokens + ), + (42, 20, 200, 0) + ); + let duration = reply.duration_ms.unwrap(); + assert!(duration < 1_800_000, "codeg's clock, got {duration}"); + lp.shutdown().await; + }) + .await; } /// A setter sent with a reply is answered from the agent's own response to diff --git a/src-tauri/src/acp/zcode_turn.rs b/src-tauri/src/acp/zcode_turn.rs index aa86af55c5..b1728e4190 100644 --- a/src-tauri/src/acp/zcode_turn.rs +++ b/src-tauri/src/acp/zcode_turn.rs @@ -3,6 +3,13 @@ //! `usage_update` describes context occupancy, not a billable turn. These //! counters are already cache-exclusive ACP usage; never subtract the cache //! again or add the deprecated alias/standard context updates to them. +//! +//! The bridge reports each busy window as a running/ended pair, and the end +//! carries that window's own usage. One prompt can hold several windows: an +//! `/auto` or `/goal` loop runs every round inside the prompt that started it, +//! a sandbox approval re-runs the prompt as a continuation round, and a +//! compaction the prompt waited behind reports a window of its own. The bridge +//! scopes each report to one backend turn, so the prompt's usage is their sum. pub(crate) const TURN_STATE_METHOD: &str = "_zcode/turnState"; @@ -17,9 +24,9 @@ pub(crate) struct TurnMetrics { } impl TurnMetrics { - /// Replace measurements on the ONE normal turn-end record, never append a - /// second end or merge token snapshots additively. Missing extras keep the - /// ordinary Codeg-measured duration and do not fabricate billing usage. + /// Replace the measurements on the ONE normal turn-end record rather than + /// append a second end. Missing extras keep the ordinary Codeg-measured + /// duration and do not fabricate billing usage. pub fn apply_to(&self, payload: &mut serde_json::Value) { if let Some(usage) = &self.usage { payload["usage"] = usage.clone(); @@ -31,13 +38,18 @@ impl TurnMetrics { } /// Constructed for each prompt, observed only inside its update drain, and -/// discarded when Codeg ends that prompt. Load replay and idle notifications never enter -/// this tracker. The protocol has no turn ID: its ordered running/ended pair -/// and the sequential prompt boundary are the available correlation markers. +/// discarded when Codeg ends that prompt. Load replay and idle notifications +/// never enter this tracker. The protocol has no turn ID: its ordered +/// running/ended pairs and the sequential prompt boundary are the available +/// correlation markers. An end with no running report before it in this +/// prompt (a repeated end, or a late end of an earlier prompt) closes nothing. #[derive(Debug, Default)] pub(crate) struct TurnTracker { running: bool, - ended: bool, + /// Busy windows that closed during this prompt. + windows: u32, + /// Set once the windows' total no longer fits; no later window revives it. + usage_unrepresentable: bool, metrics: Option, } @@ -45,7 +57,6 @@ impl TurnTracker { pub fn observe(&mut self, method: &str, params: &serde_json::Value, session_id: &str) { if method != TURN_STATE_METHOD || params.get("sessionId").and_then(|v| v.as_str()) != Some(session_id) - || self.ended { return; } @@ -53,11 +64,25 @@ impl TurnTracker { Some(true) => self.running = true, Some(false) if self.running => { self.running = false; - self.ended = true; - self.metrics = Some(TurnMetrics { - usage: params.get("usage").and_then(valid_usage), - duration_ms: params.get("turnMs").and_then(|v| v.as_u64()), - }); + self.windows += 1; + let metrics = self.metrics.get_or_insert_with(TurnMetrics::default); + if let Some(report) = params.get("usage").and_then(valid_usage) { + if !self.usage_unrepresentable { + metrics.usage = match metrics.usage.take() { + Some(total) => add_usage(&total, &report), + None => Some(report), + }; + self.usage_unrepresentable = metrics.usage.is_none(); + } + } + // One window's clock is the prompt's own. Several leave gaps + // between them (a compaction hold, judging the next round, a + // restart) that only Codeg's measurement of the prompt spans. + metrics.duration_ms = if self.windows == 1 { + params.get("turnMs").and_then(|v| v.as_u64()) + } else { + None + }; } _ => {} } @@ -100,6 +125,21 @@ fn valid_usage(value: &serde_json::Value) -> Option { billing_counter.then_some(serde_json::Value::Object(counters)) } +/// One more window's counters on top of the prompt's total, revalidated as a +/// whole. `None` when the sum no longer fits: the usage is dropped rather than +/// saturated into a number nobody reported. +fn add_usage(total: &serde_json::Value, report: &serde_json::Value) -> Option { + let mut sum = total.as_object()?.clone(); + for (key, count) in report.as_object()? { + let prior = sum + .get(key) + .and_then(serde_json::Value::as_u64) + .unwrap_or(0); + sum.insert(key.clone(), prior.checked_add(count.as_u64()?)?.into()); + } + valid_usage(&serde_json::Value::Object(sum)) +} + #[cfg(test)] mod tests { use super::*; @@ -155,20 +195,33 @@ mod tests { } #[test] - fn a_terminal_report_is_a_snapshot_and_is_consumed_once() { + fn a_repeated_end_counts_once_and_each_further_window_adds_its_report() { let mut tracker = TurnTracker::default(); start(&mut tracker); tracker.observe(TURN_STATE_METHOD, &end(), "s"); - let mut duplicate = end(); - duplicate["usage"]["inputTokens"] = json!(9999); - tracker.observe("$/zcode/turnState", &duplicate, "s"); - tracker.observe(TURN_STATE_METHOD, &duplicate, "s"); + let mut second = end(); + second["usage"]["inputTokens"] = json!(9999); + second["turnMs"] = json!(1700); + // The deprecated alias repeats the same transition, and an end with no + // running report before it closes no window: neither counts again. + tracker.observe("$/zcode/turnState", &second, "s"); + tracker.observe(TURN_STATE_METHOD, &second, "s"); + assert_eq!( + tracker.metrics().unwrap().usage.as_ref().unwrap(), + &end()["usage"] + ); + assert_eq!(tracker.metrics().unwrap().duration_ms, Some(4200)); + // A second window (the next `/auto` round, a sandbox continuation) is + // another backend turn with its own report. It adds up, and the clock + // goes back to Codeg's measurement of the whole prompt. start(&mut tracker); - tracker.observe(TURN_STATE_METHOD, &duplicate, "s"); + tracker.observe(TURN_STATE_METHOD, &second, "s"); assert_eq!( - tracker.metrics().unwrap().usage.as_ref().unwrap()["inputTokens"], - 180 + tracker.metrics().unwrap().usage.as_ref().unwrap(), + &json!({"inputTokens":10179,"outputTokens":700,"cachedReadTokens":1800, + "cachedWriteTokens":240,"totalTokens":3100,"thoughtTokens":400}) ); + assert_eq!(tracker.metrics().unwrap().duration_ms, None); let next = TurnTracker::default(); assert!( next.metrics().is_none(), @@ -176,6 +229,48 @@ mod tests { ); } + #[test] + fn a_window_without_a_report_still_hands_the_clock_back_to_codeg() { + // A compaction the prompt waited behind settles its busy window bare; + // the prompt's own round follows with the real report. + let mut tracker = TurnTracker::default(); + start(&mut tracker); + tracker.observe( + TURN_STATE_METHOD, + &json!({"sessionId":"s","running":false}), + "s", + ); + start(&mut tracker); + tracker.observe(TURN_STATE_METHOD, &end(), "s"); + assert_eq!( + tracker.metrics().unwrap().usage.as_ref().unwrap(), + &end()["usage"] + ); + assert_eq!(tracker.metrics().unwrap().duration_ms, None); + } + + #[test] + fn a_total_that_no_longer_fits_drops_the_usage_for_good() { + let window = |tracker: &mut TurnTracker, input: u64| { + start(tracker); + tracker.observe( + TURN_STATE_METHOD, + &json!({"sessionId":"s","running":false,"usage":{"inputTokens":input}}), + "s", + ); + }; + let mut tracker = TurnTracker::default(); + window(&mut tracker, u64::MAX - 1); + window(&mut tracker, 2); + assert_eq!(tracker.metrics().unwrap().usage, None); + window(&mut tracker, 5); + assert_eq!( + tracker.metrics().unwrap().usage, + None, + "a later window does not stand in for the whole prompt" + ); + } + #[test] fn absent_usage_does_not_fabricate_zero_counters() { let mut tracker = TurnTracker::default(); diff --git a/src-tauri/src/parsers/acp_native.rs b/src-tauri/src/parsers/acp_native.rs index e74de21cfe..09ed0a5fc8 100644 --- a/src-tauri/src/parsers/acp_native.rs +++ b/src-tauri/src/parsers/acp_native.rs @@ -387,8 +387,9 @@ pub fn is_recorded_update(update: &SessionUpdate) -> bool { /// not the per-turn input/output counts `MessageTurn::usage` holds — so it /// feeds the session footer's context-window trio, exactly like the built-in /// parsers derive it from their own stores, and never a turn's token usage. -/// (A turn's usage stays `None` for custom agents: ACP has no channel for it, -/// and inventing one from occupancy deltas would be a guess rendered as fact.) +/// (A turn's usage comes only from what its turn end records, today ZCode's +/// per-window reports, see `acp::zcode_turn`; otherwise it stays `None`. +/// Inventing one from occupancy deltas would be a guess rendered as fact.) /// /// Matched on the raw payload rather than by deserializing: the answer is /// almost always in the last few entries, and a string compare per entry costs