diff --git a/src-tauri/Cargo.toml b/src-tauri/Cargo.toml index 88e6867c47f..82aad5c7217 100644 --- a/src-tauri/Cargo.toml +++ b/src-tauri/Cargo.toml @@ -179,14 +179,15 @@ ignore = "0.4" # The official ACP runtime (the `sacp` crate was upstreamed into it and renamed). # Stable protocol v1 only: `unstable_protocol_v2` is deliberately left off, so # `initialize` goes out as codeg writes it, at protocol version 1, and the -# runtime refuses an answer that names any other version. The one unstable -# feature is the surface codeg consumes beyond the stable schema: -# `session/fork` (the Session Notices and Session Compaction RFDs are stable -# since schema 1.11). (`unstable_end_turn_token_usage` stays off: codeg's token -# accounting reads the agents' own transcripts, not the prompt response's -# `usage`.) No default features: codeg runs an agent through its own transport -# (`acp/agent_process.rs`), not the crate's `process`/`stdio` ones. -agent-client-protocol = { version = "3.4.0", features = ["unstable_session_fork"] } +# runtime refuses an answer that names any other version. The two unstable +# features are the surface codeg consumes beyond the stable schema: +# `session/fork`, and the prompt response's `usage`, which is read only for +# adapters whose counter contract is known (`acp/reported_usage.rs`); agents +# with native usage stores keep their own parsers. (The Session Notices and +# Session Compaction RFDs are stable since schema 1.11.) No default features: +# codeg runs an agent through its own transport (`acp/agent_process.rs`), not +# the crate's `process`/`stdio` ones. +agent-client-protocol = { version = "3.4.0", features = ["unstable_session_fork", "unstable_end_turn_token_usage"] } tokio = { version = "1", features = ["process", "io-util", "sync", "macros", "rt", "net", "rt-multi-thread"] } uuid = { version = "1", features = ["v4"] } # Version-5 UUIDs (browser profile data-store identifiers) by hand: uuid's @@ -357,4 +358,3 @@ temp-env = { version = "0.3", features = ["async_closure"] } tokio = { version = "1", features = ["test-util", "macros", "rt", "sync", "time"] } insta = { version = "1", features = ["json", "redactions"] } axum-test = { version = "17", features = ["ws"] } - diff --git a/src-tauri/src/acp/connection.rs b/src-tauri/src/acp/connection.rs index d1e40c3ab1b..13298cb9d78 100644 --- a/src-tauri/src/acp/connection.rs +++ b/src-tauri/src/acp/connection.rs @@ -1824,6 +1824,7 @@ async fn record_turn_end( stop_reason: &str, started_at_ms: u64, model: Option, + usage: Option<&agent_client_protocol::schema::v1::Usage>, metrics: Option<&crate::acp::zcode_turn::TurnMetrics>, ) { let Some(dir) = transcript_dir_for(agent_type) else { @@ -1834,6 +1835,9 @@ async fn record_turn_end( "stopReason": stop_reason, "durationMs": now.saturating_sub(started_at_ms), }); + if let Some(usage) = usage.and_then(|usage| serde_json::to_value(usage).ok()) { + payload["usage"] = usage; + } if let Some(metrics) = metrics { metrics.apply_to(&mut payload); } @@ -2029,7 +2033,7 @@ async fn build_agent( cwd: &Path, stderr_tail: &Arc, scratch: Option<&Path>, -) -> Result { +) -> Result<(AcpAgent, crate::acp::reported_usage::PromptUsageMode), AcpError> { // A conversation can outlive the custom-agent definition it was started // with (the user deleted it in settings). `get_agent_meta` cannot report // that — it is infallible — so it hands back a placeholder with an empty @@ -2044,6 +2048,10 @@ async fn build_agent( } let meta = registry::get_agent_meta(agent_type); debug_assert_eq!(meta.agent_type, agent_type); + let prompt_usage_mode = crate::acp::reported_usage::PromptUsageMode::for_distribution( + agent_type, + &meta.distribution, + ); let agent = match meta.distribution { AgentDistribution::Npx { cmd, args, env, .. } => { @@ -2534,11 +2542,12 @@ async fn build_agent( // alignment (process cwd == session cwd). Guard on an existing directory // so a not-yet-created working_dir (e.g. a worktree path) can't make the // spawn fail. - Ok(if cwd.is_dir() { + let agent = if cwd.is_dir() { agent.with_current_dir(cwd) } else { agent - }) + }; + Ok((agent, prompt_usage_mode)) } /// Stack size for the dedicated OS thread that drives each ACP connection's @@ -2664,7 +2673,7 @@ pub async fn spawn_agent_connection( // `None` (isolation off, or the directory could not be created) means the // child inherits the ambient temp dir exactly as it did before. let scratch = crate::acp::scratch_dir::create(); - let agent = match build_agent( + let (agent, prompt_usage_mode) = match build_agent( agent_type, &runtime_env, &launch_cwd, @@ -2683,6 +2692,7 @@ pub async fn spawn_agent_connection( return Err(e); } }; + session_state.write().await.prompt_usage_mode = prompt_usage_mode; let agent = agent .on_spawn({ let child_pid = Arc::clone(&child_pid); @@ -14248,6 +14258,7 @@ async fn run_conversation_loop( reason_str, turn_started_at_ms, current_session_model_id(state).await, + None, zcode_turn.metrics(), ) .await; @@ -14358,12 +14369,18 @@ async fn run_conversation_loop( // streamed chunks this one is bound-awaited: a // conversation reopened right after a turn must not // miss its tail. + let prompt_usage_mode = state.read().await.prompt_usage_mode; + let model = match prompt_usage_mode.reported_model(response.meta.as_ref()) { + Some(model) => Some(model), + None => current_session_model_id(state).await, + }; record_turn_end( agent_type, &sid.0, reason_str, turn_started_at_ms, - current_session_model_id(state).await, + model, + prompt_usage_mode.select(response.usage.as_ref()), zcode_turn.metrics(), ) .await; @@ -41987,8 +42004,37 @@ mod tests { } async fn start_for(agent_type: AgentType, turns: Vec>) -> Self { + Self::start_with_usages(agent_type, turns, Vec::new()).await + } + + async fn start_with_usages( + agent_type: AgentType, + turns: Vec>, + usages: Vec>, + ) -> Self { + let responses = usages + .into_iter() + .map(|usage| { + let mut response = agent_client_protocol::schema::v1::PromptResponse::new( + StopReason::EndTurn, + ); + response.usage = usage; + response + }) + .collect(); + Self::start_with_responses(agent_type, turns, responses).await + } + + /// Each prompt is answered with the next of `responses`, or with a bare + /// `end_turn` once they run out. + async fn start_with_responses( + agent_type: AgentType, + turns: Vec>, + responses: Vec, + ) -> Self { use agent_client_protocol::schema::v1::PromptResponse; let turns = Arc::new(std::sync::Mutex::new(VecDeque::from(turns))); + let responses = Arc::new(std::sync::Mutex::new(VecDeque::from(responses))); let respond = Arc::new(tokio::sync::Notify::new()); let agent_respond = Arc::clone(&respond); let steer_outcomes = Arc::new(std::sync::Mutex::new(VecDeque::new())); @@ -42013,13 +42059,18 @@ mod tests { responder: Responder, cx: ConnectionTo| { let frames = turns.lock().unwrap().pop_front().unwrap_or_default(); + let response = responses + .lock() + .unwrap() + .pop_front() + .unwrap_or_else(|| PromptResponse::new(StopReason::EndTurn)); send_grok_frames(&cx, &frames)?; let respond = Arc::clone(&agent_respond); // Off the handler, so a held response doesn't hold // up the agent's other incoming messages. tokio::spawn(async move { respond.notified().await; - responder.respond(PromptResponse::new(StopReason::EndTurn)) + responder.respond(response) }); Ok(()) }, @@ -42201,6 +42252,148 @@ mod tests { } } + #[tokio::test] + async fn custom_prompt_response_usage_reaches_history_without_counting_context_occupancy() { + 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 agent_type = AgentType::custom("usage-fixture").unwrap(); + let usage = serde_json::from_value(serde_json::json!({ + "inputTokens":11,"outputTokens":7,"cachedReadTokens":101, + "cachedWriteTokens":17,"thoughtTokens":5,"totalTokens":136 + })).unwrap(); + let mut lp = CustomLoop::start_with_usages(agent_type, vec![ + vec![custom_update(serde_json::json!({"sessionUpdate":"usage_update","used":9000,"size":131072})), custom_text("first")], + vec![custom_text("without counters")], + ], vec![Some(usage), None]).await; + let source = AgentDistribution::Npx { + version: "0.97.0", + package: "@agentclientprotocol/claude-agent-acp", + cmd: "claude-agent-acp", + args: &[], env: &[], node_required: None, + }; + lp.state.write().await.prompt_usage_mode = + crate::acp::reported_usage::PromptUsageMode::for_distribution(agent_type, &source); + lp.run_turn("one", "first").await; + let parser = || crate::parsers::acp_native::AcpNativeParser::new_in( + agent_type, crate::paths::codeg_acp_transcripts_root(), + ); + 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), (11, 7, 101, 17)); + assert_eq!(first.session_stats.as_ref().unwrap().total_tokens, Some(136)); + assert_eq!(first.session_stats.as_ref().unwrap().context_window_used_tokens, Some(9000)); + lp.run_turn("two", "without counters").await; + let second = parser().get_conversation("s1").unwrap(); + assert!(second.turns[3].usage.is_none()); + assert_eq!(second.session_stats.unwrap().total_tokens, Some(136)); + lp.shutdown().await; + }).await; + } + + /// claude-agent-acp's model selector values are aliases (`opus`). A turn + /// whose response attributes all of its counters to one model id is named + /// by that id; when the response lists several, the selector value stands. + #[tokio::test] + async fn a_recorded_turn_is_named_by_the_one_model_its_response_reports() { + 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 agent_type = AgentType::custom("model-fixture").unwrap(); + let response = |models: &[&str]| { + let mut response = + agent_client_protocol::schema::v1::PromptResponse::new(StopReason::EndTurn); + response.usage = Some( + serde_json::from_value(serde_json::json!({ + "inputTokens": 40, "outputTokens": 10, "totalTokens": 50 + })) + .unwrap(), + ); + let rows: Vec<_> = models + .iter() + .map(|model| { + serde_json::json!({"model": model, "token_count": {"totalTokens": 50}}) + }) + .collect(); + response.meta = Some( + serde_json::from_value(serde_json::json!({"quota": {"model_usage": rows}})) + .unwrap(), + ); + response + }; + let mut lp = CustomLoop::start_with_responses( + agent_type, + vec![vec![custom_text("first")], vec![custom_text("second")]], + vec![ + response(&["claude-opus-5-5"]), + response(&["claude-opus-5-5", "claude-haiku-5-5"]), + ], + ) + .await; + { + let mut st = lp.state.write().await; + st.prompt_usage_mode = + crate::acp::reported_usage::PromptUsageMode::PerPromptCacheExclusive; + st.config_options = Some(vec![SessionConfigOptionInfo { + id: "model".to_string(), + name: "Model".to_string(), + description: None, + category: Some("model".to_string()), + kind: SessionConfigKindInfo::Select(SessionConfigSelectInfo { + current_value: "opus".to_string(), + options: vec![SessionConfigSelectOptionInfo { + value: "opus".to_string(), + name: "Opus".to_string(), + description: None, + }], + groups: Vec::new(), + }), + recommended_value: None, + }]); + } + lp.run_turn("one", "first").await; + lp.run_turn("two", "second").await; + let detail = crate::parsers::acp_native::AcpNativeParser::new_in( + agent_type, + crate::paths::codeg_acp_transcripts_root(), + ) + .get_conversation("s1") + .unwrap(); + assert_eq!(detail.turns[1].model.as_deref(), Some("claude-opus-5-5")); + assert_eq!(detail.turns[1].usage.as_ref().map(|u| u.input_tokens), Some(40)); + assert_eq!(detail.turns[3].model.as_deref(), Some("opus")); + lp.shutdown().await; + }) + .await; + } + + #[tokio::test] + async fn unsupported_prompt_response_usage_is_not_recorded_as_per_turn_spend() { + 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 { + for agent_type in [AgentType::CommandCode, AgentType::Trae, AgentType::Custom("unknown-usage")] + { + let usage = |tokens| serde_json::from_value(serde_json::json!({ + "inputTokens":tokens,"outputTokens":0,"cachedReadTokens":100,"totalTokens":tokens, + })).unwrap(); + let mut lp = CustomLoop::start_with_usages(agent_type, vec![ + vec![custom_text("first")], vec![custom_text("second")], vec![custom_text("third")], + ], vec![Some(usage(1000)), Some(usage(2000)), Some(usage(3000))]).await; + lp.run_turn("one", "first").await; + lp.run_turn("two", "second").await; + lp.run_turn("three", "third").await; + let raw = crate::acp_transcript::read_transcript(registry::registry_id_for(agent_type), "s1"); + let ends: Vec<_> = raw.entries.iter().filter(|e| e.k == crate::acp_transcript::EntryKind::TurnEnd).collect(); + assert_eq!(ends.len(), 3, "{agent_type}"); + assert!(ends.iter().all(|e| e.p.get("usage").is_none()), "{agent_type}"); + lp.shutdown().await; + } + }).await; + } + #[tokio::test] async fn zcode_terminal_metrics_reach_history_once_before_turn_complete() { use crate::parsers::AgentParser as _; @@ -42212,7 +42405,7 @@ mod tests { "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![ + let mut lp = CustomLoop::start_with_usages(AgentType::ZCode, vec![ vec![start(), custom_update(serde_json::json!({"sessionUpdate":"usage_update","used":9999,"size":131072})), ("_zcode/turnState",end(180,4200)), @@ -42222,7 +42415,9 @@ mod tests { // 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; + ], vec![Some(serde_json::from_value(serde_json::json!({ + "inputTokens":9999,"outputTokens":1,"totalTokens":10000 + })).unwrap()), None, None]).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( diff --git a/src-tauri/src/acp/mod.rs b/src-tauri/src/acp/mod.rs index 74841093fdb..329105b0a82 100644 --- a/src-tauri/src/acp/mod.rs +++ b/src-tauri/src/acp/mod.rs @@ -46,6 +46,7 @@ pub mod provider_models; pub mod question; pub mod registry; pub mod remote_registry; +pub(crate) mod reported_usage; pub mod result_tool; pub mod scratch_dir; pub mod service_error; diff --git a/src-tauri/src/acp/reported_usage.rs b/src-tauri/src/acp/reported_usage.rs new file mode 100644 index 00000000000..898e60ddb76 --- /dev/null +++ b/src-tauri/src/acp/reported_usage.rs @@ -0,0 +1,277 @@ +//! Prompt-response counters are adapter contracts, not a universal ACP ledger. +//! +//! The official Claude adapter reports cache-exclusive, per-prompt main-loop +//! counters. Command Code reports cumulative, cache-inclusive counters, while +//! TRAE has its own native usage enrichment. Unknown adapters stay unsupported +//! until their counter contract is established. ZCode's existing metrics path +//! remains separate. This policy is captured from the launch definition so a +//! registration edit cannot change the interpretation of a running process. + +use crate::acp::registry::AgentDistribution; +use crate::models::agent::AgentType; + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub(crate) enum PromptUsageMode { + #[default] + Unsupported, + PerPromptCacheExclusive, +} + +impl PromptUsageMode { + pub(crate) fn for_distribution(agent: AgentType, distribution: &AgentDistribution) -> Self { + if !matches!(agent, AgentType::Custom(_)) { + return Self::Unsupported; + } + let AgentDistribution::Npx { + package, + cmd, + version, + .. + } = distribution + else { + return Self::Unsupported; + }; + const PACKAGE: &str = "@agentclientprotocol/claude-agent-acp"; + // Accept a version/tag of this package, never an npm alias, URL, fork, + // or a different executable installed alongside it. + let version_or_tag = |s: &str| { + !s.is_empty() + && s.bytes() + .all(|c| c.is_ascii_alphanumeric() || b"._-+^~*".contains(&c)) + }; + let official = *package == PACKAGE + || package + .strip_prefix(PACKAGE) + .and_then(|suffix| suffix.strip_prefix('@')) + .is_some_and(version_or_tag); + if official && *cmd == "claude-agent-acp" && version_or_tag(version) { + Self::PerPromptCacheExclusive + } else { + Self::Unsupported + } + } + + pub(crate) fn select( + self, + usage: Option<&agent_client_protocol::schema::v1::Usage>, + ) -> Option<&agent_client_protocol::schema::v1::Usage> { + match self { + Self::PerPromptCacheExclusive => usage, + Self::Unsupported => None, + } + } + + /// The one model a response's counters were spent on, when the adapter + /// names it. Without it a turn records the model selector's value, which + /// for this adapter is an alias (`opus`, `default`) that neither a price + /// list nor the by-model breakdown can resolve. + /// + /// Since 0.71.0 the response's `_meta.quota.model_usage` splits the + /// prompt's own tally by model id: the prompt's sub-agents are included, + /// work the agent started on its own is not, and a model with nothing to + /// report is left out. A single model there is the one every recorded + /// counter was spent on. Several cannot say which one the main loop used. + pub(crate) fn reported_model( + self, + meta: Option<&agent_client_protocol::schema::v1::Meta>, + ) -> Option { + if self != Self::PerPromptCacheExclusive { + return None; + } + let mut spent = meta? + .get("quota")? + .get("model_usage")? + .as_array()? + .iter() + .filter(|row| { + row.pointer("/token_count/totalTokens") + .and_then(serde_json::Value::as_u64) + .is_some_and(|total| total > 0) + }) + .map(|row| row.get("model").and_then(serde_json::Value::as_str)); + let model = spent.next()??.trim(); + if model.is_empty() || !spent.all(|other| other.map(str::trim) == Some(model)) { + return None; + } + Some(model.to_string()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn npx(package: &'static str, cmd: &'static str, version: &'static str) -> AgentDistribution { + AgentDistribution::Npx { + package, + cmd, + version, + args: &[], + env: &[], + node_required: None, + } + } + + #[test] + fn recognizes_the_producer_not_the_custom_id_or_display_name() { + let source = npx( + "@agentclientprotocol/claude-agent-acp", + "claude-agent-acp", + "0.97.0", + ); + assert_eq!( + PromptUsageMode::for_distribution(AgentType::Custom("review-seat"), &source), + PromptUsageMode::PerPromptCacheExclusive + ); + for source in [ + npx("command-code", "command-code", "1.86.3"), + npx("trae-cli", "trae", "0.209.1"), + npx("@example/claude-agent-acp", "claude-agent-acp", "1.0.0"), + npx( + "@agentclientprotocol/claude-agent-acp-extra", + "claude-agent-acp", + "1.0.0", + ), + npx( + "@agentclientprotocol/claude-agent-acp@npm:other", + "claude-agent-acp", + "latest", + ), + npx( + "@agentclientprotocol/claude-agent-acp", + "other-command", + "0.97.0", + ), + npx( + "@agentclientprotocol/claude-agent-acp", + "claude-agent-acp", + "npm:other", + ), + ] { + assert_eq!( + PromptUsageMode::for_distribution(AgentType::Custom("claude-code-2"), &source), + PromptUsageMode::Unsupported + ); + } + } + + #[test] + fn official_versions_keep_native_store_agents_on_their_existing_paths() { + let source = npx( + "@agentclientprotocol/claude-agent-acp@0.97.0", + "claude-agent-acp", + "latest", + ); + assert_eq!( + PromptUsageMode::for_distribution(AgentType::Custom("work"), &source), + PromptUsageMode::PerPromptCacheExclusive + ); + for agent in crate::models::agent::BUILTIN_AGENT_TYPES { + assert_eq!( + PromptUsageMode::for_distribution(*agent, &source), + PromptUsageMode::Unsupported + ); + } + } + + /// What the manual add form saves: the package spec alone, so the launch + /// command is derived from it and the version label falls back to + /// `custom`. That definition must still select the adapter's contract. + #[test] + fn a_manually_added_official_adapter_is_recognized_from_its_saved_definition() { + use crate::acp::custom_registry::{ + build_meta, CustomAgentDef, CustomAgentSpec, CustomDistributionKind, NpxSpec, + }; + let mode = |package: &str| { + let meta = build_meta(&CustomAgentDef { + registry_id: "work-claude".into(), + name: "Work Claude".into(), + description: String::new(), + version: String::new(), + distribution_kind: CustomDistributionKind::Npx, + spec: CustomAgentSpec { + npx: Some(NpxSpec { + package: package.into(), + ..Default::default() + }), + ..Default::default() + }, + icon_url: None, + skills_shared_store: false, + skills_dir: None, + source: Default::default(), + version_probe: None, + supports_mcp: true, + }) + .unwrap(); + PromptUsageMode::for_distribution(meta.agent_type, &meta.distribution) + }; + for package in [ + "@agentclientprotocol/claude-agent-acp", + "@agentclientprotocol/claude-agent-acp@0.97.0", + "@agentclientprotocol/claude-agent-acp@latest", + ] { + assert_eq!( + mode(package), + PromptUsageMode::PerPromptCacheExclusive, + "{package}" + ); + } + assert_eq!( + mode("@zed-industries/claude-code-acp@0.16.0"), + PromptUsageMode::Unsupported + ); + } + + #[test] + fn a_response_names_its_model_only_when_one_model_did_the_work() { + let meta = |rows: serde_json::Value| -> agent_client_protocol::schema::v1::Meta { + serde_json::from_value(serde_json::json!({"quota": {"model_usage": rows}})).unwrap() + }; + let row = |model: &str, total: u64| serde_json::json!({"model": model, "token_count": {"totalTokens": total}}); + let mode = PromptUsageMode::PerPromptCacheExclusive; + let one = meta(serde_json::json!([row("claude-opus-5-5", 4010)])); + assert_eq!( + mode.reported_model(Some(&one)).as_deref(), + Some("claude-opus-5-5") + ); + // A model with nothing to report does not make the answer ambiguous. + let idle = meta(serde_json::json!([ + row("claude-opus-5-5", 4010), + row("claude-haiku-5-5", 0) + ])); + assert_eq!( + mode.reported_model(Some(&idle)).as_deref(), + Some("claude-opus-5-5") + ); + // A sub-agent on another model: which one the main loop used is unknown. + let two = meta(serde_json::json!([ + row("claude-opus-5-5", 4010), + row("claude-haiku-5-5", 2005) + ])); + assert_eq!(mode.reported_model(Some(&two)), None); + assert_eq!( + mode.reported_model(Some(&meta(serde_json::json!([])))), + None + ); + assert_eq!(mode.reported_model(None), None); + assert_eq!( + PromptUsageMode::Unsupported.reported_model(Some(&one)), + None + ); + } + + #[test] + fn counters_are_not_guessed_for_unknown_or_cumulative_adapters() { + let usage: agent_client_protocol::schema::v1::Usage = serde_json::from_value(serde_json::json!({ + "inputTokens": 234, "outputTokens": 56, "cachedReadTokens": 1000, "totalTokens": 1290, + })).unwrap(); + assert!(PromptUsageMode::Unsupported.select(Some(&usage)).is_none()); + assert!(PromptUsageMode::PerPromptCacheExclusive + .select(Some(&usage)) + .is_some()); + assert!(PromptUsageMode::PerPromptCacheExclusive + .select(None) + .is_none()); + } +} diff --git a/src-tauri/src/acp/session_state.rs b/src-tauri/src/acp/session_state.rs index 0e8e2e2af97..3184adf25ec 100644 --- a/src-tauri/src/acp/session_state.rs +++ b/src-tauri/src/acp/session_state.rs @@ -266,6 +266,8 @@ pub struct SessionState { /// them as copied history and drop them. pub external_id_changed_at: Option, pub agent_type: AgentType, + /// Counter semantics of the adapter definition used for this launch. + pub(crate) prompt_usage_mode: crate::acp::reported_usage::PromptUsageMode, pub working_dir: Option, pub owner_window_label: String, pub folder_id: Option, @@ -859,6 +861,7 @@ impl SessionState { external_id: None, external_id_changed_at: None, agent_type, + prompt_usage_mode: crate::acp::reported_usage::PromptUsageMode::default(), working_dir, owner_window_label, folder_id, diff --git a/src-tauri/src/acp_transcript.rs b/src-tauri/src/acp_transcript.rs index f2fb50530e6..0997d045156 100644 --- a/src-tauri/src/acp_transcript.rs +++ b/src-tauri/src/acp_transcript.rs @@ -694,7 +694,9 @@ fn compact_batch(pending: &[Queued], compactable: bool) -> (String, bool) { // absent on both) is the only case where merging loses nothing. Some(kind) if open.as_ref().is_some_and(|(open_kind, open_entry)| { - *open_kind == kind && open_entry.p.get("_meta") == entry.p.get("_meta") + *open_kind == kind + && open_entry.p.get("_meta") == entry.p.get("_meta") + && open_entry.p.get("messageId") == entry.p.get("messageId") }) => { let (_, open_entry) = open.as_mut().expect("guard matched Some"); @@ -2092,6 +2094,36 @@ mod tests { assert_eq!(entries[2].p["_meta"], serde_json::json!({"a": 2})); } + #[test] + fn typed_message_ids_survive_recording_and_separate_compacted_messages() { + use agent_client_protocol::schema::v1::SessionUpdate; + let message = |t, id| { + // Exercise the same typed decode/encode boundary as the live + // notification handler, not only hand-built transcript JSON. + let update: SessionUpdate = serde_json::from_value(serde_json::json!({ + "sessionUpdate": "agent_message_chunk", + "content": {"type": "text", "text": "x"}, + "messageId": id, + })).unwrap(); + PendingRecord::Entry(TranscriptEntry { + t, k: EntryKind::Update, steer: false, + p: serde_json::to_value(update).unwrap(), + }) + }; + let pending: Vec<_> = vec![ + prompt_record(1, "hi"), + message(2, "msg-first"), message(3, "msg-first"), + message(4, "msg-second"), message(5, "msg-second"), + ].into_iter().map(queued).collect(); + let (lines, _) = compact_batch(&pending, false); + let entries = parse_transcript(&lines).entries; + assert_eq!(entries.len(), 3); + assert_eq!(entries[1].p["messageId"], "msg-first"); + assert_eq!(entries[2].p["messageId"], "msg-second"); + assert_eq!(entries[1].p["content"]["text"], "xx"); + assert_eq!(entries[2].p["content"]["text"], "xx"); + } + /// The frontend refetches conversation detail the moment a turn completes, /// so everything buffered during the turn has to be on disk by the time the /// turn end's ack resolves — that ack is exactly what `record_turn_end` diff --git a/src-tauri/src/commands/conversations.rs b/src-tauri/src/commands/conversations.rs index c45087e540f..27c27640754 100644 --- a/src-tauri/src/commands/conversations.rs +++ b/src-tauri/src/commands/conversations.rs @@ -785,17 +785,26 @@ pub async fn scan_importable_sessions_core( ) .await; - scan_importable_sessions_from_summaries(conn, emitter, chat_channel_manager, summaries).await + scan_importable_sessions_from_summaries( + conn, + emitter, + chat_channel_manager, + summaries, + &crate::paths::codeg_acp_transcripts_root(), + ) + .await } /// Reconcile summaries already collected by the filesystem scan. Keeping this /// boundary separate makes the DB refresh and notification behavior testable -/// without reading the developer's real agent session directories. +/// without reading the developer's real agent session directories — or codeg's +/// own ACP transcripts, read from `transcripts_root`. async fn scan_importable_sessions_from_summaries( conn: &sea_orm::DatabaseConnection, emitter: &EventEmitter, chat_channel_manager: &crate::chat_channel::manager::ChatChannelManager, summaries: Vec<(AgentType, ConversationSummary)>, + transcripts_root: &std::path::Path, ) -> Result { use sea_orm::{ColumnTrait, EntityTrait, QueryFilter}; @@ -816,6 +825,13 @@ async fn scan_importable_sessions_from_summaries( .or_insert(live); *entry = *entry || live; } + // The Claude Code sessions other conversations already count are in codeg + // already, and the import skips them (see + // `import_service::claude_sessions_counted_in`). + let claude = agent_type_db_str(&AgentType::ClaudeCode); + for session in import_service::claude_sessions_counted_in(&conv_rows, transcripts_root) { + imported_index.insert((claude.clone(), session), true); + } // Refresh the already-imported rows in place before answering, then // broadcast each one so open sidebars re-sort without a refetch. The @@ -846,12 +862,14 @@ pub async fn scan_importable_sessions( /// Batch-import the selected sessions, creating (or reopening) each target /// folder as needed. Test seam for [`import_selected_sessions_core`]: takes the -/// scanned summaries as input instead of walking the filesystem. +/// scanned summaries as input instead of walking the filesystem, and reads +/// codeg's own ACP transcripts from `transcripts_root`. pub(crate) async fn import_selected_from_summaries( conn: &sea_orm::DatabaseConnection, emitter: &EventEmitter, summaries: Vec<(AgentType, ConversationSummary)>, selections: Vec, + transcripts_root: &std::path::Path, ) -> Result { const MAX_ERRORS: usize = 10; @@ -894,6 +912,10 @@ pub(crate) async fn import_selected_from_summaries( let folder_rows = load_folder_rows(conn).await?; let folder_index = index_folder_rows(&folder_rows); + let counted_elsewhere = + import_service::claude_sessions_counted_elsewhere(conn, transcripts_root) + .await + .map_err(AppCommandError::from)?; let mut result = ImportSelectedResult { imported: 0, @@ -1070,6 +1092,7 @@ pub(crate) async fn import_selected_from_summaries( folder_id, &items, import_service::DeletedPolicy::Restore, + &counted_elsewhere, ) .await; result.imported += tally.imported; @@ -1153,7 +1176,14 @@ pub async fn import_selected_sessions_core( let envs = agent_setting_service::load_agent_envs(conn).await; let summaries = import_service::collect_local_summaries(&envs, |_, _, _, _| {}).await; - import_selected_from_summaries(conn, emitter, summaries, selections).await + import_selected_from_summaries( + conn, + emitter, + summaries, + selections, + &crate::paths::codeg_acp_transcripts_root(), + ) + .await } #[cfg(feature = "tauri-runtime")] @@ -4958,6 +4988,7 @@ mod tests { &emitter, &chat_channel_manager, vec![summary], + no_transcripts(), ) .await .expect("scan summaries"); @@ -4986,6 +5017,115 @@ mod tests { ); } + /// A custom agent on `package`, saved the way the add form saves one. + fn npx_agent_def(id: &str, package: &str) -> crate::acp::custom_registry::CustomAgentDef { + use crate::acp::custom_registry::{ + CustomAgentDef, CustomAgentSpec, CustomDistributionKind, NpxSpec, + }; + CustomAgentDef { + registry_id: id.into(), + name: id.into(), + description: String::new(), + version: String::new(), + distribution_kind: CustomDistributionKind::Npx, + spec: CustomAgentSpec { + npx: Some(NpxSpec { + package: package.into(), + ..Default::default() + }), + ..Default::default() + }, + icon_url: None, + skills_shared_store: false, + skills_dir: None, + source: Default::default(), + version_probe: None, + supports_mcp: true, + } + } + + /// The import skips a Claude Code session whose usage a custom agent's + /// conversation counts, so the picker must not offer it as new. + #[tokio::test] + // `hydrate` replaces a process-global map; the lock is held throughout. + #[allow(clippy::await_holding_lock)] + async fn scan_shows_a_claude_session_a_custom_agent_counts_as_imported() { + use crate::acp::custom_registry::{hydrate, hydrate_test_guard}; + let _guard = hydrate_test_guard(); + hydrate(&[ + npx_agent_def( + "scan-claude-adapter", + "@agentclientprotocol/claude-agent-acp", + ), + npx_agent_def("scan-other-adapter", "@zed-industries/claude-code-acp"), + ]); + let db = fresh_in_memory_db().await; + let cwd = "/tmp/codeg-scan-counted-elsewhere"; + let folder_id = seed_folder(&db, cwd).await; + for (agent, session) in [ + ("scan-claude-adapter", "counted"), + ("scan-other-adapter", "uncounted"), + ("scan-claude-adapter", "deleted"), + ("scan-claude-adapter", "carried-on"), + ] { + let agent = AgentType::custom(agent).unwrap(); + let row = conversation_service::create(&db.conn, folder_id, agent, None, None) + .await + .expect("create"); + conversation_service::bind_external_id(&db.conn, row.id, session, &[]) + .await + .expect("bind"); + if session == "deleted" { + conversation_service::soft_delete(&db.conn, row.id) + .await + .expect("soft delete"); + } + } + // The adapter could not load `carried-from` again, so its conversation + // went on in `carried-on`, whose transcript continues the first one. + let transcripts = tempfile::tempdir().expect("transcripts root"); + let carried = AgentType::custom("scan-claude-adapter").unwrap(); + let header = + crate::acp_transcript::TranscriptHeader::new(&carried.as_wire(), "carried-on", cwd, 0) + .continuing("carried-from"); + crate::acp_transcript::record_header_in( + transcripts.path(), + crate::acp::registry::registry_id_for(carried), + &header, + ) + .await + .expect("record header"); + let (_broadcaster, emitter) = sync_test_emitter(); + + let result = scan_importable_sessions_from_summaries( + &db.conn, + &emitter, + &crate::chat_channel::manager::ChatChannelManager::new(), + vec![ + scan_summary("counted", AgentType::ClaudeCode, Some(cwd), at(0)), + scan_summary("uncounted", AgentType::ClaudeCode, Some(cwd), at(1)), + scan_summary("deleted", AgentType::ClaudeCode, Some(cwd), at(2)), + scan_summary("carried-from", AgentType::ClaudeCode, Some(cwd), at(3)), + ], + transcripts.path(), + ) + .await + .expect("scan summaries"); + hydrate(&[]); + + let by_id: HashMap<&str, ScanSessionStatus> = result.folders[0] + .sessions + .iter() + .map(|s| (s.external_id.as_str(), s.status)) + .collect(); + assert_eq!(by_id["counted"], ScanSessionStatus::Imported); + assert_eq!(by_id["carried-from"], ScanSessionStatus::Imported); + assert_eq!(by_id["uncounted"], ScanSessionStatus::New); + // A deleted conversation counts nothing, so the session imports. + assert_eq!(by_id["deleted"], ScanSessionStatus::New); + assert_eq!(result.importable_count, 2); + } + /// The channel half of a title notification must not be on the caller's /// critical path: `edit_thread_title` reaches Telegram with a 60s per-call /// timeout, and `list_all_conversations` is the sidebar's primary read. @@ -5671,6 +5811,13 @@ mod tests { // directly against a test broadcaster and assert the emitted JSON. // ────────────────────────────────────────────────────────────────────── + /// The ACP transcripts root for import tests with no transcript to read: + /// one is read only for a conversation of a custom agent on the official + /// Claude adapter. + fn no_transcripts() -> &'static std::path::Path { + std::path::Path::new("/nonexistent/codeg-test-acp-transcripts") + } + fn sync_test_emitter() -> ( std::sync::Arc, EventEmitter, @@ -6266,6 +6413,7 @@ mod tests { key_of(AgentType::ClaudeCode, "s1"), key_of(AgentType::Codex, "s2"), ], + no_transcripts(), ) .await .expect("batch import"); @@ -6289,6 +6437,57 @@ mod tests { assert!(convs.iter().all(|c| c.folder_id == folder_rows[0].id)); } + /// A Claude Code session whose usage a custom agent's conversation already + /// counts stays out even when a request names it. + #[tokio::test] + // `hydrate` replaces a process-global map; the lock is held throughout. + #[allow(clippy::await_holding_lock)] + async fn batch_import_skips_a_claude_session_a_custom_agent_counts() { + use crate::acp::custom_registry::{hydrate, hydrate_test_guard}; + use sea_orm::{ColumnTrait, EntityTrait, QueryFilter}; + let _guard = hydrate_test_guard(); + hydrate(&[npx_agent_def( + "batch-claude-adapter", + "@agentclientprotocol/claude-agent-acp", + )]); + let db = fresh_in_memory_db().await; + let cwd = "/tmp/proj-counted"; + let folder_id = seed_folder(&db, cwd).await; + let agent = AgentType::custom("batch-claude-adapter").unwrap(); + let row = conversation_service::create(&db.conn, folder_id, agent, None, None) + .await + .expect("create"); + conversation_service::bind_external_id(&db.conn, row.id, "counted", &[]) + .await + .expect("bind"); + + let result = import_selected_from_summaries( + &db.conn, + &EventEmitter::Noop, + vec![ + scan_summary("counted", AgentType::ClaudeCode, Some(cwd), at(0)), + scan_summary("fresh", AgentType::ClaudeCode, Some(cwd), at(1)), + ], + vec![ + key_of(AgentType::ClaudeCode, "counted"), + key_of(AgentType::ClaudeCode, "fresh"), + ], + no_transcripts(), + ) + .await + .expect("batch import"); + hydrate(&[]); + + assert_eq!((result.imported, result.skipped), (1, 1)); + let claude_rows = conversation::Entity::find() + .filter(conversation::Column::AgentType.eq("claude_code")) + .all(&db.conn) + .await + .unwrap(); + assert_eq!(claude_rows.len(), 1); + assert_eq!(claude_rows[0].external_id.as_deref(), Some("fresh")); + } + // A session whose cwd is a linked git worktree of an open repo belongs // UNDER that repo. `parent_id` is exactly the sidebar's test for "worktree // of" versus "one more top-level folder", and the branch-label backfill @@ -6313,6 +6512,7 @@ mod tests { &EventEmitter::Noop, summaries, vec![key_of(AgentType::Codex, "s1")], + no_transcripts(), ) .await .expect("batch import"); @@ -6362,6 +6562,7 @@ mod tests { key_of(AgentType::Codex, "s0"), key_of(AgentType::Codex, "s1"), ], + no_transcripts(), ) .await .expect("batch import"); @@ -6420,6 +6621,7 @@ mod tests { &EventEmitter::Noop, summaries, vec![key_of(AgentType::Codex, "s1")], + no_transcripts(), ) .await .expect("batch import"); @@ -6469,6 +6671,7 @@ mod tests { &EventEmitter::Noop, summaries, vec![key_of(AgentType::Codex, "s1")], + no_transcripts(), ) .await .expect("batch import"); @@ -6528,6 +6731,7 @@ mod tests { &EventEmitter::Noop, summaries, vec![key_of(AgentType::Codex, "s1")], + no_transcripts(), ) .await .expect("batch import"); @@ -6585,6 +6789,7 @@ mod tests { &EventEmitter::Noop, summaries, vec![key_of(AgentType::Codex, "s1")], + no_transcripts(), ) .await .expect("batch import"); @@ -6630,6 +6835,7 @@ mod tests { &EventEmitter::Noop, summaries, vec![key_of(AgentType::Codex, "s1")], + no_transcripts(), ) .await .expect("batch import"); @@ -6674,6 +6880,7 @@ mod tests { &EventEmitter::Noop, summaries, vec![key_of(AgentType::Codex, "s1")], + no_transcripts(), ) .await .expect("batch import"); @@ -6712,6 +6919,7 @@ mod tests { &EventEmitter::Noop, summaries, vec![key_of(AgentType::Codex, "s1")], + no_transcripts(), ) .await .expect("batch import"); @@ -6744,6 +6952,7 @@ mod tests { &EventEmitter::Noop, summaries, vec![key_of(AgentType::ClaudeCode, "s1")], + no_transcripts(), ) .await .expect("batch import"); @@ -6791,6 +7000,7 @@ mod tests { &EventEmitter::Noop, summaries, vec![key_of(AgentType::ClaudeCode, "s1")], + no_transcripts(), ) .await .expect("batch import"); @@ -6827,6 +7037,7 @@ mod tests { &EventEmitter::Noop, make(), vec![key_of(AgentType::ClaudeCode, "s1")], + no_transcripts(), ) .await .unwrap(); @@ -6840,6 +7051,7 @@ mod tests { key_of(AgentType::ClaudeCode, "s1"), key_of(AgentType::Codex, "does-not-exist"), ], + no_transcripts(), ) .await .unwrap(); @@ -6867,6 +7079,7 @@ mod tests { &EventEmitter::Noop, make(), vec![key_of(AgentType::ClaudeCode, "s1")], + no_transcripts(), ) .await .unwrap(); @@ -6891,6 +7104,7 @@ mod tests { &EventEmitter::Noop, make(), vec![key_of(AgentType::ClaudeCode, "s1")], + no_transcripts(), ) .await .unwrap(); @@ -6923,6 +7137,7 @@ mod tests { folder_id, &items, import_service::DeletedPolicy::Skip, + &HashSet::new(), ) .await .unwrap(); @@ -6944,6 +7159,7 @@ mod tests { folder_id, &items, import_service::DeletedPolicy::Skip, + &HashSet::new(), ) .await .unwrap(); @@ -6967,6 +7183,7 @@ mod tests { &EventEmitter::Noop, vec![(agent, child)], vec![key_of(AgentType::Hermes, "child")], + no_transcripts(), ) .await .unwrap(); @@ -7002,6 +7219,7 @@ mod tests { key_of(AgentType::ClaudeCode, "s1"), key_of(AgentType::Codex, "s2"), ], + no_transcripts(), ) .await .unwrap(); @@ -7094,6 +7312,7 @@ mod tests { 999_999, &items, import_service::DeletedPolicy::Skip, + &HashSet::new(), ) .await; assert_eq!(failed, 2, "both rows fail the folder FK and are counted"); @@ -7110,6 +7329,7 @@ mod tests { folder_id, &items, import_service::DeletedPolicy::Skip, + &HashSet::new(), ) .await; assert_eq!(failed2, 0); @@ -7134,7 +7354,8 @@ mod tests { &db.conn, 999_999, &items, - import_service::DeletedPolicy::Skip + import_service::DeletedPolicy::Skip, + &HashSet::new() ) .await .is_err(), diff --git a/src-tauri/src/commands/token_usage.rs b/src-tauri/src/commands/token_usage.rs index d9c304146b6..59af52e1a22 100644 --- a/src-tauri/src/commands/token_usage.rs +++ b/src-tauri/src/commands/token_usage.rs @@ -144,6 +144,13 @@ const TOP_CONVERSATIONS: usize = 8; /// codeg's ACP-projected history; each sub-agent's own rollout is read too, /// and its spend lands on the turn that launched it. See /// `parsers::trae_usage`. +/// * `9` — the usage a turn end reports is kept when no open reply is there +/// to carry it: a prompt that produced no text or tools, or a reply that a +/// relayed user message closed before the turn ended. Such usage used to be +/// dropped. It comes from ZCode's turn metrics and, for custom agents, from +/// the official Claude adapter's prompt response. See +/// `parsers::acp_native::project_turns`. (`8` is skipped: preview builds of +/// an earlier proposal wrote it.) /// /// Only the accounting stored in `token_usage_turn` counts: the four token /// counters, the duration and the timestamp. A conversation's context WINDOW is @@ -158,7 +165,7 @@ const TOP_CONVERSATIONS: usize = 8; /// counters predate it, since a custom/BYO model has always exposed them and /// the parser reads every session under `~/.qoder/projects`, not just the ones /// codeg launched. -const FACT_SCHEMA_VERSION: &str = "7"; +const FACT_SCHEMA_VERSION: &str = "9"; const FACT_SCHEMA_VERSION_KEY: &str = "token_usage_fact_schema_version"; @@ -911,10 +918,17 @@ pub async fn token_usage_status_core( let counts = usage_service::sync_counts(conn) .await .map_err(AppCommandError::from)?; + let schema = app_metadata_service::get_value(conn, FACT_SCHEMA_VERSION_KEY) + .await + .map_err(AppCommandError::from)?; Ok(TokenUsageSyncStatus { total_conversations: counts.total_conversations, synced_conversations: counts.synced_conversations, - stale_conversations: counts.stale_conversations, + stale_conversations: if schema.as_deref() != Some(FACT_SCHEMA_VERSION) { + counts.total_conversations + } else { + counts.stale_conversations + }, fact_rows: counts.fact_rows, last_synced_at: counts.last_synced_at, running: SYNC_RUNNING.load(std::sync::atomic::Ordering::Relaxed), diff --git a/src-tauri/src/db/service/import_service.rs b/src-tauri/src/db/service/import_service.rs index a010173f765..2fa17eac67d 100644 --- a/src-tauri/src/db/service/import_service.rs +++ b/src-tauri/src/db/service/import_service.rs @@ -1,3 +1,6 @@ +use std::collections::HashSet; +use std::path::Path; + use sea_orm::{ ActiveModelTrait, ActiveValue::NotSet, ColumnTrait, DatabaseConnection, EntityTrait, QueryFilter, Set, @@ -153,12 +156,14 @@ pub(crate) enum DeletedPolicy { /// re-querying. Strict: the first row error aborts and propagates — this is the /// legacy per-folder import's contract, so its public command still surfaces DB /// failures rather than silently returning `0/0/0`. The batch importer uses the -/// resilient [`import_summaries_resilient`] instead. +/// resilient [`import_summaries_resilient`] instead. `counted_elsewhere` is +/// [`claude_sessions_counted_elsewhere`] for this import. pub(crate) async fn import_summaries( conn: &DatabaseConnection, folder_id: i32, items: &[(AgentType, ConversationSummary)], deleted: DeletedPolicy, + counted_elsewhere: &HashSet, ) -> Result<(ImportResult, Vec), DbError> { let mut imported = 0u32; let mut updated = 0u32; @@ -167,7 +172,16 @@ pub(crate) async fn import_summaries( let mut updated_ids: Vec = Vec::new(); for (agent_type, summary) in items { - match import_one(conn, folder_id, agent_type, summary, deleted).await? { + match import_one( + conn, + folder_id, + agent_type, + summary, + deleted, + counted_elsewhere, + ) + .await? + { ImportOutcome::Imported => imported += 1, ImportOutcome::Updated(id) => { updated += 1; @@ -193,6 +207,7 @@ pub(crate) async fn import_summaries_resilient( folder_id: i32, items: &[(AgentType, ConversationSummary)], deleted: DeletedPolicy, + counted_elsewhere: &HashSet, ) -> (ImportResult, Vec, u32) { let mut imported = 0u32; let mut updated = 0u32; @@ -202,7 +217,16 @@ pub(crate) async fn import_summaries_resilient( let mut updated_ids: Vec = Vec::new(); for (agent_type, summary) in items { - match import_one(conn, folder_id, agent_type, summary, deleted).await { + match import_one( + conn, + folder_id, + agent_type, + summary, + deleted, + counted_elsewhere, + ) + .await + { Ok(ImportOutcome::Imported) => imported += 1, Ok(ImportOutcome::Updated(id)) => { updated += 1; @@ -253,7 +277,17 @@ pub async fn import_local_conversations( }) .collect(); - import_summaries(conn, folder_id, &matched, DeletedPolicy::Skip).await + let counted_elsewhere = + claude_sessions_counted_elsewhere(conn, &crate::paths::codeg_acp_transcripts_root()) + .await?; + import_summaries( + conn, + folder_id, + &matched, + DeletedPolicy::Skip, + &counted_elsewhere, + ) + .await } /// Outcome of reconciling a single parsed session against the DB. @@ -272,7 +306,8 @@ enum ImportOutcome { Restored, /// Already imported and nothing changed — or the row is one the sidebar /// never shows (a delegation child, or a soft-deleted row under - /// [`DeletedPolicy::Skip`]). + /// [`DeletedPolicy::Skip`]), or another agent's conversation already counts + /// the session's usage (see [`claude_sessions_counted_in`]). Skipped, } @@ -284,6 +319,77 @@ fn agent_type_db_str(agent_type: &AgentType) -> String { .unwrap_or_default() } +/// The custom agent behind `agent_type` (the column's string form) when it is +/// the official Claude adapter, whose per-prompt counters codeg records (see +/// [`crate::acp::reported_usage`]) — told from the agent's current +/// registration, the definition its next launch records under. +fn claude_adapter_agent(agent_type: &str) -> Option { + use crate::acp::reported_usage::PromptUsageMode; + let agent = AgentType::from_wire(agent_type)?; + let meta = crate::acp::custom_registry::get(agent.custom_id()?)?; + (PromptUsageMode::for_distribution(agent, &meta.distribution) + == PromptUsageMode::PerPromptCacheExclusive) + .then_some(agent) +} + +/// The Claude Code sessions whose token usage a live conversation in `rows` +/// counts, by id. +/// +/// A custom agent on the official Claude adapter runs Claude Code itself. When +/// it shares Claude Code's config directory, each session it runs is also a +/// Claude Code session under the same id, and importing that session as a +/// Claude Code conversation too would count the same work twice. Its +/// conversation counts the session it holds and every session its transcript +/// continues (see [`crate::acp_transcript::TranscriptHeader::continues_from`]): +/// when a session cannot be loaded again — its config directory changed, say +/// — codeg carries the conversation on in a new one, and the earlier session +/// stays on disk under its own id. +/// +/// Only the agent's current registration is consulted, so conversations +/// recorded under a definition the agent no longer has (removed, or moved to +/// another adapter) are not recognized. Nor is the work after a plan approval +/// that clears the context: the adapter carries on in a new Claude Code +/// session, under an id no conversation or transcript holds. +pub(crate) fn claude_sessions_counted_in( + rows: &[conversation::Model], + transcripts_root: &Path, +) -> HashSet { + let mut sessions = HashSet::new(); + for row in rows.iter().filter(|row| row.deleted_at.is_none()) { + let Some(external_id) = row.external_id.as_deref() else { + continue; + }; + let Some(agent) = claude_adapter_agent(&row.agent_type) else { + continue; + }; + sessions.extend(crate::acp_transcript::continuation_ancestors_in( + transcripts_root, + crate::acp::registry::registry_id_for(agent), + external_id, + )); + sessions.insert(external_id.to_string()); + } + sessions +} + +/// [`claude_sessions_counted_in`] over every live conversation of a custom +/// agent. +pub(crate) async fn claude_sessions_counted_elsewhere( + conn: &DatabaseConnection, + transcripts_root: &Path, +) -> Result, DbError> { + let rows = conversation::Entity::find() + .filter( + conversation::Column::AgentType + .starts_with(crate::models::agent::CUSTOM_AGENT_WIRE_PREFIX), + ) + .filter(conversation::Column::ExternalId.is_not_null()) + .filter(conversation::Column::DeletedAt.is_null()) + .all(conn) + .await?; + Ok(claude_sessions_counted_in(&rows, transcripts_root)) +} + /// Reconcile ONE already-imported conversation against a fresh parse of its /// agent-side session file. Returns `true` when a row was written, so the /// caller can count it and broadcast a sidebar upsert. Never inserts, never @@ -410,12 +516,15 @@ pub(crate) async fn sync_imported_sessions( /// place from the freshly parsed session file (see [`refresh_existing`]). A /// soft-deleted row is restored first when `deleted` says so (see /// [`DeletedPolicy`]), and only then refreshed like any other existing row. +/// `counted_elsewhere` holds the Claude Code sessions other conversations +/// already count ([`claude_sessions_counted_elsewhere`]). async fn import_one( conn: &DatabaseConnection, folder_id: i32, agent_type: &AgentType, summary: &ConversationSummary, deleted: DeletedPolicy, + counted_elsewhere: &HashSet, ) -> Result { let at_str = agent_type_db_str(agent_type); @@ -425,6 +534,17 @@ async fn import_one( .one(conn) .await?; + // A session another agent's live conversation already counts is not added + // under this agent — neither as a new row nor by bringing back a deleted + // one — or the same work would be counted twice. A live row of this agent + // that holds the session already is refreshed below as before. + if *agent_type == AgentType::ClaudeCode + && exists.as_ref().is_none_or(|row| row.deleted_at.is_some()) + && counted_elsewhere.contains(&summary.id) + { + return Ok(ImportOutcome::Skipped); + } + if let Some(mut existing) = exists { // A delegation child is never a sidebar row, so it is left completely // alone whatever the policy — including no token-usage invalidation. @@ -547,7 +667,15 @@ mod tests { agent_type: &AgentType, summary: &ConversationSummary, ) -> Result { - import_one(conn, folder_id, agent_type, summary, DeletedPolicy::Skip).await + import_one( + conn, + folder_id, + agent_type, + summary, + DeletedPolicy::Skip, + &HashSet::new(), + ) + .await } fn summary(id: &str, title: Option<&str>) -> ConversationSummary { @@ -859,6 +987,7 @@ mod tests { &at, &timed_summary("ext-1", Some("AI Summary"), later, 9), DeletedPolicy::Restore, + &HashSet::new(), ) .await .expect("restore"); @@ -901,17 +1030,31 @@ mod tests { .expect("soft delete"); assert_eq!( - import_one(&db.conn, folder, &at, &summary("ext-1", Some("kept")), DeletedPolicy::Restore) - .await - .expect("restore"), + import_one( + &db.conn, + folder, + &at, + &summary("ext-1", Some("kept")), + DeletedPolicy::Restore, + &HashSet::new() + ) + .await + .expect("restore"), ImportOutcome::Restored ); // A second pass has nothing left to restore: the row is live, so it is // an ordinary already-imported row and must not be counted again. assert_eq!( - import_one(&db.conn, folder, &at, &summary("ext-1", Some("kept")), DeletedPolicy::Restore) - .await - .expect("re-run"), + import_one( + &db.conn, + folder, + &at, + &summary("ext-1", Some("kept")), + DeletedPolicy::Restore, + &HashSet::new() + ) + .await + .expect("re-run"), ImportOutcome::Skipped ); } @@ -931,9 +1074,16 @@ mod tests { .await .expect("soft delete"); - import_one(&db.conn, target, &at, &summary("ext-1", Some("t")), DeletedPolicy::Restore) - .await - .expect("restore"); + import_one( + &db.conn, + target, + &at, + &summary("ext-1", Some("t")), + DeletedPolicy::Restore, + &HashSet::new(), + ) + .await + .expect("restore"); // The picker's group folder is the one `add_folder` just made live AND // open; the row's own folder may have been removed or closed since, @@ -985,9 +1135,16 @@ mod tests { // A delegation child is not a sidebar row, so "restore" has no meaning // for it — it would surface a sub-session as a root conversation. assert_eq!( - import_one(&db.conn, folder, &at, &summary("child-ext", Some("x")), DeletedPolicy::Restore) - .await - .expect("restore child"), + import_one( + &db.conn, + folder, + &at, + &summary("child-ext", Some("x")), + DeletedPolicy::Restore, + &HashSet::new() + ) + .await + .expect("restore child"), ImportOutcome::Skipped ); let child = find_row(&db.conn, "child-ext").await; @@ -1293,4 +1450,195 @@ mod tests { "child title untouched" ); } + + /// A custom agent on `package`, saved the way the add form saves one. + fn custom_def(id: &str, package: &str) -> crate::acp::custom_registry::CustomAgentDef { + use crate::acp::custom_registry::{ + CustomAgentDef, CustomAgentSpec, CustomDistributionKind, NpxSpec, + }; + CustomAgentDef { + registry_id: id.into(), + name: id.into(), + description: String::new(), + version: String::new(), + distribution_kind: CustomDistributionKind::Npx, + spec: CustomAgentSpec { + npx: Some(NpxSpec { + package: package.into(), + ..Default::default() + }), + ..Default::default() + }, + icon_url: None, + skills_shared_store: false, + skills_dir: None, + source: Default::default(), + version_probe: None, + supports_mcp: true, + } + } + + /// A conversation of `agent` bound to the session `external_id`. + async fn bound_row( + conn: &DatabaseConnection, + folder_id: i32, + agent: AgentType, + external_id: &str, + ) -> i32 { + let row = conversation_service::create(conn, folder_id, agent, None, None) + .await + .expect("create"); + conversation_service::bind_external_id(conn, row.id, external_id, &[]) + .await + .expect("bind"); + row.id + } + + async fn claude_rows(conn: &DatabaseConnection, ext: &str) -> Vec { + conversation::Entity::find() + .filter(conversation::Column::ExternalId.eq(ext)) + .filter(conversation::Column::AgentType.eq("claude_code")) + .all(conn) + .await + .expect("query") + } + + /// A custom agent on the official Claude adapter runs Claude Code itself, + /// so with a shared config directory each of its sessions is also listed + /// as a Claude Code session under the same id — one whose usage the custom + /// agent's conversation already counts. + #[tokio::test] + // `hydrate` replaces a process-global map; the lock is held throughout. + #[allow(clippy::await_holding_lock)] + async fn a_claude_session_a_custom_agent_counts_is_not_imported_again() { + use crate::acp::custom_registry::{hydrate, hydrate_test_guard}; + let _guard = hydrate_test_guard(); + hydrate(&[ + custom_def( + "import-claude-adapter", + "@agentclientprotocol/claude-agent-acp@0.97.0", + ), + custom_def( + "import-other-adapter", + "@zed-industries/claude-code-acp@0.16.0", + ), + ]); + let official = AgentType::custom("import-claude-adapter").unwrap(); + let other = AgentType::custom("import-other-adapter").unwrap(); + let transcripts = tempfile::tempdir().expect("transcripts root"); + let db = fresh_in_memory_db().await; + let folder = seed_folder(&db, "/tmp/codeg-import").await; + let at = AgentType::ClaudeCode; + let import = |ext: &'static str, deleted: DeletedPolicy| { + let (conn, root) = (db.conn.clone(), transcripts.path().to_path_buf()); + async move { + let counted = claude_sessions_counted_elsewhere(&conn, &root) + .await + .expect("counted sessions"); + import_one( + &conn, + folder, + &at, + &summary(ext, Some("native")), + deleted, + &counted, + ) + .await + .expect("import") + } + }; + + bound_row(&db.conn, folder, official, "counted").await; + assert_eq!( + import("counted", DeletedPolicy::Restore).await, + ImportOutcome::Skipped + ); + assert!(claude_rows(&db.conn, "counted").await.is_empty()); + // Only a Claude Code session is that agent's. + let counted = claude_sessions_counted_elsewhere(&db.conn, transcripts.path()) + .await + .expect("counted sessions"); + assert_eq!( + import_one( + &db.conn, + folder, + &AgentType::Codex, + &summary("counted", None), + DeletedPolicy::Skip, + &counted, + ) + .await + .expect("import"), + ImportOutcome::Imported + ); + + // The adapter could not load `carried-from` again, so the conversation + // went on in `carried-on`; its transcript still counts the first one. + bound_row(&db.conn, folder, official, "carried-on").await; + let header = crate::acp_transcript::TranscriptHeader::new( + &official.as_wire(), + "carried-on", + "/tmp/codeg-import", + 0, + ) + .continuing("carried-from"); + crate::acp_transcript::record_header_in( + transcripts.path(), + crate::acp::registry::registry_id_for(official), + &header, + ) + .await + .expect("record header"); + assert_eq!( + import("carried-from", DeletedPolicy::Skip).await, + ImportOutcome::Skipped + ); + + // codeg records no usage for this adapter: the native session is the + // only place its usage is counted, so it imports as before. + bound_row(&db.conn, folder, other, "uncounted").await; + assert_eq!( + import("uncounted", DeletedPolicy::Skip).await, + ImportOutcome::Imported + ); + + // A deleted conversation counts nothing. + let deleted = bound_row(&db.conn, folder, official, "deleted-holder").await; + conversation_service::soft_delete(&db.conn, deleted) + .await + .expect("soft delete"); + assert_eq!( + import("deleted-holder", DeletedPolicy::Skip).await, + ImportOutcome::Imported + ); + + // Nor is a deleted Claude Code copy brought back beside a live holder… + bound_row(&db.conn, folder, official, "deleted-copy").await; + let copy = bound_row(&db.conn, folder, at, "deleted-copy").await; + conversation_service::soft_delete(&db.conn, copy) + .await + .expect("soft delete"); + assert_eq!( + import("deleted-copy", DeletedPolicy::Restore).await, + ImportOutcome::Skipped + ); + assert!(claude_rows(&db.conn, "deleted-copy").await[0] + .deleted_at + .is_some()); + + // …while a live one imported before is refreshed as it always was. + bound_row(&db.conn, folder, official, "live-copy").await; + let copy = bound_row(&db.conn, folder, at, "live-copy").await; + assert_eq!( + import("live-copy", DeletedPolicy::Skip).await, + ImportOutcome::Updated(copy) + ); + + // Told from the registration: with the agent removed, it imports. + hydrate(&[]); + assert_eq!( + import("counted", DeletedPolicy::Skip).await, + ImportOutcome::Imported + ); + } } diff --git a/src-tauri/src/models/conversation.rs b/src-tauri/src/models/conversation.rs index da6d62599f3..b1527d4326c 100644 --- a/src-tauri/src/models/conversation.rs +++ b/src-tauri/src/models/conversation.rs @@ -283,7 +283,9 @@ pub struct ImportResult { pub enum ScanSessionStatus { /// Not in the DB — importable. New, - /// At least one live (non-deleted) row exists — already imported. + /// At least one live (non-deleted) row exists — already imported. Also a + /// Claude Code session whose usage a live conversation of a custom agent + /// on the official Claude adapter counts: the import skips it. Imported, /// Only soft-deleted rows exist. Deletion is a soft delete, so the row (and /// its whole history) is still there: the picker offers such a session for diff --git a/src-tauri/src/parsers/acp_native.rs b/src-tauri/src/parsers/acp_native.rs index 3ea91f9b56b..7703cdef088 100644 --- a/src-tauri/src/parsers/acp_native.rs +++ b/src-tauri/src/parsers/acp_native.rs @@ -433,8 +433,8 @@ 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 comes only from what its turn end records, today ZCode's -/// per-window reports, see `acp::zcode_turn`; otherwise it stays `None`. +/// (A turn's usage comes only from its recorded prompt response or ZCode's +/// native turn metrics, 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 @@ -532,10 +532,13 @@ pub fn project_turns(entries: &[TranscriptEntry]) -> Vec { // prompt; the part's own span (`flush`) is the honest one. apply_turn_end(p, &entry.p, !steered); p.last_at_ms = entry.t; - } else if steered { - // Nothing came after the last steer, so the usage and model - // the turn end reports go on the part before it — never on - // a reply to an earlier prompt. + } else { + // Nothing came after the last user message of this turn — + // a steer, or one the agent relayed from another client + // (claude-agent-acp forwards Remote Control messages as + // `user_message_chunk`). The usage and model the turn end + // reports go on the part before it — never on a reply to an + // earlier prompt. if let Some(last) = turns .iter_mut() .skip(prompt_start.unwrap_or(usize::MAX)) @@ -551,6 +554,20 @@ pub fn project_turns(entries: &[TranscriptEntry]) -> Vec { } } } + // A paid attempt can finish without emitting any text or + // tools. Keep its reported usage without inventing content, + // and never attach an unprompted footer to an earlier turn. + if pending.is_none() + && prompt_start.is_some() + && !turns.iter().skip(prompt_start.unwrap_or(usize::MAX)) + .any(|turn| matches!(turn.role, TurnRole::Assistant)) + && turn_end_model_and_usage(&entry.p).1.is_some() + { + let mut turn = PendingTurn::new(turn_start_hint.unwrap_or(entry.t)); + apply_turn_end(&mut turn, &entry.p, !steered); + turn.last_at_ms = entry.t; + pending = Some(turn); + } flush(&mut pending, &mut turns, &mut seq); prompt_just_recorded = false; turn_start_hint = None; @@ -593,7 +610,7 @@ pub fn project_turns(entries: &[TranscriptEntry]) -> Vec { fn flush(pending: &mut Option, turns: &mut Vec, seq: &mut usize) { let Some(p) = pending.take() else { return }; - if !p.has_content { + if !p.has_content && p.usage.is_none() { return; } turns.push(MessageTurn { @@ -641,7 +658,7 @@ fn turn_end_model_and_usage(payload: &serde_json::Value) -> (Option, Opt (model, usage) } -/// Read an ACP usage object. Field names follow the `unstable_session_usage` +/// Read an ACP usage object. Field names follow `unstable_end_turn_token_usage` /// shape, with snake_case accepted as well since agents differ. fn parse_usage(value: &serde_json::Value) -> Option { let get = |keys: &[&str]| -> u64 { @@ -1168,6 +1185,46 @@ mod tests { }) } + #[test] + fn reported_usage_survives_a_turn_without_output_but_not_an_orphan_footer() { + let footer = entry(200, EntryKind::TurnEnd, serde_json::json!({ + "stopReason":"end_turn", "durationMs":100, + "usage":{"inputTokens":3,"outputTokens":0,"cachedReadTokens":4,"totalTokens":7} + })); + let turns = project_turns(&[prompt(100, "attempt"), footer.clone()]); + assert_eq!(turns.len(), 2); + assert!(turns[1].blocks.is_empty()); + assert_eq!(turns[1].usage.as_ref().unwrap().input_tokens, 3); + assert_eq!(turns[1].usage.as_ref().unwrap().cache_read_input_tokens, 4); + assert_eq!(turns[1].duration_ms, Some(100)); + assert!(project_turns(&[footer]).is_empty()); + } + + /// A message typed in another client of the session reaches codeg as a + /// `user_message_chunk` while the prompt is still open. It closes the + /// reply so far, and when nothing follows it the turn end's usage still + /// belongs to that reply. + #[test] + fn a_relayed_user_message_does_not_drop_the_turn_end_usage() { + let turns = project_turns(&[ + prompt(100, "P"), + update(150, text_chunk("agent_message_chunk", "Reply")), + update(200, text_chunk("user_message_chunk", "Remote follow-up")), + entry( + 300, + EntryKind::TurnEnd, + serde_json::json!({ + "stopReason": "end_turn", "model": "opus", + "usage": {"inputTokens": 4000, "outputTokens": 10, "totalTokens": 4010} + }), + ), + ]); + assert_eq!(roles(&turns), ["User", "Assistant", "User"]); + assert_eq!(turns[1].usage.as_ref().map(|u| u.input_tokens), Some(4000)); + assert_eq!(turns[1].model.as_deref(), Some("opus")); + assert_eq!(turns[2].usage, None); + } + fn attachment_prompt() -> serde_json::Value { serde_json::json!([ {"type":"text", "text":"Review these files"}, @@ -2071,9 +2128,13 @@ mod tests { steer(2_500, "make it blue"), reported_turn_end(3_000, 990), ]); - assert_eq!(roles(&silent), ["User", "Assistant", "User", "User"]); + // The second prompt's reported spend belongs to its own empty result, + // never to the reply produced for the first prompt. + assert_eq!(roles(&silent), ["User", "Assistant", "User", "User", "Assistant"]); assert_eq!(silent[1].usage, None); assert_eq!(silent[1].model, None); + assert!(silent[4].blocks.is_empty()); + assert_eq!(silent[4].usage.as_ref().map(|usage| usage.input_tokens), Some(10)); } /// A call the reply started before a steer can report after it: a command