Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions docs/agents/zcode.md
Original file line number Diff line number Diff line change
Expand Up @@ -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. 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.

## Current scope
Expand Down
234 changes: 232 additions & 2 deletions src-tauri/src/acp/connection.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1811,6 +1811,7 @@ async fn record_turn_end(
stop_reason: &str,
started_at_ms: u64,
model: Option<String>,
metrics: Option<&crate::acp::zcode_turn::TurnMetrics>,
) {
let Some(dir) = transcript_dir_for(agent_type) else {
return;
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -13477,6 +13482,26 @@ async fn run_conversation_loop(
let runtime = terminal_runtime.clone();
let session_id = sid.clone();
let cwd_opt = Some(cwd);
// 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(
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
Expand Down Expand Up @@ -13729,6 +13754,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
Expand Down Expand Up @@ -13844,6 +13870,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
Expand Down Expand Up @@ -14152,6 +14179,16 @@ async fn run_conversation_loop(
.release_all_for_session(sid.0.as_ref())
.await;
tracked_terminal_tool_calls.clear();
// 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
Expand Down Expand Up @@ -40495,9 +40532,12 @@ mod tests {

impl CustomLoop {
async fn start(turns: Vec<Vec<GrokFrame>>) -> 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<Vec<GrokFrame>>) -> 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);
Expand Down Expand Up @@ -40711,6 +40751,196 @@ 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;
}

/// 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_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 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;
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;
}

/// 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
/// that request. Here the agent first re-sends its option list showing the
/// requested value — the unprompted echo that makes the broadcasts useless
Expand Down
1 change: 1 addition & 0 deletions src-tauri/src/acp/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down
3 changes: 2 additions & 1 deletion src-tauri/src/acp/registry.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
Loading
Loading