From 3e3d711f065e279dc0b983d2469b97e361a96256 Mon Sep 17 00:00:00 2001 From: daniel Date: Sat, 5 Sep 2026 23:29:21 +0100 Subject: [PATCH 1/7] feat(tui): browse and resume conversation branches --- docs/user/tui-and-sessions.md | 23 +- src/acp_child.rs | 230 ++++- src/events.rs | 258 +++++- src/main.rs | 2 + src/protocols/acp/prompt_branches.rs | 25 + src/protocols/acp/v2.rs | 419 ++++++++- src/runtime.rs | 5 +- src/runtime/diagnostics.rs | 453 ++++++++++ src/session.rs | 40 + src/session/branch.rs | 463 +++++++++- src/tools/subagent.rs | 12 +- src/tools/subagent/tests.rs | 211 +++++ src/tui/app.rs | 916 ++++++++++++++++--- src/tui/branches.rs | 501 +++++++++++ src/tui/command.rs | 10 +- src/tui/mod.rs | 1250 +++++++++++++++++++++++--- src/tui/ui.rs | 262 +++++- 17 files changed, 4693 insertions(+), 387 deletions(-) create mode 100644 src/runtime/diagnostics.rs create mode 100644 src/tui/branches.rs diff --git a/docs/user/tui-and-sessions.md b/docs/user/tui-and-sessions.md index 130c60e..36bf977 100644 --- a/docs/user/tui-and-sessions.md +++ b/docs/user/tui-and-sessions.md @@ -44,6 +44,7 @@ Use the `artifact` tool to read spilled output, including memory-only artifacts; | `Esc` | Close the transcript navigator; cancel a pending-message edit and restore the previous draft; leave the queue selector; otherwise interrupt a running turn or dismiss an idle notice | | `F2` | Focus the pending-message queue (or return to the composer) | | `F3` | Open or close the read-only transcript navigator, including while streaming | +| `F4` | Browse saved conversation branches without replacing the editor draft | | `Up` / `Down`, `Enter`, `Backspace` / `Delete` in the queue | Select a pending message, edit it if supported, or remove it | | `Command+B` | Move the newest running foreground top-level compose call to the background | | `Ctrl+C` | Interrupt a running turn; clear a non-empty idle prompt; quit when idle with an empty prompt | @@ -80,7 +81,7 @@ Press `F3` to open the read-only transcript navigator, even while a response is | `Ctrl+Up` / `Ctrl+Down` | Select the previous / next User prompt, switching to the User filter while retaining the query | | `Enter` | Reveal and highlight the selected rendered block in the transcript | | `Esc` / `F3` | Close without changing transcript scroll position, editor text, or attachments | -| Type `/branch`, then `Enter` | While idle, open the backend-approved conversation checkpoint chooser | +| `F2` (or type `/branch`, then `Enter`) | While idle, open the backend-approved conversation checkpoint chooser; the selected display block is not a checkout address | Navigator queries are limited to 4,096 UTF-8 bytes. Oversized pastes keep a bounded whole-grapheme prefix; pasted newlines and tabs are ignored rather than activating navigator controls. @@ -92,9 +93,21 @@ Revealing a Thought block temporarily shows it even when reasoning is hidden. Br The navigator searches only the currently displayed or replayed history, not a compacted archive or undelivered messages in the pending queue. Tool searches include display text such as the title, script, and output; media labels are searchable, but binary media payloads are not. Block identities are local and ephemeral, not durable addresses for forking. This is navigation only: it does not fork a session, write history, cancel a turn, or send a prompt. +### Browse saved conversation futures + +Press **F4** to open the existing sessions dialog as a read-only branch tree. `/branches` is an alias for `/sessions`; neither creates a branch. **`/branch` (singular) creates a new conversation future from an authoritative checkpoint.** F4 preserves editor text, cursor, attachments, and any parked pending-message draft. The loading dialog opens immediately: typing and paste already belong to its search field while the catalog is scanned, and Enter cannot activate a session until the scan finishes. Esc cancels the loading dialog; a late result cannot reopen it. The typed `/sessions` or `/branches` command can consume its own command text. + +Each durable workspace session appears once, using its existing name, short ID, and latest preview. Indentation follows validated persisted parent lineage, not visible transcript block positions. The detail shows the persisted branch point separately from latest activity. Legacy sessions are independent roots. Missing parents are labeled orphans; malformed, incomplete, or cyclic lineage is shown with a warning rather than hiding unrelated sessions. Reading the tree does not repair or write transcripts or locks. + +`●` marks the **current session**; `›` marks only the **selected row**. Use Up/Down to select, `/` to edit a case-insensitive name/ID/preview search, and Enter or Esc to finish editing the search. Matching ancestors remain as labeled context. Queries are bounded to 4,096 UTF-8 bytes and indentation and scrolling are bounded for large trees. Press `r` to rename the selected visible session using the existing naming dialog; Esc cancels. Esc or F4 closes the tree without changing the draft or active session. + +Enter resumes the selected visible session only while **idle**, outside provisional checkout, submission, pending-message edits, or conflicting queued work. Browsing is still available while working. A different session cannot replace an unsent draft or attachments: send or clear them first. Selecting the current session explicitly closes and reloads committed history while retaining its unsent editor draft and attachments. Reload never reconstructs or submits cancelled input. If loading fails after the current actor closes, Kit makes one restoration attempt. If that also fails, the TUI shows **disconnected**, retains the draft and attachments, and blocks backend-dependent actions. Open F4 and select the current session again to retry; only a successful load reconnects it. + +Live diagnostic compatibility: after a session reload or switch back, Kit rejects delayed diagnostics from earlier activations. A reused legacy ACP v1 child shares one stderr stream without per-operation identity, so its nested live stderr details remain tied to its original activation and are also rejected. Parent-owned subagent roster events, explicit closes, and tool results still update normally. Newly launched children report live details under the new activation. + ### Branch from a text checkpoint in a new session -While idle, type the exact local command `/branch`, or open `/transcript` (`F3`), type `/branch` in its search field, and press Enter. The separate checkout chooser lists only conversation checkpoints approved by the backend, labeled `[user]`, `[assistant]`, or `[tool]`. Archived checkpoints also show `[archived]`; this list is authoritative, not inferred from the visible transcript. Press `0` for All, `1` for User, `2` for Assistant, or `3` for Tool. Changing the filter selects its first match. Use Up/Down to select a checkpoint and Enter to prepare a draft; an empty filter has nothing to prepare. Filtering never changes the backend address used for checkout. Unsupported agents or ineligible checkpoints produce an error without changing the source session. `/transcript` remains a display-only navigator: its local block positions are not checkout addresses. +While idle, type the exact local command `/branch`, or open `/transcript` (`F3`) and press `F2` (typing `/branch` in its search field and pressing Enter also works). The separate checkout chooser lists only conversation checkpoints approved by the backend, labeled `[user]`, `[assistant]`, or `[tool]`. Archived checkpoints also show `[archived]`; this list is authoritative, not inferred from the visible transcript. Press `0` for All, `1` for User, `2` for Assistant, or `3` for Tool. Changing the filter selects its first match. Use Up/Down to select a checkpoint and Enter to prepare a draft; an empty filter has nothing to prepare. Filtering never changes the backend address used for checkout. Unsupported agents or ineligible checkpoints produce an error without changing the source session. `/transcript` remains a display-only navigator: its local block positions are not checkout addresses. > Only conversation context changes. Filesystem changes, running processes, and external effects are not rolled back. @@ -108,7 +121,7 @@ The child retains the safe prefix before the selected user prompt or after the s An unsubmitted checkout becomes stale when its source conversation or configuration changes, including compaction, or when the backend restarts. Abandon it and list checkpoints again. Committed submissions survive restart: retrying the same checkout and text finds the same child without generating a second response. Errors after a durable commit identify the child so it remains discoverable even if activation or response delivery failed. -After successful submission, the source remains loaded and unchanged. Use `/sessions` to return to it. If the child is cancelled before execution starts, its committed history remains available but its connection can close to prevent the cancelled prompt from running later. Select the source, then the child in `/sessions` to reload it without rerunning that prompt. Checkout does not restore files, stop processes, reverse tool calls, or undo any other external effect. +After successful submission, the source remains loaded and unchanged. Use `/sessions` to return to it. If the child is cancelled before execution starts, its committed history remains available but its connection can close to prevent the cancelled prompt from running later. Select the current child in `/branches` and press Enter to explicitly close and reload its committed history without rerunning that prompt. An unsent editor draft and its attachments survive this same-session reload; they are not submitted automatically. Checkout does not restore files, stop processes, reverse tool calls, or undo any other external effect. ### Edit or remove a pending message @@ -158,7 +171,7 @@ When the agent roster is visible, terminals at least 108 columns wide show the t ## Manage sessions and compact from the TUI -The TUI handles `/new`, `/resume`, `/sessions`, `/close`, `/model`, `/effort`, `/agents`, `/transcript`, and `/branch` as exact local slash-command tokens. It also discovers agent commands through ACP and highlights them without interpreting them locally: +The TUI handles `/new`, `/resume`, `/sessions`, `/branches`, `/close`, `/model`, `/effort`, `/agents`, `/transcript`, and `/branch` as exact local slash-command tokens. It also discovers agent commands through ACP and highlights them without interpreting them locally: ```text /new @@ -175,7 +188,7 @@ The TUI handles `/new`, `/resume`, `/sessions`, `/close`, `/model`, `/effort`, ` /transcript ``` -Except for `/transcript`, which also opens while streaming, these local commands are available only while the session is idle. `/agents` toggles the agent roster without starting a model turn. `/new` closes the current session and starts a fresh persisted session. It clears the visible transcript but does not delete or alter the previous session, which remains resumable by its ID. Text following `/new` becomes the new session's first prompt. `/resume ` closes the current session, resumes the requested durable session, and replays its transcript; selecting the already-active ID is a no-op. `/sessions` opens a visible newest-first selector for the same workspace. Up and Down move, Enter uses the existing resume flow, `R` opens an inline rename field, and Esc cancels renaming or closes the dialog. Submit an empty rename and confirm to clear the custom name. After a save, the picker remains open on the selected session and refreshes its displayed name. `/close` closes the current session and exits the TUI. +Read-only `/transcript`, `/sessions`, and `/branches` also open while streaming; switching sessions and creating a branch require idle state. `/agents` toggles the agent roster without starting a model turn. `/new` closes the current session and starts a fresh persisted session. It clears the visible transcript but does not delete or alter the previous session, which remains resumable by its ID. Text following `/new` becomes the new session's first prompt. `/resume ` closes the current session, resumes the requested durable session, and replays its transcript; selecting the already-active ID explicitly closes and reloads committed history while preserving the unsent editor draft and attachments. `/sessions` (alias `/branches`, shortcut F4) opens the derived branch tree for the same workspace, with newest-activity sibling ordering. Up and Down move, Enter uses the existing resume flow, `R` opens an inline rename field, and Esc cancels renaming or closes the dialog. Submit an empty rename and confirm to clear the custom name. After a save, the picker remains open on the selected session and refreshes its displayed name. `/close` closes the current session and exits the TUI. `/model` opens the model selector. `/effort` opens the advertised ACP reasoning-effort selector; `/effort default|low|medium|high` selects directly. In either dialog, Tab toggles saving the selection to `~/.kit/config.toml`, Enter selects, and Esc closes. Saving `default` removes top-level `reasoning_effort`; other values update it without replacing unrelated TOML. A new or resumed process starts from the resolved CLI/TOML default unless the selection was saved. diff --git a/src/acp_child.rs b/src/acp_child.rs index 1d49918..c81839f 100644 --- a/src/acp_child.rs +++ b/src/acp_child.rs @@ -555,7 +555,7 @@ impl ChildSession { let (ready_tx, mut ready_rx) = oneshot::channel(); let (closed_tx, closed_rx) = watch::channel(false); let actor_tx = tx.clone(); - let mut task = tokio::spawn(async move { + let mut task = tokio::spawn(crate::events::inherit_diagnostics(async move { let result = run( RunConfig { config, @@ -572,7 +572,7 @@ impl ChildSession { .await; let _ = closed_tx.send(true); result - }); + })); let result = tokio::select! { ready = &mut ready_rx => match ready { Ok(Ok(ready)) => Ok(ready), @@ -621,6 +621,9 @@ impl ChildSession { pub async fn close(&self) -> Result<(), ChildError> { if let Some(ancestor_id) = &self.descendant_parent { + // Explicit close belongs to its caller, which may be a new turn + // using a reused or native-forked child. Deferred cleanup callers + // retain their originating scope across their own spawn boundary. crate::events::emit(&crate::events::RuntimeEvent::SubagentDescendantsRemoved { ancestor_id: ancestor_id.clone(), }); @@ -784,23 +787,8 @@ async fn run( .stderr .take() .ok_or("could not open ACP harness stderr")?; - let label = harness.clone(); - let ancestor_id = config.parent_id.clone(); - let descendant_parent = ancestor_id.clone(); - tokio::spawn(async move { - forward_stderr( - stderr, - &label, - ancestor_id.as_deref(), - |output| match output { - ForwardedStderr::RuntimeLine(line) | ForwardedStderr::Diagnostic(line) => { - eprintln!("{line}"); - } - ForwardedStderr::Cleanup(event) => crate::events::emit(&event), - }, - ) - .await; - }); + let descendant_parent = config.parent_id.clone(); + spawn_forward_stderr(stderr, harness.clone(), config.parent_id.clone()); let transport = ByteStreams::new(stdin.compat_write(), stdout.compat()); let routes = Arc::new(Mutex::new( HashMap::>>::new(), @@ -1116,6 +1104,41 @@ async fn run( }) } +// Capture before spawn. ACP v1 stderr is one process-wide stream shared by +// reused sessions and native forks. Neither a prompt response nor a new request +// identifies the operation that produced buffered bytes; even child activation +// IDs belong to another process and have no parent-operation mapping here. +// Keep the launch scope for this stream and its EOF cleanup. After reactivation, +// the reader must reject these events, including otherwise-current nested +// diagnostics. This intentionally sacrifices nested live detail rather than +// retagging delayed old bytes. Parent-owned Subagents lifecycle events and +// explicit ChildSession::close use their invocation scope and remain current. +fn spawn_forward_stderr( + stderr: impl AsyncRead + Unpin + Send + 'static, + label: String, + ancestor_id: Option, +) -> tokio::task::JoinHandle<()> { + let diagnostics = crate::events::DiagnosticScope::capture(); + tokio::spawn(async move { + forward_stderr(stderr, &label, ancestor_id.as_deref(), |output| { + match output { + ForwardedStderr::RuntimeLine(line) => { + // Child epochs belong to another process. Re-envelope the + // event with this parent's captured activation at the IPC + // boundary; recursively forwarded roster events keep their + // own subagent/ancestor IDs in the unchanged payload. + if let Some(event) = crate::events::parse(&line) { + diagnostics.emit(&event); + } + } + ForwardedStderr::Diagnostic(line) => eprintln!("{line}"), + ForwardedStderr::Cleanup(event) => diagnostics.emit(&event), + } + }) + .await; + }) +} + #[derive(Debug, PartialEq)] enum ForwardedStderr { RuntimeLine(String), @@ -2199,6 +2222,175 @@ mod tests { } } + #[test] + fn spawned_stderr_keeps_origin_activation_and_unscoped_is_none() { + use crate::events::{self, DiagnosticActivation, RuntimeEvent}; + const PROBE: &str = "KIT_TEST_STDERR_ACTIVATION_PROBE"; + if std::env::var_os(PROBE).is_some() { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap() + .block_on(async { + use tokio::io::AsyncWriteExt; + for route in ["same", "roundtrip", "construction", "unscoped", "current"] { + let session = format!("stderr-{route}"); + events::activate_diagnostics(&session); + if route == "current" { + events::activate_diagnostics("stderr-current-other"); + events::activate_diagnostics(&session); + } + let (mut writer, reader) = tokio::io::duplex(4096); + // Also exercise a future constructed before reactivation + // but first polled afterward (the construction case). + // The parent must reactivate and supply stderr bytes + // before awaiting the returned forwarder task. + #[allow(clippy::async_yields_async)] + let producer = events::scope_diagnostics(&session, async move { + spawn_forward_stderr(reader, "acp.kit".into(), Some(route.into())) + }); + let (task, delayed) = if route == "construction" { + (None, Some(producer)) + } else if route == "unscoped" { + drop(producer); + let (new_writer, reader) = tokio::io::duplex(4096); + writer = new_writer; + ( + Some(spawn_forward_stderr( + reader, + "acp.kit".into(), + Some(route.into()), + )), + None, + ) + } else { + (Some(producer.await), None) + }; + if route == "roundtrip" { + events::activate_diagnostics("stderr-other"); + } + if route != "current" { + events::activate_diagnostics(&session); + } + let task = match task { + Some(task) => task, + None => delayed.unwrap().await, + }; + // Actual stderr transport gates both forwarded events + // and EOF cleanup until the new activation is active. + let event = RuntimeEvent::SubagentDescendantsRemoved { + ancestor_id: format!("forward-{route}"), + }; + let child_line = events::DiagnosticEvent { + event, + activation: Some(DiagnosticActivation { + session_id: "child-process".into(), + epoch: 999, + }), + }; + writer + .write_all( + format!( + "{}{}\n", + events::EVENT_MARKER, + serde_json::to_string(&child_line).unwrap() + ) + .as_bytes(), + ) + .await + .unwrap(); + writer.shutdown().await.unwrap(); + task.await.unwrap(); + } + events::scope_diagnostics("stderr-same", async { + events::emit(&RuntimeEvent::StorageStatus { + pending: true, + exhausted: false, + }); + }) + .await; + events::emit(&RuntimeEvent::CompactionStarted { + reason: "unscoped".into(), + at: 1, + }); + }); + return; + } + let output = std::process::Command::new(std::env::current_exe().unwrap()) + .args([ + "--exact", + "acp_child::tests::spawned_stderr_keeps_origin_activation_and_unscoped_is_none", + "--nocapture", + ]) + .env(PROBE, "1") + .env(events::EVENTS_ENV, "1") + .output() + .unwrap(); + assert!( + output.status.success(), + "{}", + String::from_utf8_lossy(&output.stderr) + ); + let stderr = String::from_utf8(output.stderr).unwrap(); + let diagnostics = stderr + .lines() + .filter_map(events::parse_diagnostic) + .collect::>(); + for route in ["same", "roundtrip", "construction", "unscoped", "current"] { + let session = format!("stderr-{route}"); + let first = diagnostics + .iter() + .find_map(|line| match &line.event { + RuntimeEvent::SessionStarted { session_id } if session_id == &session => { + line.activation.clone() + } + _ => None, + }) + .expect("missing initial activation"); + let expected = if route == "current" { + diagnostics + .iter() + .rev() + .find_map(|line| match &line.event { + RuntimeEvent::SessionStarted { session_id } if session_id == &session => { + line.activation.clone() + } + _ => None, + }) + .unwrap() + } else { + first + }; + for ancestor in [route.to_string(), format!("forward-{route}")] { + let line = diagnostics.iter().find(|line| matches!(&line.event, + RuntimeEvent::SubagentDescendantsRemoved { ancestor_id } if ancestor_id == &ancestor)) + .expect("missing forwarded event or EOF cleanup"); + assert_eq!( + line.activation, + (route != "unscoped").then(|| expected.clone()) + ); + } + } + for line in diagnostics.iter().filter(|line| { + matches!( + line.event, + RuntimeEvent::StorageStatus { .. } | RuntimeEvent::CompactionStarted { .. } + ) + }) { + assert_eq!(line.activation, None); + } + assert!( + diagnostics + .iter() + .any(|line| matches!(line.event, RuntimeEvent::StorageStatus { .. })) + ); + assert!( + diagnostics + .iter() + .any(|line| matches!(line.event, RuntimeEvent::CompactionStarted { .. })) + ); + } + mod removes_descendants_on_exit { use super::*; diff --git a/src/events.rs b/src/events.rs index bf59995..7f1ab66 100644 --- a/src/events.rs +++ b/src/events.rs @@ -14,8 +14,9 @@ //! ACP hosts never see the extra chatter. use std::{ + collections::HashMap, io::Write, - sync::OnceLock, + sync::{Mutex, OnceLock}, time::{SystemTime, UNIX_EPOCH}, }; @@ -101,7 +102,7 @@ pub enum GenerationOutcome { } impl RuntimeEvent { - /// Whether a parent Kit runtime should forward this child event unchanged. + /// Whether a parent Kit runtime should forward this child event payload. pub(crate) fn forward_from_child(&self) -> bool { matches!( self, @@ -135,21 +136,191 @@ pub fn enabled() -> bool { *ENABLED.get_or_init(|| std::env::var_os(EVENTS_ENV).is_some()) } -/// Writes one event to stderr when emission is enabled. -pub fn emit(event: &RuntimeEvent) { - if !enabled() { - return; +/// Private, ephemeral stderr envelope. Persisted/replayed `RuntimeEvent` shapes +/// are unchanged. Legacy lines decode without an activation and cannot establish +/// live routing authority; historical replay still reads the unchanged event. +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] +pub struct DiagnosticEvent { + #[serde(flatten)] + pub event: RuntimeEvent, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub activation: Option, +} + +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] +pub struct DiagnosticActivation { + pub session_id: String, + pub epoch: u64, +} + +pub const ACTIVATION_META_KEY: &str = "kitRuntimeActivation"; + +#[derive(Default)] +struct DiagnosticRoutes { + next_epoch: u64, + sessions: HashMap, + active: Option, +} + +fn diagnostic_routes() -> &'static Mutex { + static ROUTES: OnceLock> = OnceLock::new(); + ROUTES.get_or_init(Mutex::default) +} + +/// Commit a new activation only after session admission has succeeded. The ACP +/// response returns this exact epoch; stderr alone is never activation authority. +pub fn activate_diagnostics(session_id: &str) -> u64 { + let mut routes = diagnostic_routes().lock().expect("diagnostic routes lock"); + routes.next_epoch = routes + .next_epoch + .checked_add(1) + .expect("activation epoch overflow"); + let epoch = routes.next_epoch; + routes.sessions.insert(session_id.to_owned(), epoch); + routes.active = Some(DiagnosticActivation { + session_id: session_id.to_owned(), + epoch, + }); + write_diagnostic_marker(&routes); + epoch +} + +#[cfg(test)] +pub(crate) fn diagnostic_epoch(session_id: &str) -> Option { + diagnostic_routes() + .lock() + .expect("diagnostic routes lock") + .sessions + .get(session_id) + .copied() +} + +/// Restore an actor's route without minting a new activation on each turn. +pub fn restore_diagnostics(session_id: &str) { + let mut routes = diagnostic_routes().lock().expect("diagnostic routes lock"); + routes.active = routes + .sessions + .get(session_id) + .map(|epoch| DiagnosticActivation { + session_id: session_id.to_owned(), + epoch: *epoch, + }); + write_diagnostic_marker(&routes); +} + +tokio::task_local! { + static DIAGNOSTIC_ACTIVATION: Option; +} + +/// Preserve the turn's activation even if another actor restores its route, or +/// a loaded-session resume reactivates this durable ID while the turn is live. +/// Capture synchronously at construction, before the future can be spawned. +pub fn scope_diagnostics( + session_id: &str, + future: F, +) -> impl std::future::Future + use { + let activation = { + let routes = diagnostic_routes().lock().expect("diagnostic routes lock"); + routes + .sessions + .get(session_id) + .map(|epoch| DiagnosticActivation { + session_id: session_id.to_owned(), + epoch: *epoch, + }) + }; + DiagnosticScope(activation).scope(future) +} + +fn write_diagnostic_marker(routes: &DiagnosticRoutes) { + if enabled() + && let Some(activation) = &routes.active + { + write_diagnostic( + &mut std::io::stderr().lock(), + &DiagnosticEvent { + event: RuntimeEvent::SessionStarted { + session_id: activation.session_id.clone(), + }, + activation: Some(activation.clone()), + }, + ); } - let mut stderr = std::io::stderr().lock(); - write_event(&mut stderr, event); } -fn write_event(writer: &mut impl Write, event: &RuntimeEvent) { +/// Immutable producer identity. Capture before constructing a spawned task, not +/// when that task is first polled. Unscoped producers never borrow the active UI +/// route, including after a durable session ID is reactivated. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub(crate) struct DiagnosticScope(Option); + +impl DiagnosticScope { + pub(crate) fn capture() -> Self { + Self(DIAGNOSTIC_ACTIVATION.try_with(Clone::clone).unwrap_or(None)) + } + + pub(crate) fn scope( + &self, + future: F, + ) -> impl std::future::Future + use { + DIAGNOSTIC_ACTIVATION.scope(self.0.clone(), future) + } + + pub(crate) fn emit(&self, event: &RuntimeEvent) { + if !enabled() { + return; + } + write_diagnostic( + &mut std::io::stderr().lock(), + &DiagnosticEvent { + event: event.clone(), + activation: if matches!(event, RuntimeEvent::StorageStatus { .. }) { + None + } else { + self.0.clone() + }, + }, + ); + } +} + +/// Inherit the caller's scope synchronously, before the returned future can be +/// moved into another task. Tokio task locals do not propagate through spawn. +pub(crate) fn inherit_diagnostics( + future: F, +) -> impl std::future::Future + use { + DiagnosticScope::capture().scope(future) +} + +/// Writes with the originating producer's scope. Truly unscoped events remain +/// unscoped; process-wide StorageStatus is global even inside a session scope. +pub fn emit(event: &RuntimeEvent) { + DiagnosticScope::capture().emit(event); +} + +fn write_diagnostic(writer: &mut impl Write, event: &DiagnosticEvent) { if let Ok(line) = serde_json::to_string(event) { let _ = writeln!(writer, "{EVENT_MARKER}{line}"); } } +#[cfg(test)] +fn write_event(writer: &mut impl Write, event: &RuntimeEvent) { + write_diagnostic( + writer, + &DiagnosticEvent { + event: event.clone(), + activation: None, + }, + ); +} + +/// Decode both current activation-aware and historical unscoped stderr lines. +#[must_use] +pub fn parse_diagnostic(line: &str) -> Option { + serde_json::from_str(line.strip_prefix(EVENT_MARKER)?).ok() +} + /// Parses one stderr line, returning an event when the line carries one. #[must_use] pub fn parse(line: &str) -> Option { @@ -227,10 +398,75 @@ mod tests { use serde_json::json; use super::{ - EVENT_MARKER, GenerationOutcome, RuntimeEvent, SubagentStatus, parse, summarize_input, - summarize_output, write_event, + DIAGNOSTIC_ACTIVATION, DiagnosticActivation, DiagnosticEvent, EVENT_MARKER, + GenerationOutcome, RuntimeEvent, SubagentStatus, activate_diagnostics, parse, + parse_diagnostic, restore_diagnostics, scope_diagnostics, summarize_input, + summarize_output, write_diagnostic, write_event, }; + #[test] + fn diagnostic_envelope_preserves_legacy_replay_and_rejects_malformed_identity() { + let legacy = format!( + "{EVENT_MARKER}{}", + r#"{"event":"child_started","call":"p:compose:c","tool":"shell","summary":"ok","at":1}"# + ); + let event = parse(&legacy).unwrap(); + assert_eq!(parse_diagnostic(&legacy).unwrap().activation, None); + let current = DiagnosticEvent { + event: event.clone(), + activation: Some(DiagnosticActivation { + session_id: "source".into(), + epoch: 42, + }), + }; + let mut bytes = Vec::new(); + write_diagnostic(&mut bytes, ¤t); + let line = String::from_utf8(bytes).unwrap(); + assert_eq!(parse_diagnostic(line.trim_end()), Some(current)); + // Old child-forwarding/replay readers retain exactly their event shape. + assert_eq!(parse(line.trim_end()), Some(event)); + assert!(line.contains(r#""activation":{"session_id":"source","epoch":42}"#)); + for activation in [ + r#"{"session_id":"source","epoch":"42"}"#, + r#"{"session_id":"source","epoch":-1}"#, + r#"{"epoch":42}"#, + ] { + let malformed = format!( + "{EVENT_MARKER}{{\"event\":\"session_started\",\"session_id\":\"source\",\"activation\":{activation}}}" + ); + assert!(parse_diagnostic(&malformed).is_none()); + } + } + + #[tokio::test] + async fn diagnostic_epochs_change_on_reactivation_not_turn_restoration() { + let session_id = "event-epoch-test"; + let first = activate_diagnostics(session_id); + restore_diagnostics(session_id); + scope_diagnostics(session_id, async { + assert_eq!( + DIAGNOSTIC_ACTIVATION.with(|activation| activation.as_ref().unwrap().epoch), + first + ); + let second = activate_diagnostics(session_id); + assert!(second > first); + // Already-running futures keep their original activation identity. + assert_eq!( + DIAGNOSTIC_ACTIVATION.with(|activation| activation.as_ref().unwrap().epoch), + first + ); + restore_diagnostics(session_id); + scope_diagnostics(session_id, async { + assert_eq!( + DIAGNOSTIC_ACTIVATION.with(|activation| activation.as_ref().unwrap().epoch), + second + ); + }) + .await; + }) + .await; + } + #[test] fn storage_status_events_round_trip_without_session_affinity() { for (pending, exhausted) in [(true, false), (false, false), (true, true)] { diff --git a/src/main.rs b/src/main.rs index 7c0a466..1b4427b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1928,6 +1928,8 @@ future_option = true title: Some("OAuth token bug".into()), preview: Some("Fix tests in the catalog".into()), is_subagent: false, + lineage: kit::session::CatalogLineage::Root, + branch_point: None, updated_at: 0, }]); assert_eq!( diff --git a/src/protocols/acp/prompt_branches.rs b/src/protocols/acp/prompt_branches.rs index 92aada0..c4da78a 100644 --- a/src/protocols/acp/prompt_branches.rs +++ b/src/protocols/acp/prompt_branches.rs @@ -63,6 +63,9 @@ pub(crate) struct SubmitPromptBranchRequest { #[derive(Debug, Clone, Serialize, Deserialize, agent_client_protocol::JsonRpcResponse)] pub(crate) struct SubmitPromptBranchResponse { + /// Ephemeral diagnostic identity; old servers omit this private extension. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub runtime_activation: Option, pub session_id: wire::SessionId, pub config_options: Vec, } @@ -694,6 +697,28 @@ mod tests { assert!(checkout_candidate(&source[5]).is_none()); } + #[test] + fn submit_response_activation_is_optional_and_roundtrips_without_durable_changes() { + let legacy = json!({"session_id": "child", "config_options": []}); + let parsed: SubmitPromptBranchResponse = serde_json::from_value(legacy.clone()).unwrap(); + assert_eq!(parsed.runtime_activation, None); + assert_eq!(serde_json::to_value(parsed).unwrap(), legacy); + + let mut current = legacy.clone(); + current["runtime_activation"] = json!(42); + let parsed: SubmitPromptBranchResponse = serde_json::from_value(current.clone()).unwrap(); + assert_eq!(parsed.runtime_activation, Some(42)); + assert_eq!(serde_json::to_value(parsed).unwrap(), current); + + current["runtime_activation"] = serde_json::Value::Null; + let parsed: SubmitPromptBranchResponse = serde_json::from_value(current.clone()).unwrap(); + assert_eq!(serde_json::to_value(parsed).unwrap(), legacy); + for malformed in [json!(-1), json!("42"), json!({"epoch": 42})] { + current["runtime_activation"] = malformed; + assert!(serde_json::from_value::(current.clone()).is_err()); + } + } + #[test] fn role_wire_defaults_legacy_and_writes_current_shape() { let legacy = json!({"address":"opaque", "text":"preview", "historical":false}); diff --git a/src/protocols/acp/v2.rs b/src/protocols/acp/v2.rs index 1b5da66..218d61b 100644 --- a/src/protocols/acp/v2.rs +++ b/src/protocols/acp/v2.rs @@ -55,9 +55,13 @@ static NEXT_ERROR_MESSAGE_ID: AtomicU64 = AtomicU64::new(1); static NEXT_THOUGHT_MESSAGE_ID: AtomicU64 = AtomicU64::new(1); fn establish_diagnostic_route(session_id: &wire::SessionId) { - crate::events::emit(&crate::events::RuntimeEvent::SessionStarted { - session_id: session_id.to_string(), - }); + crate::events::restore_diagnostics(&session_id.to_string()); +} + +fn activation_meta(epoch: u64) -> serde_json::Map { + [(crate::events::ACTIVATION_META_KEY.to_owned(), epoch.into())] + .into_iter() + .collect() } fn validate_resume_location( @@ -707,6 +711,7 @@ struct AttachedSession { session_id: wire::SessionId, config_options: Vec, canonical_transcript: Vec, + runtime_activation: u64, activation: oneshot::Sender<()>, } @@ -943,11 +948,14 @@ impl Server { let AttachedSession { session_id, config_options, + runtime_activation, activation, .. } = attached; Ok(( - wire::NewSessionResponse::new(session_id).config_options(config_options), + wire::NewSessionResponse::new(session_id) + .config_options(config_options) + .meta(activation_meta(runtime_activation)), activation, )) } @@ -987,7 +995,9 @@ impl Server { Vec::new() }; return Ok(( - wire::ResumeSessionResponse::new().config_options(attached.config_options), + wire::ResumeSessionResponse::new() + .config_options(attached.config_options) + .meta(activation_meta(attached.runtime_activation)), updates, attached.activation, )); @@ -1025,7 +1035,9 @@ impl Server { Vec::new() }; Ok(( - wire::ResumeSessionResponse::new().config_options(attached.config_options), + wire::ResumeSessionResponse::new() + .config_options(attached.config_options) + .meta(activation_meta(attached.runtime_activation)), updates, attached.activation, )) @@ -1242,12 +1254,13 @@ impl Server { SessionPublicationError::Commit(error) => error, }); } - establish_diagnostic_route(&session_id); + let runtime_activation = crate::events::activate_diagnostics(&session_id.to_string()); drop(actor_task); Ok(AttachedSession { session_id, config_options, canonical_transcript, + runtime_activation, activation, }) } @@ -1280,12 +1293,13 @@ impl Server { .map_err(|_| AcpRuntimeError::ClientClosed)??; // Reattachment can switch the client's diagnostic route without creating // an actor. Write the source marker before returning the load response. - establish_diagnostic_route(session_id); + let runtime_activation = crate::events::activate_diagnostics(&session_id.to_string()); let (activation, _already_active) = oneshot::channel(); Ok(Some(AttachedSession { session_id: session_id.clone(), config_options, canonical_transcript, + runtime_activation, activation, })) } @@ -1357,6 +1371,7 @@ impl Server { let response = SubmitPromptBranchResponse { session_id: attached.session_id.clone(), config_options: attached.config_options.clone(), + runtime_activation: Some(attached.runtime_activation), }; return Ok((response, Some(attached))); } @@ -1377,6 +1392,7 @@ impl Server { let response = SubmitPromptBranchResponse { session_id: attached.session_id.clone(), config_options: attached.config_options.clone(), + runtime_activation: Some(attached.runtime_activation), }; // A durable retry may reattach, but never starts another generation. return Ok((response, Some(attached))); @@ -1415,6 +1431,7 @@ impl Server { let response = SubmitPromptBranchResponse { session_id: attached.session_id.clone(), config_options: attached.config_options.clone(), + runtime_activation: Some(attached.runtime_activation), }; Ok((response, Some(attached))) } @@ -2410,8 +2427,9 @@ async fn run_active_turn( // after the client returns to its source without sending session/resume. // Re-establish identity before activity, model, or cleanup diagnostics. establish_diagnostic_route(session_id); - activity - .execute( + crate::events::scope_diagnostics( + &session_id.to_string(), + activity.execute( origin, async { let result = drive_prompt( @@ -2441,8 +2459,9 @@ async fn run_active_turn( result }, |reason| Some(reason.clone()), - ) - .await + ), + ) + .await } // False means admission was denied; the actor must retain the selected wake. @@ -3952,6 +3971,257 @@ mod tests { cancelled_gated_branch_activation(false, GatedTurn::Autonomous).await; } + #[tokio::test] + async fn cancelled_child_recovers_through_close_and_actual_resume_then_executes_once() { + const CHILD: &str = "KIT_TEST_CANCELLED_CHILD_ACTUAL_RESUME"; + if std::env::var_os(CHILD).is_none() { + // Isolate provider environment overrides in a subprocess. The real + // runtime/load driver uses a local model endpoint, not a substituted + // Agent::builder transcript or a test-only attachment implementation. + let requests = Arc::new(Mutex::new(Vec::::new())); + let recorded = requests.clone(); + let model = axum::Router::new() + .route("/models", axum::routing::get(|| async { + axum::Json(json!({"data": [{"id": "test-model", "name": "test-model", "context_length": 8192}]})) + })) + .route("/chat/completions", axum::routing::post(move |axum::Json(request): axum::Json| { + let recorded = recorded.clone(); + async move { + recorded.lock().unwrap().push(request); + ([ (axum::http::header::CONTENT_TYPE, "text/event-stream") ], concat!( + "data: {\"id\":\"reloaded\",\"object\":\"chat.completion.chunk\",\"model\":\"test-model\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"recovered answer\"},\"finish_reason\":null}]}\n\n", + "data: {\"id\":\"reloaded\",\"object\":\"chat.completion.chunk\",\"model\":\"test-model\",\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n", + "data: [DONE]\n\n" + )) + } + })); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + axum::serve(listener, model).await.unwrap(); + }); + let output = tokio::process::Command::new(std::env::current_exe().unwrap()) + .args(["--exact", "protocols::acp::v2::tests::cancelled_child_recovers_through_close_and_actual_resume_then_executes_once", "--nocapture"]) + .env(CHILD, "1") + .env("OPENROUTER_API_KEY", "test-key") + .env("OPENROUTER_BASE_URL", format!("http://{address}/chat/completions")) + .kill_on_drop(true) + .output(); + let output = timeout(Duration::from_secs(45), output) + .await + .unwrap() + .unwrap(); + server.abort(); + assert!( + output.status.success(), + "{}\n{}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); + let requests = requests.lock().unwrap(); + assert_eq!( + requests.len(), + 1, + "cancelled prompt must make zero model calls; new prompt exactly one" + ); + let last_user = requests[0]["messages"] + .as_array() + .unwrap() + .iter() + .rev() + .find(|message| message["role"] == "user") + .unwrap(); + assert!( + last_user["content"] + .to_string() + .contains("new prompt after recovery") + ); + return; + } + + let root = tempfile::tempdir().unwrap(); + let source_id = crate::session::new_id(); + drop( + crate::session::open( + root.path(), + &source_id, + false, + false, + vec![ + Item::text(ItemKind::System, "system"), + Item::text(ItemKind::User, "original prompt"), + Item::text(ItemKind::Assistant, "original answer"), + ], + ) + .unwrap(), + ); + let runtime = Runtime::new_with_provider_and_credentials( + root.path(), + "test-model", + ProviderKind::OpenRouter, + crate::credentials::CredentialStorage::Memory, + ) + .unwrap(); + let state = Arc::new(Server::new(runtime, SessionRegistry::new())); + let (client_transport, agent_transport) = agent_client_protocol::Channel::duplex(); + let (done, finished) = oneshot::channel(); + let client = agent_client_protocol::Client + .v2() + .on_receive_notification( + async move |_update: wire::UpdateSessionNotification, _cx| Ok(()), + agent_client_protocol::on_receive_notification!(), + ) + .connect_with(client_transport, async move |_cx| { + let _ = finished.await; + Ok(()) + }); + let agent = agent_client_protocol::Agent.v2().connect_with( + agent_transport, + async move |connection| { + let workspace = root.path().to_path_buf(); + let (_, _, activate_source) = state + .resume_session( + wire::ResumeSessionRequest::new(source_id.clone(), workspace.clone()), + connection.clone(), + ) + .await + .unwrap(); + activate_source.send(()).unwrap(); + let listed = state + .list_prompt_branches(ListPromptBranchesRequest { + session_id: wire::SessionId::new(source_id.clone()), + }) + .await + .unwrap(); + let prepared = state + .prepare_prompt_branch(PreparePromptBranchRequest { + session_id: wire::SessionId::new(source_id.clone()), + address: listed.boundaries[0].address.clone(), + }) + .await + .unwrap(); + let (submitted, attached) = state + .submit_prompt_branch( + SubmitPromptBranchRequest { + session_id: wire::SessionId::new(source_id.clone()), + checkout_token: prepared.checkout_token, + text: "cancelled edited prompt".into(), + }, + connection.clone(), + ) + .await + .unwrap(); + let child_id = submitted.session_id; + let retired_commands = state.sender(&child_id).unwrap(); + state + .cancel(wire::CancelSessionNotification::new(child_id.clone())) + .await + .unwrap(); + attached.unwrap().activation.send(()).unwrap(); + timeout(Duration::from_secs(5), retired_commands.closed()) + .await + .unwrap(); + // An actor may already have removed itself from the map. The common + // client close path treats precisely this not-found case as closed. + let closed = state + .close(wire::CloseSessionRequest::new(child_id.clone())) + .await; + assert!(matches!( + closed, + Ok(_) | Err(AcpRuntimeError::SessionNotFound(_)) + )); + let (resumed, replay, activation) = state + .resume_session( + wire::ResumeSessionRequest::new(child_id.clone(), workspace.clone()) + .replay_from(wire::ReplayFrom::Start(wire::ReplayFromStart::new())), + connection.clone(), + ) + .await + .unwrap(); + assert_ne!( + resumed.meta.as_ref().unwrap()[crate::events::ACTIVATION_META_KEY].as_u64(), + submitted.runtime_activation + ); + assert!(!replay.is_empty()); + activation.send(()).unwrap(); + let busy = state + .sessions + .lock() + .unwrap() + .get(&child_id) + .unwrap() + .busy + .clone(); + assert!( + !busy.load(Ordering::Acquire), + "actual resume must be passive" + ); + let history = + crate::session::branch::load_history(&workspace, &child_id.to_string()) + .unwrap(); + assert!( + !history + .last() + .unwrap() + .iter() + .any(|item| item.kind == ItemKind::Assistant) + ); + let start = state + .prepare_prompt(wire::PromptRequest::new( + child_id.clone(), + vec![wire::ContentBlock::Text(wire::TextContent::new( + "new prompt after recovery", + ))], + )) + .await + .unwrap(); + assert!(busy.load(Ordering::Acquire)); + start.send(()).unwrap(); + timeout(Duration::from_secs(10), async { + while busy.load(Ordering::Acquire) { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + state + .close(wire::CloseSessionRequest::new(child_id.clone())) + .await + .unwrap(); + state + .close(wire::CloseSessionRequest::new(source_id)) + .await + .unwrap(); + let history = + crate::session::branch::load_history(&workspace, &child_id.to_string()) + .unwrap(); + assert_eq!( + history + .last() + .unwrap() + .iter() + .filter(|item| item.kind == ItemKind::Assistant) + .count(), + 1 + ); + assert!( + serde_json::to_string(&history) + .unwrap() + .contains("recovered answer") + ); + let _ = done.send(()); + Ok(()) + }, + ); + let (agent, client) = timeout(Duration::from_secs(30), async { + tokio::join!(agent, client) + }) + .await + .unwrap(); + agent.unwrap(); + client.unwrap(); + } + #[test] fn resume_location_rejects_invalid_routes_before_loaded_lookup() { let root = tempfile::tempdir().unwrap(); @@ -7236,12 +7506,61 @@ mod tests { )) .block_task() .await?; - cx.send_request( - wire::ResumeSessionRequest::new(source_id.clone(), workspace.clone()) - .replay_from(wire::ReplayFrom::Start(wire::ReplayFromStart::new())), - ) - .block_task() - .await?; + let opened = cx + .send_request( + wire::ResumeSessionRequest::new(source_id.clone(), workspace.clone()) + .replay_from(wire::ReplayFrom::Start(wire::ReplayFromStart::new())), + ) + .block_task() + .await?; + let first_epoch = opened.meta.as_ref().unwrap()[crate::events::ACTIVATION_META_KEY] + .as_u64() + .unwrap(); + assert_eq!( + crate::events::diagnostic_epoch(&source_id), + Some(first_epoch) + ); + // Rejected requests must not reactivate a loaded actor. Exercise + // every pre-lookup validation, including unknown replay cursors. + let mut directories = + wire::ResumeSessionRequest::new(source_id.clone(), workspace.clone()); + directories + .additional_directories + .push(workspace.clone().into()); + let mut cursor = + wire::ResumeSessionRequest::new(source_id.clone(), workspace.clone()); + cursor.replay_from = + Some(serde_json::from_value(json!({"type": "_unsupported"})).unwrap()); + for invalid in [ + directories, + cursor, + wire::ResumeSessionRequest::new(source_id.clone(), workspace.join("missing")), + ] { + cx.send_request(invalid) + .block_task() + .await + .expect_err("invalid loaded resume"); + assert_eq!( + crate::events::diagnostic_epoch(&source_id), + Some(first_epoch) + ); + } + let same_id = cx + .send_request(wire::ResumeSessionRequest::new( + source_id.clone(), + workspace.clone(), + )) + .block_task() + .await?; + let same_id_epoch = same_id.meta.as_ref().unwrap() + [crate::events::ACTIVATION_META_KEY] + .as_u64() + .unwrap(); + assert!(same_id_epoch > first_epoch); + assert_eq!( + crate::events::diagnostic_epoch(&source_id), + Some(same_id_epoch) + ); let listed = cx .send_request(ListPromptBranchesRequest { session_id: wire::SessionId::new(source_id.clone()), @@ -7315,8 +7634,13 @@ mod tests { async { cx.send_request(request.clone()).block_task().await }, async { cx.send_request(request.clone()).block_task().await }, ); - assert_eq!(first?.session_id.to_string(), child_id); - assert_eq!(second?.session_id.to_string(), child_id); + let first = first?; + let second = second?; + assert_eq!(first.session_id.to_string(), child_id); + assert_eq!(second.session_id.to_string(), child_id); + assert!(first.runtime_activation.is_some()); + assert!(second.runtime_activation.is_some()); + assert_ne!(first.runtime_activation, second.runtime_activation); let changed = cx .send_request(wire::SetSessionConfigOptionRequest::new( child_id.clone(), @@ -7343,11 +7667,56 @@ mod tests { ) .block_task() .await?; - assert!( - serde_json::to_string(&resumed.config_options) - .unwrap() - .contains("high") + let returned_epoch = resumed.meta.as_ref().unwrap() + [crate::events::ACTIVATION_META_KEY] + .as_u64() + .unwrap(); + assert!(returned_epoch > same_id_epoch); + assert_eq!( + crate::events::diagnostic_epoch(&source_id), + Some(returned_epoch) ); + assert_eq!(serde_json::to_value(&resumed.config_options).unwrap()[1]["currentValue"], "high"); + // A distinct durable sibling has its own captured defaults and + // replay, even after visiting a still-loaded source and child. + let sibling_id = crate::session::new_id(); + let sibling_initial = crate::session::branch::prepare( + prefix.clone(), source_id.clone(), + crate::session::branch::Boundary::new(0, &prefix).unwrap(), + "sibling-checkout".into(), + crate::session::branch::SubmittedRequest { + id: prompt_branches::submitted_request_id(&source_id, "sibling prompt"), + selection: crate::session::branch::CapturedSelection::new(&selection, None), + }, + Item::text(ItemKind::User, "sibling prompt"), + ).unwrap(); + let sibling = crate::session::open_uncommitted(&workspace, &sibling_id, false, sibling_initial).unwrap(); + crate::session::branch::commit(&sibling.observer, &sibling.transcript).unwrap(); + let mut sibling_future = sibling.transcript.clone(); + sibling_future.push(Item::text(ItemKind::Assistant, "distinct sibling answer")); + sibling.observer.replace(&sibling_future).unwrap(); + drop(sibling); + let sibling_resume = cx.send_request( + wire::ResumeSessionRequest::new(sibling_id.clone(), workspace.clone()) + .replay_from(wire::ReplayFrom::Start(wire::ReplayFromStart::new())), + ).block_task().await?; + let sibling_epoch = sibling_resume.meta.as_ref().unwrap()[crate::events::ACTIVATION_META_KEY].as_u64().unwrap(); + assert!(sibling_epoch > returned_epoch); + assert_eq!(crate::events::diagnostic_epoch(&sibling_id), Some(sibling_epoch)); + assert_eq!(serde_json::to_value(&sibling_resume.config_options).unwrap()[1]["currentValue"], "default"); + cx.send_request(wire::CloseSessionRequest::new(sibling_id.clone())).block_task().await?; + let source_again = cx.send_request( + wire::ResumeSessionRequest::new(source_id.clone(), workspace.clone()) + .replay_from(wire::ReplayFrom::Start(wire::ReplayFromStart::new())), + ).block_task().await?; + let source_epoch = source_again.meta.as_ref().unwrap()[crate::events::ACTIVATION_META_KEY].as_u64().unwrap(); + assert!(source_epoch > sibling_epoch); + assert_eq!(crate::events::diagnostic_epoch(&source_id), Some(source_epoch)); + assert_eq!(serde_json::to_value(&source_again.config_options).unwrap()[1]["currentValue"], "high"); + assert!(updates.lock().unwrap().iter().any(|update| { + update.session_id.to_string() == sibling_id + && serde_json::to_string(&update.update).unwrap().contains("distinct sibling answer") + })); cx.send_request(wire::CloseSessionRequest::new(source_id.clone())) .block_task() .await?; @@ -8266,6 +8635,8 @@ mod tests { title: Some("Saved session".into()), preview: Some("Saved session preview".into()), is_subagent: true, + lineage: crate::session::CatalogLineage::Root, + branch_point: None, updated_at: 0, }, &PathBuf::from("/workspace"), diff --git a/src/runtime.rs b/src/runtime.rs index 0abb61d..33b4d1f 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -44,6 +44,7 @@ use crate::{ }, }; +mod diagnostics; mod input_settlement; pub(crate) use input_settlement::{InputSettlement, InputSettlingDriver}; @@ -2452,8 +2453,8 @@ fn backgroundable_spec(mut spec: ToolSpec) -> ToolSpec { spec } -fn background_task_manager() -> AsyncTaskManager { - AsyncTaskManager::new().routing(background_route) +fn background_task_manager() -> diagnostics::DiagnosticTaskManager { + diagnostics::DiagnosticTaskManager(AsyncTaskManager::new().routing(background_route)) } fn background_route(request: &ToolRequest) -> RoutingDecision { diff --git a/src/runtime/diagnostics.rs b/src/runtime/diagnostics.rs new file mode 100644 index 0000000..2921e67 --- /dev/null +++ b/src/runtime/diagnostics.rs @@ -0,0 +1,453 @@ +//! Carry producer identity through dependency-owned task and Runlet spawns. + +use std::sync::Arc; + +use agentkit_core::{TurnCancellation, TurnId}; +use agentkit_task_manager::{ + AsyncTaskManager, PendingLoopUpdates, TaskLaunchRequest, TaskManager, TaskManagerError, + TaskManagerHandle, TaskStartContext, TaskStartOutcome, TurnTaskUpdate, +}; +use agentkit_tools_core::{ + ApprovalRequest, OwnedToolContext, ToolCatalogEvent, ToolContext, ToolExecutionOutcome, + ToolExecutor, ToolRequest, ToolSpec, +}; +use async_trait::async_trait; + +use crate::events::DiagnosticScope; + +pub(super) struct DiagnosticTaskManager(pub(super) AsyncTaskManager); + +#[async_trait] +impl TaskManager for DiagnosticTaskManager { + async fn start_task( + &self, + request: TaskLaunchRequest, + mut ctx: TaskStartContext, + ) -> Result { + // This runs inside the originating turn, before AsyncTaskManager spawns. + // Never recover identity from session IDs: the same ID can be reactivated + // while an old producer is still running. + let origin = DiagnosticScope::capture(); + ctx.executor = Arc::new(DiagnosticExecutor { + inner: ctx.executor, + origin: origin.clone(), + }); + if let Some(scope) = &mut ctx.tool_context.execution_scope { + // Runlet crosses spawn_blocking and Handle::block_on before invoking + // this executor. Scoping only the outer compose future is insufficient. + scope.executor = Arc::new(DiagnosticExecutor { + inner: scope.executor.clone(), + origin, + }); + } + self.0.start_task(request, ctx).await + } + + async fn wait_for_turn( + &self, + turn_id: &TurnId, + cancellation: Option, + ) -> Result, TaskManagerError> { + self.0.wait_for_turn(turn_id, cancellation).await + } + + async fn take_pending_loop_updates(&self) -> Result { + self.0.take_pending_loop_updates().await + } + + async fn wait_for_loop_update(&self) -> Result<(), TaskManagerError> { + self.0.wait_for_loop_update().await + } + + async fn on_turn_interrupted(&self, turn_id: &TurnId) -> Result<(), TaskManagerError> { + self.0.on_turn_interrupted(turn_id).await + } + + fn handle(&self) -> TaskManagerHandle { + self.0.handle() + } +} + +struct DiagnosticExecutor { + inner: Arc, + origin: DiagnosticScope, +} + +#[async_trait] +impl ToolExecutor for DiagnosticExecutor { + fn specs(&self) -> Vec { + self.inner.specs() + } + + fn drain_catalog_events(&self) -> Vec { + self.inner.drain_catalog_events() + } + + async fn execute( + &self, + request: ToolRequest, + ctx: &mut ToolContext<'_>, + ) -> ToolExecutionOutcome { + self.origin.scope(self.inner.execute(request, ctx)).await + } + + async fn execute_owned( + &self, + request: ToolRequest, + ctx: OwnedToolContext, + ) -> ToolExecutionOutcome { + self.origin + .scope(self.inner.execute_owned(request, ctx)) + .await + } + + async fn execute_approved( + &self, + request: ToolRequest, + approval: &ApprovalRequest, + ctx: &mut ToolContext<'_>, + ) -> ToolExecutionOutcome { + self.origin + .scope(self.inner.execute_approved(request, approval, ctx)) + .await + } + + async fn execute_approved_owned( + &self, + request: ToolRequest, + approval: &ApprovalRequest, + ctx: OwnedToolContext, + ) -> ToolExecutionOutcome { + self.origin + .scope(self.inner.execute_approved_owned(request, approval, ctx)) + .await + } +} + +#[cfg(test)] +mod tests { + use std::{process::Command, time::Duration}; + + use agentkit_core::{MetadataMap, SessionId, ToolCallId, ToolOutput, ToolResultPart}; + use agentkit_task_manager::{TaskEvent, TaskLaunchKind}; + use agentkit_tool_compose::ComposeTool; + use agentkit_tools_core::{ + ApprovalReason, BasicToolExecutor, PermissionChecker, PermissionDecision, + PermissionRequest, Tool, ToolError, ToolExecutionScope, ToolName, ToolRegistry, ToolResult, + ToolSource, + }; + use serde_json::json; + use tokio::sync::Barrier; + + use super::*; + use crate::{events, tools::Observed}; + + // The barriers belong to a real tool, not to the production propagation path. + struct Gate { + spec: ToolSpec, + entered: Arc, + release: Arc, + } + + #[async_trait] + impl Tool for Gate { + fn spec(&self) -> &ToolSpec { + &self.spec + } + + fn proposed_requests( + &self, + request: &ToolRequest, + ) -> Result>, ToolError> { + Ok(vec![Box::new(GatePermission { + metadata: MetadataMap::new(), + needs_approval: request.input["approved"] == true, + })]) + } + + async fn invoke( + &self, + request: ToolRequest, + ctx: &mut ToolContext<'_>, + ) -> Result { + if request.input["approved"] == true { + assert_eq!(ctx.approved_request.as_ref().unwrap().id, approval().id); + } + if request.input["wait"] == true { + self.entered.wait().await; + self.release.wait().await; + } + Ok(ToolResult::new(ToolResultPart::success( + request.call_id, + ToolOutput::structured(json!({"done": true})), + ))) + } + } + + struct GatePermission { + metadata: MetadataMap, + needs_approval: bool, + } + + impl PermissionRequest for GatePermission { + fn kind(&self) -> &'static str { + "diagnostic.gate" + } + fn summary(&self) -> String { + "run gate".into() + } + fn metadata(&self) -> &MetadataMap { + &self.metadata + } + fn as_any(&self) -> &dyn std::any::Any { + self + } + } + + struct GatePermissions; + + impl PermissionChecker for GatePermissions { + fn evaluate(&self, request: &dyn PermissionRequest) -> PermissionDecision { + if request + .as_any() + .downcast_ref::() + .unwrap() + .needs_approval + { + PermissionDecision::RequireApproval(approval()) + } else { + PermissionDecision::Allow + } + } + } + + fn approval() -> ApprovalRequest { + ApprovalRequest::new( + "diagnostic-approval", + "diagnostic.gate", + ApprovalReason::PolicyRequiresConfirmation, + "run gate", + ) + } + + fn context(executor: Arc) -> OwnedToolContext { + let scope = ToolExecutionScope { + executor, + session_id: SessionId::new("A"), + turn_id: TurnId::new("turn"), + permissions: Arc::new(GatePermissions), + resources: Arc::new(()), + cancellation: None, + }; + scope.nested_context(MetadataMap::new()) + } + + fn request(call: &str, tool: &str, input: serde_json::Value) -> ToolRequest { + ToolRequest::new( + ToolCallId::new(call), + ToolName::new(tool), + input, + SessionId::new("A"), + TurnId::new("turn"), + ) + } + + async fn completed(manager: &DiagnosticTaskManager) { + let handle = manager.handle(); + loop { + match handle.next_event().await.expect("task event") { + TaskEvent::Completed(_, result) => { + assert!(!result.is_error, "{result:?}"); + return; + } + TaskEvent::Failed(_, error) => panic!("task failed: {error}"), + TaskEvent::Cancelled(_) => panic!("task cancelled"), + _ => {} + } + } + } + + #[test] + fn dependency_spawns_keep_origin_on_stderr() { + for mode in [ + "manager", + "approved-manager", + "runlet", + "approved-runlet", + "borrowed", + "approved-borrowed", + "owned", + "approved-owned", + "unscoped", + ] { + for switch in [false, true] { + let output = Command::new(std::env::current_exe().unwrap()) + .args([ + "--exact", + "runtime::diagnostics::tests::diagnostic_transport_child", + "--nocapture", + ]) + .env(events::EVENTS_ENV, "1") + .env("KIT_DIAGNOSTIC_RUNTIME_TEST", mode) + .env("KIT_DIAGNOSTIC_SWITCH_TEST", switch.to_string()) + .output() + .unwrap(); + assert!( + output.status.success(), + "{mode}/{switch}: {}\n{}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); + let diagnostics = String::from_utf8(output.stderr) + .unwrap() + .lines() + .filter_map(events::parse_diagnostic) + .collect::>(); + let markers = diagnostics + .iter() + .filter(|event| { + matches!(event.event, events::RuntimeEvent::SessionStarted { .. }) + }) + .collect::>(); + let first = markers.first().unwrap().activation.clone(); + let latest = markers.last().unwrap().activation.clone(); + assert_ne!(first, latest); + let expected = if mode == "unscoped" { None } else { first }; + let mut old_started = 0; + let mut old_finished = 0; + let mut current = 0; + for diagnostic in &diagnostics { + let (call, started) = match &diagnostic.event { + events::RuntimeEvent::ChildStarted { call, .. } => (call, true), + events::RuntimeEvent::ChildFinished { call, ok, .. } => { + assert!(*ok); + (call, false) + } + _ => continue, + }; + if call.starts_with("current") { + assert_eq!(diagnostic.activation, latest, "{mode}/{switch}"); + current += 1; + } else { + assert!(call.starts_with("old"), "{call}"); + assert_eq!(diagnostic.activation, expected, "{mode}/{switch}: {call}"); + if started { + old_started += 1; + } else { + old_finished += 1; + } + } + } + let calls = if mode.contains("runlet") { 2 } else { 1 }; + assert_eq!( + (old_started, old_finished, current), + (calls, calls, 2), + "{mode}/{switch}" + ); + // Delayed finishes and Runlet's dependent second start occur + // after the final activation marker, not just before it. + let last_marker = diagnostics + .iter() + .rposition(|event| { + matches!(event.event, events::RuntimeEvent::SessionStarted { .. }) + }) + .unwrap(); + assert!(diagnostics[last_marker + 1..].iter().any(|event| { + matches!(&event.event, events::RuntimeEvent::ChildFinished { call, .. } if call.starts_with("old")) + })); + if mode.contains("runlet") { + assert!(diagnostics[last_marker + 1..].iter().any(|event| { + matches!(&event.event, events::RuntimeEvent::ChildStarted { call, .. } if call.starts_with("old")) + })); + } + } + } + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn diagnostic_transport_child() { + let Ok(mode) = std::env::var("KIT_DIAGNOSTIC_RUNTIME_TEST") else { + return; + }; + tokio::time::timeout(Duration::from_secs(15), async { + let switch = std::env::var("KIT_DIAGNOSTIC_SWITCH_TEST").unwrap() == "true"; + let entered = Arc::new(Barrier::new(2)); + let release = Arc::new(Barrier::new(2)); + let mut children = ToolRegistry::new(); + children.register(Observed::new(Gate { + spec: ToolSpec::new("gate", "barrier gate", json!({"type": "object"})), + entered: entered.clone(), + release: release.clone(), + })); + let compose = ComposeTool::wrap(children.clone()) + .with_backend(crate::runtime::HiddenRunletBackend(children)); + let executor: Arc = Arc::new(BasicToolExecutor::new([ + Arc::new(compose) as Arc, + ])); + let manager = crate::runtime::background_task_manager(); + // The same real permission checker must interrupt an unapproved + // call; approved executor entry points must resume, not bypass it. + let denied = executor.execute_owned( + request("approval-check", "gate", json!({"approved": true})), + context(executor.clone()), + ).await; + assert!(matches!(denied, ToolExecutionOutcome::Interrupted(_))); + events::activate_diagnostics("A"); + let approved = mode.starts_with("approved-"); + let runlet = mode.contains("runlet"); + let req = if runlet { + request("old", "compose", json!({ + "script": "first = gate({wait: true})\nreturn gate({wait: false, previous: first})", + "background": true, + })) + } else { + request("old", "gate", json!({"wait": true, "approved": approved})) + }; + let start = async { + let ctx = context(executor.clone()); + if mode.contains("borrowed") || mode.ends_with("owned") { + let wrapped = DiagnosticExecutor { + inner: executor.clone(), + origin: DiagnosticScope::capture(), + }; + let mode = mode.clone(); + Some(tokio::spawn(async move { + let approval = approval(); + match mode.as_str() { + "borrowed" => wrapped.execute(req, &mut ctx.borrowed()).await, + "approved-borrowed" => wrapped.execute_approved(req, &approval, &mut ctx.borrowed()).await, + "owned" => wrapped.execute_owned(req, ctx).await, + "approved-owned" => wrapped.execute_approved_owned(req, &approval, ctx).await, + _ => unreachable!(), + } + })) + } else { + let outcome = manager.start_task(TaskLaunchRequest { + task_id: None, + request: req, + kind: if approved { TaskLaunchKind::Approved(approval()) } else { TaskLaunchKind::Plain }, + }, TaskStartContext { executor: executor.clone(), tool_context: ctx }).await.unwrap(); + assert!(matches!(outcome, TaskStartOutcome::Pending { .. })); + None + } + }; + let task = if mode == "unscoped" { start.await } else { events::scope_diagnostics("A", start).await }; + entered.wait().await; + if switch { events::activate_diagnostics("B"); } + events::activate_diagnostics("A"); + release.wait().await; + if let Some(task) = task { + assert!(matches!(task.await.unwrap(), ToolExecutionOutcome::Completed(_))); + } else { + completed(&manager).await; + } + // A newly-created producer still gets the current activation. + events::scope_diagnostics("A", manager.start_task(TaskLaunchRequest { + task_id: None, + request: request("current", "gate", json!({"wait": false})), + kind: TaskLaunchKind::Plain, + }, TaskStartContext { executor: executor.clone(), tool_context: context(executor) })).await.unwrap(); + completed(&manager).await; + }).await.expect("barrier-controlled task timed out"); + } +} diff --git a/src/session.rs b/src/session.rs index 62c58d5..e25ef34 100644 --- a/src/session.rs +++ b/src/session.rs @@ -23,6 +23,22 @@ use serde::{Deserialize, Serialize}; pub(crate) mod branch; +pub use branch::{ + Boundary as BranchBoundary, BranchMetadata as BranchProvenance, + CapturedSelection as BranchSelection, Completion as BranchCompletion, + SubmittedRequest as BranchRequest, +}; + +/// Validated checkout provenance, derived from the already-read transcript. +/// Invalid or incomplete checkouts remain visible, but must not supply edges. +#[derive(Clone, Debug, Default, PartialEq, Eq, PartialOrd, Ord)] +pub enum CatalogLineage { + #[default] + Root, + Branch(Box), + Warning(String), +} + pub const SCHEMA_VERSION: u32 = 3; const REDIRECT_SCHEMA_VERSION: u32 = 4; const PREVIOUS_SCHEMA_VERSION: u32 = 2; @@ -69,6 +85,11 @@ pub struct CatalogEntry { pub title: Option, pub preview: Option, pub is_subagent: bool, + pub lineage: CatalogLineage, + /// Ephemeral parent-context detail, verified against the original historical + /// prefix, or explicitly labeled provenance-only when it cannot be verified. + /// Absent for roots and incomplete/malformed checkouts. + pub branch_point: Option, /// Last activity as milliseconds since the Unix epoch. pub updated_at: u64, } @@ -1129,6 +1150,9 @@ fn catalog_for_workspace( } let ids = list_ids_for_workspace(&root, global_directory)?; let mut entries = Vec::with_capacity(ids.len()); + // Keep the already-read histories until every child's parent is available. + // Moving these vectors avoids extra transcript copies and all per-row rereads. + let mut histories = std::collections::HashMap::with_capacity(ids.len()); for id in ids { // Discovery is best-effort per transcript: a damaged file or one caught // mid-append must not hide every other session in the workspace. @@ -1159,14 +1183,28 @@ fn catalog_for_workspace( .unwrap_or(0); let directory = workspace_storage_directory(global_directory, &root); let title = read_display_name(&directory, &id).or(title); + let lineage = branch::catalog_lineage(&authority, &id); + histories.insert(id.clone(), authority.historical_items); entries.push(CatalogEntry { id, title, preview, is_subagent, + lineage, + branch_point: None, updated_at: item_updated.max(file_updated), }); } + for entry in &mut entries { + if let CatalogLineage::Branch(metadata) = &entry.lineage { + entry.branch_point = Some(branch::catalog_branch_point( + metadata, + histories + .get(&metadata.parent_session_id) + .map(Vec::as_slice), + )); + } + } entries.sort_by(|left, right| { right .updated_at @@ -1396,6 +1434,7 @@ pub(crate) fn list_ids_in(directory: &Path) -> Result, String> { struct Authority { items: Vec, historical_items: Vec>, + replacement_boundaries: Vec<(usize, usize)>, path: PathBuf, legacy_histories: Vec, } @@ -1574,6 +1613,7 @@ fn select_authority_with( Ok(authority.map(|candidate| Authority { items: candidate.history.items, historical_items: candidate.history.states, + replacement_boundaries: candidate.history.replacement_boundaries, path: candidate.path, legacy_histories, })) diff --git a/src/session/branch.rs b/src/session/branch.rs index ff4cd24..a93f5d3 100644 --- a/src/session/branch.rs +++ b/src/session/branch.rs @@ -10,9 +10,9 @@ use crate::{ReasoningEffort, provider::ModelSelection}; pub(crate) const METADATA_KEY: &str = "dev.kit.session.prompt_checkout"; const VERSION: u32 = 1; -#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] #[serde(deny_unknown_fields)] -pub(crate) struct Boundary { +pub struct Boundary { /// Index in `load_history`, including states preceding compaction. pub state_index: usize, /// Number of items in the retained prefix (exclusive for a selected user @@ -36,9 +36,9 @@ impl Boundary { } /// Canonical IDs from the actual runtime types, not UI labels or adapter defaults. -#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] #[serde(deny_unknown_fields)] -pub(crate) struct CapturedSelection { +pub struct CapturedSelection { pub provider: String, pub model: String, /// `default` is explicit: absence is not silently interpreted as default. @@ -62,16 +62,16 @@ impl CapturedSelection { } } -#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] #[serde(deny_unknown_fields)] -pub(crate) struct SubmittedRequest { +pub struct SubmittedRequest { pub id: String, pub selection: CapturedSelection, } -#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] #[serde(deny_unknown_fields)] -pub(crate) struct Completion { +pub struct Completion { pub prefix_len: usize, /// Hash of every retained item, with only this payload removed. pub prefix_hash: String, @@ -79,9 +79,9 @@ pub(crate) struct Completion { pub prompt_hash: String, } -#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] #[serde(deny_unknown_fields)] -pub(crate) struct BranchMetadata { +pub struct BranchMetadata { pub version: u32, pub parent_session_id: String, pub boundary: Boundary, @@ -169,6 +169,114 @@ impl BranchMetadata { } } +/// Classify a catalog row without opening, repairing, flushing, or rereading it. +/// Exact replacement boundaries (not subsequently appended state) are the same +/// completion evidence used by checkout recovery. This is provenance validation, +/// not a new durability barrier for process-resident filesystem overlays. +pub(super) fn catalog_lineage(authority: &Authority, session_id: &str) -> CatalogLineage { + match catalog_provenance(authority, session_id) { + Ok(Some(metadata)) => CatalogLineage::Branch(Box::new(metadata)), + Ok(None) => CatalogLineage::Root, + Err(error) => { + CatalogLineage::Warning(truncate_catalog_text(&normalize_catalog_text(&error), 160)) + } + } +} + +fn catalog_provenance( + authority: &Authority, + session_id: &str, +) -> Result, String> { + let mut provenance = BranchMetadata::read(&authority.items)?; + for state in &authority.historical_items { + if let Some(metadata) = BranchMetadata::read(state)? { + if provenance + .as_ref() + .is_some_and(|previous| previous != &metadata) + { + return Err("prompt checkout metadata changed across history".into()); + } + provenance = Some(metadata); + } + } + // Compaction may remove the bootstrap payload; a validated historical + // completion still proves the branch's original lineage. + let Some(metadata) = provenance else { + return Ok(None); + }; + if metadata.parent_session_id == session_id { + return Err("prompt checkout cannot parent itself".into()); + } + let mut completed: Option<&[Item]> = None; + for &(state_index, length) in &authority.replacement_boundaries { + let snapshot = &authority.historical_items[state_index][..length]; + if BranchMetadata::read(snapshot)?.as_ref() == Some(&metadata) + && metadata.matches_snapshot(snapshot)? + { + if completed.is_some_and(|previous| previous != snapshot) { + return Err("conflicting prompt checkout completion records".into()); + } + completed = Some(snapshot); + } + } + if completed.is_none() { + return Err("incomplete prompt checkout: no matching completion record".into()); + } + Ok(Some(metadata)) +} + +/// Describe the actual retained parent context only after verifying the original +/// prefix with the same canonical digest as `Boundary::new`. In particular, do +/// not hash the sanitized child prefix or substitute the parent's current state. +/// The catalog owns this history snapshot; this helper performs no filesystem IO. +pub(super) fn catalog_branch_point( + metadata: &BranchMetadata, + parent_history: Option<&[Vec]>, +) -> String { + let Some(history) = parent_history else { + return "Provenance only: parent history unavailable".into(); + }; + let Some(prefix) = history + .get(metadata.boundary.state_index) + .and_then(|state| state.get(..metadata.boundary.prefix_len)) + .filter(|prefix| !prefix.is_empty()) + else { + return "Provenance only: recorded parent prefix unavailable".into(); + }; + match digest(prefix) { + Ok(hash) if hash == metadata.boundary.prefix_hash => {} + Ok(_) => return "Provenance only: recorded parent prefix hash does not match".into(), + Err(_) => return "Provenance only: parent prefix could not be verified".into(), + } + let retained = prefix.last().expect("nonempty verified parent prefix"); + let role = match retained.kind { + ItemKind::System => "System", + ItemKind::Developer => "Developer", + ItemKind::Context => "Context", + ItemKind::Notification => "Notification", + ItemKind::User => "User", + ItemKind::Assistant => "Assistant", + ItemKind::Tool => "Tool", + }; + let preview = retained + .parts + .iter() + .find_map(|part| { + let text = match part { + Part::Text(text) => Some(text.text.as_str()), + Part::ToolResult(result) => match &result.output { + agentkit_core::ToolOutput::Text(text) => Some(text.as_str()), + _ => None, + }, + _ => None, + }?; + let text = normalize_catalog_text(text); + (!text.is_empty()).then(|| truncate_catalog_text(&text, 160)) + }) + .unwrap_or_else(|| "[no visible text in retained item]".into()); + format!("Validated parent context · {role}: {preview}") +} + /// Validate a branch reload before an opener can repair history or an adapter /// can start. Root/legacy transcripts retain their ordinary defaults. pub(crate) fn validate_resume( @@ -680,6 +788,341 @@ mod tests { assert!(metadata.matches_snapshot(&descendant).unwrap()); } + fn catalog_row(root: &Path) -> CatalogEntry { + let mut entries = catalog_for_workspace(root, &root.join("sessions")).unwrap(); + assert_eq!(entries.len(), 1); + entries.remove(0) + } + + fn disk_snapshot(root: &Path) -> Vec<(PathBuf, SystemTime, Option>)> { + let mut pending = vec![root.to_path_buf()]; + let mut snapshot = Vec::new(); + while let Some(path) = pending.pop() { + let metadata = std::fs::metadata(&path).unwrap(); + let bytes = if metadata.is_dir() { + pending.extend( + std::fs::read_dir(&path) + .unwrap() + .map(|entry| entry.unwrap().path()), + ); + None + } else { + Some(std::fs::read(&path).unwrap()) + }; + snapshot.push((path, metadata.modified().unwrap(), bytes)); + } + snapshot.sort_by(|left, right| left.0.cmp(&right.0)); + snapshot + } + + #[test] + fn catalog_lineage_validates_completion_survives_compaction_and_never_writes() { + let root = tempfile::tempdir().unwrap(); + let transcript = prepared(); + let expected = BranchMetadata::read(&transcript).unwrap().unwrap(); + let opened = open_branch(root.path(), transcript.clone()); + let before = disk_snapshot(root.path()); + assert!(matches!( + catalog_row(root.path()).lineage, + CatalogLineage::Warning(_) + )); + assert_eq!(disk_snapshot(root.path()), before); + commit(&opened.observer, &transcript).unwrap(); + let mut later = Item::text(ItemKind::Assistant, "later answer"); + stamp_item(&mut later, Timestamp::now()); + opened.observer.0.lock().unwrap().append(&later).unwrap(); + let before = disk_snapshot(root.path()); + for _ in 0..2 { + assert_eq!( + catalog_row(root.path()).lineage, + CatalogLineage::Branch(Box::new(expected.clone())) + ); + } + assert_eq!(disk_snapshot(root.path()), before); + opened + .observer + .replace(&[Item::text(ItemKind::System, "compacted")]) + .unwrap(); + drop(opened); + let before = disk_snapshot(root.path()); + assert_eq!( + catalog_row(root.path()).lineage, + CatalogLineage::Branch(Box::new(expected.clone())) + ); + assert_eq!(disk_snapshot(root.path()), before); + // Provenance extraction consumes the authority's existing parsed records. + // It must not reopen the transcript or introduce a disk/flush barrier. + let authority = select_authority_with( + &root.path().join("sessions"), + &canonical_workspace(root.path()), + "branch", + false, + false, + ) + .unwrap() + .unwrap(); + std::fs::remove_file(path(root.path())).unwrap(); + assert_eq!( + catalog_lineage(&authority, "branch"), + CatalogLineage::Branch(Box::new(expected)) + ); + } + + #[test] + fn catalog_lineage_keeps_malformed_unknown_and_conflicting_metadata_visible() { + for value in [ + serde_json::Value::Null, + json!({"version": 999}), + json!({"version": 1}), + ] { + let root = tempfile::tempdir().unwrap(); + let mut transcript = prepared(); + transcript[0].metadata.insert(METADATA_KEY.into(), value); + let opened = open_branch(root.path(), transcript); + opened.observer.commit_creation(); + drop(opened); + let before = disk_snapshot(root.path()); + assert!(matches!( + catalog_row(root.path()).lineage, + CatalogLineage::Warning(_) + )); + assert_eq!(disk_snapshot(root.path()), before); + } + let root = tempfile::tempdir().unwrap(); + let mut transcript = prepared(); + let opened = open_branch(root.path(), transcript.clone()); + commit(&opened.observer, &transcript).unwrap(); + transcript[0].metadata.get_mut(METADATA_KEY).unwrap()["parent_session_id"] = + json!("other-parent"); + opened.observer.replace(&transcript).unwrap(); + drop(opened); + assert!(matches!( + catalog_row(root.path()).lineage, + CatalogLineage::Warning(_) + )); + } + + #[test] + fn catalog_lineage_rejects_append_only_and_short_replacement_completion() { + for short_replacement in [false, true] { + let root = tempfile::tempdir().unwrap(); + let transcript = prepared(); + let opened = open_branch(root.path(), transcript[..1].to_vec()); + if short_replacement { + opened.observer.replace(&transcript[..1]).unwrap(); + } + opened + .observer + .0 + .lock() + .unwrap() + .append(&transcript[1]) + .unwrap(); + opened.observer.commit_creation(); + drop(opened); + let before = disk_snapshot(root.path()); + assert!(matches!( + catalog_row(root.path()).lineage, + CatalogLineage::Warning(_) + )); + assert_eq!(disk_snapshot(root.path()), before); + } + } + + #[test] + fn catalog_lineage_rejects_mismatched_completion_and_self_parenting() { + for self_parent in [false, true] { + let root = tempfile::tempdir().unwrap(); + let mut transcript = prepared(); + if self_parent { + transcript[0].metadata.get_mut(METADATA_KEY).unwrap()["parent_session_id"] = + json!("branch"); + } else { + transcript[1].parts = Item::text(ItemKind::User, "changed prompt").parts; + } + let opened = open_branch(root.path(), transcript.clone()); + opened.observer.replace(&transcript).unwrap(); + opened.observer.commit_creation(); + drop(opened); + let before = disk_snapshot(root.path()); + assert!(matches!( + catalog_row(root.path()).lineage, + CatalogLineage::Warning(_) + )); + assert_eq!(disk_snapshot(root.path()), before); + } + } + + #[test] + fn catalog_lineage_preserves_legacy_roots_without_rewriting() { + let root = tempfile::tempdir().unwrap(); + let opened = open_branch(root.path(), vec![Item::text(ItemKind::System, "legacy")]); + opened.observer.commit_creation(); + drop(opened); + let before = disk_snapshot(root.path()); + assert_eq!(catalog_row(root.path()).lineage, CatalogLineage::Root); + assert_eq!(disk_snapshot(root.path()), before); + } + + fn catalog_child(root: &Path) -> CatalogEntry { + catalog_for_workspace(root, &root.join("sessions")) + .unwrap() + .into_iter() + .find(|entry| entry.id == "branch") + .unwrap() + } + + fn prepared_from_parent(prefix: Vec, state_index: usize) -> Vec { + let boundary = Boundary::new(state_index, &prefix).unwrap(); + let request = BranchMetadata::read(&prepared()).unwrap().unwrap().request; + // Deliberately model a sanitized child copy: only the original parent + // history, never this changed child text, can provide the context detail. + let mut sanitized = prefix; + sanitized.last_mut().unwrap().parts = + Item::text(ItemKind::Assistant, "sanitized child copy").parts; + prepare( + sanitized, + "parent".into(), + boundary, + "checkout-1".into(), + request, + Item::text(ItemKind::User, "new submitted branch prompt"), + ) + .unwrap() + } + + #[test] + fn catalog_branch_point_uses_original_historical_prefix_after_parent_advances_and_compacts() { + let root = tempfile::tempdir().unwrap(); + let directory = root.path().join("sessions"); + let parent = open_in( + root.path(), + &directory, + "parent", + false, + false, + vec![Item::text(ItemKind::System, "initial bootstrap")], + ) + .unwrap(); + let retained = vec![ + Item::text(ItemKind::System, "compacted bootstrap"), + Item::text(ItemKind::User, "original parent prompt"), + Item::text(ItemKind::Assistant, "retained parent answer"), + ]; + parent.observer.replace(&retained).unwrap(); + let transcript = prepared_from_parent(retained, 1); + let child = open_branch(root.path(), transcript.clone()); + let before = disk_snapshot(root.path()); + let row = catalog_child(root.path()); + assert!(matches!(row.lineage, CatalogLineage::Warning(_))); + assert_eq!(row.branch_point, None); + assert_eq!(disk_snapshot(root.path()), before); + commit(&child.observer, &transcript).unwrap(); + let expected = Some("Validated parent context · Assistant: retained parent answer".into()); + assert_eq!(catalog_child(root.path()).branch_point, expected); + let mut later = Item::text( + ItemKind::User, + "later parent prompt must not become the branch point", + ); + stamp_item(&mut later, Timestamp::now()); + parent.observer.0.lock().unwrap().append(&later).unwrap(); + assert_eq!(catalog_child(root.path()).branch_point, expected); + parent + .observer + .replace(&[ + Item::text(ItemKind::System, "new compaction"), + Item::text( + ItemKind::Assistant, + "newest answer must not become the branch point", + ), + ]) + .unwrap(); + child + .observer + .replace(&[Item::text(ItemKind::System, "child compaction")]) + .unwrap(); + drop(parent); + drop(child); + let before = disk_snapshot(root.path()); + assert_eq!(catalog_child(root.path()).branch_point, expected); + assert_eq!(disk_snapshot(root.path()), before); + } + + #[test] + fn catalog_branch_point_labels_missing_and_mismatched_parent_prefix_as_provenance_only() { + for mode in [ + "missing parent", + "missing state", + "short prefix", + "hash mismatch", + ] { + let root = tempfile::tempdir().unwrap(); + let directory = root.path().join("sessions"); + let parent = open_in( + root.path(), + &directory, + "parent", + false, + false, + vec![ + Item::text(ItemKind::System, "bootstrap"), + Item::text(ItemKind::User, "parent prompt"), + Item::text(ItemKind::Assistant, "retained answer"), + ], + ) + .unwrap(); + let mut transcript = prepared_from_parent(parent.transcript.clone(), 0); + if mode == "missing state" { + transcript[0].metadata.get_mut(METADATA_KEY).unwrap()["boundary"]["state_index"] = + json!(99); + } else if mode == "hash mismatch" { + transcript[0].metadata.get_mut(METADATA_KEY).unwrap()["boundary"]["prefix_hash"] = + json!("0".repeat(64)); + } + let child = open_branch(root.path(), transcript.clone()); + commit(&child.observer, &transcript).unwrap(); + drop(parent); + drop(child); + let parent_path = transcript_path( + &workspace_storage_directory(&directory, &canonical_workspace(root.path())), + "parent", + ); + if mode == "missing parent" { + std::fs::remove_file(&parent_path).unwrap(); + } else if mode == "short prefix" { + let bytes = std::fs::read(&parent_path).unwrap(); + let end = bytes.iter().position(|byte| *byte == b'\n').unwrap() + 1; + std::fs::write(&parent_path, &bytes[..end]).unwrap(); + } + let before = disk_snapshot(root.path()); + let row = catalog_child(root.path()); + assert!(matches!(row.lineage, CatalogLineage::Branch(_)), "{mode}"); + let detail = row.branch_point.unwrap(); + assert!(detail.starts_with("Provenance only:"), "{mode}: {detail}"); + assert!(!detail.contains("retained answer"), "{mode}"); + assert_eq!(disk_snapshot(root.path()), before, "{mode}"); + } + } + + #[test] + fn catalog_branch_point_is_bounded_sanitized_and_has_no_filesystem_dependency() { + let mut prefix = vec![Item::text(ItemKind::System, "bootstrap")]; + prefix.push(Item::text( + ItemKind::Assistant, + format!("safe\u{1b}\u{202e} {}", "界".repeat(500)), + )); + let transcript = prepared_from_parent(prefix.clone(), 0); + let metadata = BranchMetadata::read(&transcript).unwrap().unwrap(); + let detail = catalog_branch_point(&metadata, Some(&[prefix])); + assert!(detail.starts_with("Validated parent context · Assistant: safe ")); + assert!(detail.chars().count() <= 210); + assert!(!detail.contains(['\u{1b}', '\u{202e}'])); + assert_eq!( + catalog_branch_point(&metadata, None), + "Provenance only: parent history unavailable" + ); + } + #[test] fn writer_uses_existing_replacement_shape_and_restart_lookup_survives_compaction() { let root = tempfile::tempdir().unwrap(); diff --git a/src/tools/subagent.rs b/src/tools/subagent.rs index 8e989e0..3e44f7b 100644 --- a/src/tools/subagent.rs +++ b/src/tools/subagent.rs @@ -610,7 +610,7 @@ impl Subagents { let (reply, response) = oneshot::channel(); let manager = self.clone(); tokio::spawn( - async move { + crate::events::inherit_diagnostics(async move { let source_state = Arc::clone(&operation.source_state); let reservation = operation.id.clone(); let result = manager.run_fork(operation, &reply).await; @@ -621,7 +621,7 @@ impl Subagents { let _ = reply.send(Err(error)); } } - } + }) .instrument(tracing::Span::current()), ); match response.await.map_err(|_| { @@ -1092,11 +1092,11 @@ impl Subagents { ) -> ChildError { let manager = self.clone(); match tokio::spawn( - async move { + crate::events::inherit_diagnostics(async move { manager .cleanup_installed_child(&id, &state, &child, error) .await - } + }) .instrument(tracing::Span::current()), ) .await @@ -1140,7 +1140,7 @@ impl Subagents { let parent_id = self.config.parent_id.clone(); let parent_name = self.config.parent_name.clone(); let mut closed = child.closed_signal(); - tokio::spawn(async move { + tokio::spawn(crate::events::inherit_diagnostics(async move { if !*closed.borrow() { let _ = closed.changed().await; } @@ -1179,7 +1179,7 @@ impl Subagents { *event_parent_name = parent_name; } let _ = event_sink(&event); - }); + })); } async fn fail_removed_and_remove(&self, id: &str, state: &Arc>) { diff --git a/src/tools/subagent/tests.rs b/src/tools/subagent/tests.rs index d3a62a8..e94cb6f 100644 --- a/src/tools/subagent/tests.rs +++ b/src/tools/subagent/tests.rs @@ -478,6 +478,217 @@ async fn dropping_a_session_manager_terminates_its_children() { .unwrap(); } +#[tokio::test] +async fn child_exit_monitor_keeps_spawn_activation_across_reactivation() { + for route in ["same", "roundtrip", "unscoped"] { + let directory = tempfile::tempdir().unwrap(); + let (mut manager, state, _) = manager_with_disconnected_session(directory.path()); + let (events, mut received) = tokio::sync::mpsc::unbounded_channel(); + manager.event_sink = Arc::new(move |event| { + events + .send((event.clone(), events::DiagnosticScope::capture())) + .map_err(|_| ()) + }); + let child = ChildSession::disconnected_for_test(); + // The real exit monitor must acquire the state lock before publishing. + // Hold it across the route change instead of adding a test-only pause. + let locked = state.lock().await; + let session = format!("monitor-{route}"); + events::activate_diagnostics(&session); + let expected = if route == "unscoped" { + manager.monitor_child_exit("source".into(), &state, &child); + events::DiagnosticScope::default() + } else { + events::scope_diagnostics(&session, async { + let expected = events::DiagnosticScope::capture(); + manager.monitor_child_exit("source".into(), &state, &child); + expected + }) + .await + }; + if route == "roundtrip" { + events::activate_diagnostics("monitor-other"); + } + events::activate_diagnostics(&session); + drop(locked); + let (event, actual) = + tokio::time::timeout(std::time::Duration::from_secs(2), received.recv()) + .await + .expect("child-exit monitor did not publish") + .unwrap(); + assert!(matches!( + event, + events::RuntimeEvent::SubagentStateChanged { + status: SubagentStatus::Removed, + .. + } + )); + assert_eq!(actual, expected, "route: {route}"); + } +} + +#[test] +fn reused_and_native_forked_children_emit_current_roster_activation() { + const PROBE: &str = "KIT_TEST_CURRENT_SUBAGENT_DIAGNOSTICS"; + if std::env::var_os(PROBE).is_some() { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap() + .block_on(async { + for route in ["same", "roundtrip"] { + let root = tempfile::tempdir().unwrap(); + let mut config = + manager_with_generic_harness(root.path(), Vec::new()).child_config(); + // Use the existing mock transport as a configured Kit profile + // so native forks carry Kit's descendant-parent metadata. + config.harnesses = + crate::acp_child::AcpHarnesses::new(std::collections::BTreeMap::from([( + "kit".into(), + crate::acp_child::AcpHarnessProfile { + command: "python3".into(), + args: vec![format!( + "{}/fixtures/mock-acp.py", + env!("CARGO_MANIFEST_DIR") + )], + permissions: Default::default(), + }, + )])) + .unwrap(); + config.default_harness = "acp.kit".into(); + let manager = Subagents::new(config, 2); + let session = format!("reuse-{route}"); + events::activate_diagnostics(&session); + let (source, branch) = events::scope_diagnostics(&session, async { + let source = manager + .create( + format!("initial-source-{route}"), + CreateOptions::default(), + 0, + TurnCancellation::default(), + None, + ) + .await + .unwrap(); + let branch = manager + .fork( + source.clone(), + format!("initial-fork-{route}"), + None, + 0, + TurnCancellation::default(), + None, + ) + .await + .unwrap(); + (source, branch) + }) + .await; + if route == "roundtrip" { + events::activate_diagnostics("reuse-other"); + } + events::activate_diagnostics(&session); + events::scope_diagnostics(&session, async { + // These are actual ACP prompts on logical sessions sharing + // the old child process, not synthetic roster events. + let source = manager + .prompt( + source, + format!("current-source-{route}"), + TurnCancellation::default(), + None, + ) + .await + .unwrap(); + let branch = manager + .prompt( + branch, + format!("current-fork-{route}"), + TurnCancellation::default(), + None, + ) + .await + .unwrap(); + let new_branch = manager + .fork( + source.clone(), + format!("current-new-fork-{route}"), + None, + 0, + TurnCancellation::default(), + None, + ) + .await + .unwrap(); + for child in [branch, new_branch, source] { + manager + .close(&child.id, &TurnCancellation::default()) + .await + .unwrap(); + } + }) + .await; + } + }); + return; + } + let output = std::process::Command::new(std::env::current_exe().unwrap()) + .args(["--exact", "tools::subagent::tests::reused_and_native_forked_children_emit_current_roster_activation", "--nocapture"]) + .env(PROBE, "1").env(events::EVENTS_ENV, "1").output().unwrap(); + assert!( + output.status.success(), + "{}", + String::from_utf8_lossy(&output.stderr) + ); + let stderr = String::from_utf8(output.stderr).unwrap(); + let lines = stderr + .lines() + .filter_map(events::parse_diagnostic) + .collect::>(); + for route in ["same", "roundtrip"] { + let session = format!("reuse-{route}"); + let activations = lines + .iter() + .filter_map(|line| match &line.event { + events::RuntimeEvent::SessionStarted { session_id } if session_id == &session => { + line.activation.clone() + } + _ => None, + }) + .collect::>(); + assert_eq!(activations.len(), 2); + assert_ne!(activations[0], activations[1]); + let mut current_ids = HashSet::new(); + for line in &lines { + if let events::RuntimeEvent::SubagentStateChanged { id, task, .. } = &line.event { + if !task.ends_with(route) { + continue; + } + let expected = if task.starts_with("initial-") { + &activations[0] + } else { + assert!(task.starts_with("current-")); + current_ids.insert(id.clone()); + &activations[1] + }; + assert_eq!(line.activation.as_ref(), Some(expected), "task: {task}"); + } + } + assert_eq!( + current_ids.len(), + 3, + "missing reused source/fork or current native fork" + ); + for id in current_ids { + // Explicit closes of both old and new handles use the current + // invocation, unlike lifetime-owned shared-stream EOF cleanup. + assert!(lines.iter().any(|line| matches!(&line.event, + events::RuntimeEvent::SubagentDescendantsRemoved { ancestor_id } if ancestor_id == &id) + && line.activation.as_ref() == Some(&activations[1])), "missing current close for {id}"); + } + } +} + #[test] fn close_input_accepts_a_handle_subagent_id_or_background_call_id() { let handle = json!({"id": "child", "output": "done", "generation": 1}); diff --git a/src/tui/app.rs b/src/tui/app.rs index 6f922d4..82d164a 100644 --- a/src/tui/app.rs +++ b/src/tui/app.rs @@ -45,7 +45,10 @@ pub enum Update { /// The actual dynamically allocated A2A listen address. A2aAddress(String), /// Result of listing sessions without blocking the terminal event loop. - SessionCatalog(Result, String>), + SessionCatalog { + epoch: u64, + result: Result, String>, + }, FileMatches { revision: u64, result: Result, String>, @@ -208,6 +211,8 @@ pub struct EffortDialog { } pub struct SessionDialog { + pub query: String, + pub searching: bool, pub selected: usize, pub rename: Option, } @@ -374,7 +379,9 @@ pub enum Action { id: String, }, New(Option), - ListSessions, + ListSessions { + epoch: u64, + }, RenameSession { session_id: String, display_name: Option, @@ -711,10 +718,12 @@ pub struct App { pub reasoning_effort: String, pub effort_choices: Vec, pub effort_dialog: Option, - pub session_choices: Vec, + pub(super) session_forest: super::branches::BranchForest, + pub(super) session_matches: Vec, pub session_dialog: Option, pub file_picker: Option, - session_catalog_pending: bool, + pub(super) session_catalog_pending: bool, + pub(super) session_catalog_epoch: u64, pub auth_methods: Vec, pub available_commands: Vec, pub command_completion_selected: usize, @@ -722,6 +731,7 @@ pub struct App { command_completion_dismissed: Option, pub a2a: String, pub session_id: Option, + pub(super) session_connected: bool, /// Session currently associated with the ordered runtime side channel. runtime_session_id: Option, pub blocks: Vec, @@ -964,10 +974,12 @@ impl App { reasoning_effort: "default".into(), effort_choices: Vec::new(), effort_dialog: None, - session_choices: Vec::new(), + session_forest: super::branches::BranchForest::default(), + session_matches: Vec::new(), session_dialog: None, file_picker: None, session_catalog_pending: false, + session_catalog_epoch: 0, branch_epoch: 0, branch_chooser: None, branch_draft: None, @@ -979,6 +991,7 @@ impl App { command_completion_dismissed: None, a2a, session_id: None, + session_connected: true, runtime_session_id: None, blocks: Vec::new(), transcript_cache: Vec::new(), @@ -1705,7 +1718,12 @@ impl App { pub fn apply(&mut self, update: Update) { // The source stays loaded while its provisional replacement is visible. - if let Some(draft) = &mut self.branch_draft { + if let Some(draft) = &mut self.branch_draft + && !matches!( + &update, + Update::SessionCatalog { .. } | Update::SessionRenamed { .. } + ) + { match update { Update::RoutedRuntime { session_id, event } => { if draft.original.session_id.as_ref() == Some(&session_id) { @@ -1765,19 +1783,18 @@ impl App { if self.session_id.as_ref() == Some(session_id)); match update { Update::A2aAddress(address) => self.a2a = address, - Update::SessionCatalog(result) => { + Update::SessionCatalog { epoch, result } => { + if epoch != self.session_catalog_epoch || !self.session_catalog_pending { + return; + } self.session_catalog_pending = false; match result { - Ok(entries) if entries.is_empty() => { - self.toast("no sessions found for this workspace"); - } Ok(entries) => { - self.session_choices = entries; - self.file_picker = None; - self.session_dialog = Some(SessionDialog { - selected: 0, - rename: None, - }); + let empty = entries.is_empty(); + self.set_session_choices(entries); + if empty { + self.toast("no sessions found for this workspace"); + } } Err(error) => self.toast(format!("could not list sessions: {error}")), } @@ -1786,29 +1803,48 @@ impl App { session_id, display_name, result, - } => match result { - Ok(title) => { - if let Some(dialog) = &mut self.session_dialog { - dialog.rename = None; - } - if let Some(entry) = self - .session_choices - .iter_mut() - .find(|entry| entry.id == session_id) - { - entry.title = title; + } => { + let selected_id = self + .session_dialog + .as_ref() + .and_then(|dialog| self.session_matches.get(dialog.selected)) + .map(|matched| self.session_forest.rows[matched.index].entry.id.clone()); + let saving_selected = selected_id.as_ref() == Some(&session_id) + && self + .session_dialog + .as_ref() + .is_some_and(|dialog| matches!(dialog.rename, Some(SessionRename::Saving))); + match result { + Ok(title) => { + if saving_selected && let Some(dialog) = &mut self.session_dialog { + dialog.rename = None; + } + let entries = self + .session_forest + .rows + .iter() + .map(|row| { + let mut entry = row.entry.clone(); + if entry.id == session_id { + entry.title = title.clone(); + } + entry + }) + .collect(); + self.set_session_choices(entries); + self.filter_sessions(selected_id.as_deref()); } - } - Err(error) => { - if let Some(dialog) = &mut self.session_dialog { - dialog.rename = Some(match display_name { - Some(name) => SessionRename::Editing(name), - None => SessionRename::ConfirmClear, - }); + Err(error) => { + if saving_selected && let Some(dialog) = &mut self.session_dialog { + dialog.rename = Some(match display_name { + Some(name) => SessionRename::Editing(name), + None => SessionRename::ConfirmClear, + }); + } + self.toast(format!("could not rename session: {error}")); } - self.toast(format!("could not rename session: {error}")); } - }, + } Update::FileMatches { revision, result } => { if self .file_picker @@ -2165,8 +2201,10 @@ impl App { self.queue_handoff = false; self.retired_steers.clear(); self.steer_mutations.clear(); - self.session_catalog_pending = false; + self.dismiss_sessions(); + let same_session = self.session_id.as_ref() == Some(&session_id); self.session_id = Some(session_id); + self.session_connected = true; self.file_picker = None; self.available_commands.clear(); self.command_completion_selected = 0; @@ -2184,7 +2222,9 @@ impl App { self.transcript_cache_width = 0; self.retained_image_source_bytes = 0; self.transcript_focus_index = None; - self.clear_attachments(); + if !same_session { + self.clear_attachments(); + } self.latest_agent_source.clear(); self.phase = Phase::Idle; self.turn_started = None; @@ -2465,29 +2505,10 @@ impl App { /// A paste never sends: the newlines in it are part of the text. Multi-line /// pastes say so, because the prompt box shows only its last rows and the /// rest is easy to miss. - pub fn session_rename_active(&self) -> bool { - self.session_dialog - .as_ref() - .is_some_and(|dialog| dialog.rename.is_some()) - } - pub fn paste(&mut self, text: &str) { if self.model_switch.is_some() { return; } - if self.branch_chooser.is_some() || self.branch_submitting() { - return; - } - if self.branch_draft.is_some() { - self.last_key = None; - self.editor.insert_str(text); - return; - } - if let Some(dialog) = &mut self.navigation.dialog { - dialog.insert(text); - self.sync_navigation(); - return; - } // An explicit bracketed paste is not part of the unbracketed key-burst heuristic. self.last_key = None; if let Some(rename) = self @@ -2507,6 +2528,24 @@ impl App { } return; } + if self.session_dialog.is_some() { + self.session_dialog.as_mut().unwrap().searching = true; + self.insert_session_query(text); + return; + } + if self.branch_chooser.is_some() || self.branch_submitting() { + return; + } + if self.branch_draft.is_some() { + self.last_key = None; + self.editor.insert_str(text); + return; + } + if let Some(dialog) = &mut self.navigation.dialog { + dialog.insert(text); + self.sync_navigation(); + return; + } self.file_picker = None; self.queue_handoff = false; self.editor.insert_str(text); @@ -2652,10 +2691,135 @@ impl App { Action::None } + pub(super) fn set_session_choices(&mut self, entries: Vec) { + self.session_forest = super::branches::BranchForest::new(entries); + self.filter_sessions(None); + } + + fn filter_sessions(&mut self, preferred_id: Option<&str>) { + let query = self + .session_dialog + .as_ref() + .map_or("", |dialog| dialog.query.as_str()); + self.session_matches = self.session_forest.search(query); + if let Some(dialog) = &mut self.session_dialog { + dialog.selected = preferred_id + .and_then(|id| { + self.session_matches + .iter() + .position(|matched| self.session_forest.rows[matched.index].entry.id == id) + }) + .unwrap_or_else(|| { + dialog + .selected + .min(self.session_matches.len().saturating_sub(1)) + }); + } + } + + fn insert_session_query(&mut self, text: &str) { + if let Some(dialog) = &mut self.session_dialog { + for grapheme in text.graphemes(true) { + if grapheme.chars().any(char::is_control) { + continue; + } + if dialog.query.len() + grapheme.len() > 4096 { + break; + } + dialog.query.push_str(grapheme); + } + dialog.selected = 0; + } + self.filter_sessions(None); + } + + fn open_sessions(&mut self) -> Action { + if self.session_catalog_pending { + self.toast("session catalog scan is already in progress — Esc cancels"); + return Action::None; + } + self.session_catalog_epoch = self.session_catalog_epoch.wrapping_add(1); + self.session_catalog_pending = true; + self.file_picker = None; + self.session_forest = super::branches::BranchForest::default(); + self.session_matches.clear(); + self.session_dialog = Some(SessionDialog { + query: String::new(), + searching: false, + selected: 0, + rename: None, + }); + self.toast("loading conversation branches — Esc cancels"); + Action::ListSessions { + epoch: self.session_catalog_epoch, + } + } + + fn dismiss_sessions(&mut self) { + self.session_dialog = None; + self.session_catalog_pending = false; + self.session_catalog_epoch = self.session_catalog_epoch.wrapping_add(1); + } + + pub(super) fn require_session_connection(&mut self) -> bool { + if !self.session_connected { + self.toast("session disconnected — F4, select current session, Enter to retry; draft preserved"); + } + self.session_connected + } + + pub(super) fn session_switch_allowed(&mut self, requested_id: &str) -> bool { + if self.phase != Phase::Idle + || self.editing_branch() + || self.branch_chooser.is_some() + || self.editing_steer() + || !self.pending_steers.is_empty() + || !self.steer_mutations.is_empty() + || self.queue_handoff + { + self.toast( + "session switching requires idle state with no checkout or pending-message work", + ); + return false; + } + let cross_session = self.session_id.as_deref() != Some(requested_id); + if cross_session { + // Deleting a placeholder explicitly removes that attachment from + // the draft. Use the same cleanup as submission, not invisible + // stale metadata that could make an empty draft unswitchable. + self.prune_attachments(); + } + if cross_session && (!self.editor.is_empty() || !self.attachments.is_empty()) { + self.toast( + "unsent draft: send or clear the prompt and attachments before switching sessions", + ); + return false; + } + true + } + fn handle_session_key(&mut self, key: KeyEvent, pasted: bool) -> Action { + if self.session_catalog_pending { + match key.code { + KeyCode::Esc | KeyCode::F(4) => { + self.dismiss_sessions(); + return Action::None; + } + KeyCode::Enter => return Action::None, + _ => {} + } + } + let selected_id = self + .session_dialog + .as_ref() + .and_then(|dialog| self.session_matches.get(dialog.selected)) + .map(|matched| self.session_forest.rows[matched.index].entry.id.clone()); let Some(dialog) = &mut self.session_dialog else { return Action::None; }; + if pasted && matches!(key.code, KeyCode::Enter | KeyCode::Tab) { + return Action::None; + } if let Some(rename) = &mut dialog.rename { match rename { @@ -2666,11 +2830,10 @@ impl App { dialog.rename = Some(SessionRename::Editing(String::new())); } KeyCode::Enter => { - let selected = dialog.selected; dialog.rename = Some(SessionRename::Saving); - if let Some(entry) = self.session_choices.get(selected) { + if let Some(session_id) = selected_id { return Action::RenameSession { - session_id: entry.id.clone(), + session_id, display_name: None, }; } @@ -2689,12 +2852,11 @@ impl App { dialog.rename = Some(SessionRename::ConfirmClear); } KeyCode::Enter => { - let selected = dialog.selected; let display_name = input.clone(); dialog.rename = Some(SessionRename::Saving); - if let Some(entry) = self.session_choices.get(selected) { + if let Some(session_id) = selected_id { return Action::RenameSession { - session_id: entry.id.clone(), + session_id, display_name: Some(display_name), }; } @@ -2712,23 +2874,54 @@ impl App { return Action::None; } + if dialog.searching { + match key.code { + KeyCode::Esc | KeyCode::Enter => dialog.searching = false, + KeyCode::Backspace => { + if let Some((index, _)) = dialog.query.grapheme_indices(true).next_back() { + dialog.query.truncate(index); + } + dialog.selected = 0; + self.filter_sessions(None); + } + KeyCode::Char(character) + if key.modifiers.is_empty() || key.modifiers == KeyModifiers::SHIFT => + { + self.insert_session_query(&character.to_string()); + } + _ => {} + } + return Action::None; + } match key.code { - KeyCode::Esc => self.session_dialog = None, + KeyCode::Esc | KeyCode::F(4) => self.dismiss_sessions(), KeyCode::Up => dialog.selected = dialog.selected.saturating_sub(1), KeyCode::Down => { dialog.selected = - (dialog.selected + 1).min(self.session_choices.len().saturating_sub(1)); + (dialog.selected + 1).min(self.session_matches.len().saturating_sub(1)); } - KeyCode::Char('r' | 'R') => { + KeyCode::Char('/') if key.modifiers.is_empty() => dialog.searching = true, + KeyCode::Char('r' | 'R') + if !pasted + && !key.modifiers.contains(KeyModifiers::CONTROL) + && selected_id.is_some() => + { dialog.rename = Some(SessionRename::Editing(String::new())); } KeyCode::Enter => { - if let Some(entry) = self.session_choices.get(dialog.selected) { - let id = entry.id.clone(); - self.session_dialog = None; + if let Some(id) = selected_id + && self.session_switch_allowed(&id) + { + self.dismiss_sessions(); return Action::Resume(id); } } + KeyCode::Char(character) + if key.modifiers.is_empty() || key.modifiers == KeyModifiers::SHIFT => + { + dialog.searching = true; + self.insert_session_query(&character.to_string()); + } _ => {} } Action::None @@ -3127,6 +3320,9 @@ impl App { }); return Action::None; } + if key.code == KeyCode::F(2) && key.modifiers.is_empty() && !pasted { + return self.open_branch_chooser(); + } let matches = self.sync_navigation(); let selected = self .navigation @@ -3267,7 +3463,15 @@ impl App { } fn open_branch_chooser(&mut self) -> Action { - if self.working() || self.editing_steer() || !self.pending_steers.is_empty() { + if !self.require_session_connection() { + return Action::None; + } + if self.working() + || self.editing_branch() + || self.editing_steer() + || !self.pending_steers.is_empty() + || !self.steer_mutations.is_empty() + { self.toast( "prompt checkout is available only while idle and outside pending-message edits", ); @@ -3343,6 +3547,7 @@ impl App { provisional.auth_methods = self.auth_methods.clone(); provisional.show_thoughts = self.show_thoughts; provisional.branch_epoch = epoch; + provisional.session_catalog_epoch = self.session_catalog_epoch.wrapping_add(1); for update in response.prefix { let (_, updates) = super::translate( agent_client_protocol::schema::v2::UpdateSessionNotification::new( @@ -3368,6 +3573,7 @@ impl App { pub(super) fn abandon_branch(&mut self) { let epoch = self.branch_epoch.wrapping_add(1); + let catalog_epoch = self.session_catalog_epoch; if let Some(draft) = self.branch_draft.take() { let input_state = ( self.input_overflow, @@ -3380,6 +3586,9 @@ impl App { self.input_recovery_ready, self.last_key, ) = input_state; + // View restoration must not restore a spent catalog request epoch. + self.session_catalog_epoch = catalog_epoch; + self.dismiss_sessions(); } self.branch_chooser = None; self.branch_epoch = epoch; @@ -3558,6 +3767,18 @@ impl App { } return Action::Redraw; } + if self.session_dialog.is_some() { + // Terminals without bracketed paste deliver a paste as a key burst, so + // the arrival gap is the only thing separating it from typing. + let pasted = self + .last_key + .is_some_and(|last| received_at.saturating_duration_since(last) < PASTE_GAP); + self.last_key = Some(received_at); + return self.handle_session_key(key, pasted); + } + if key.code == KeyCode::F(4) && key.modifiers.is_empty() { + return self.open_sessions(); + } if self.branch_chooser.is_some() || self.branch_draft.is_some() { let action = self.handle_branch_key_at(key, received_at); self.last_key = Some(received_at); @@ -3570,15 +3791,6 @@ impl App { self.last_key = Some(received_at); return self.handle_navigation_key(key, pasted); } - if self.session_dialog.is_some() { - // Terminals without bracketed paste deliver a paste as a key burst, so - // the arrival gap is the only thing separating it from typing. - let pasted = self - .last_key - .is_some_and(|last| received_at.saturating_duration_since(last) < PASTE_GAP); - self.last_key = Some(received_at); - return self.handle_session_key(key, pasted); - } if self.model_dialog.is_some() { return self.handle_model_key(key); } @@ -3800,12 +4012,30 @@ impl App { self.open_navigation(); return Action::None; } + if matches!( + parse(self.editor.text(), !self.auth_methods.is_empty()), + Parsed::Sessions + ) { + if !self.editing_steer() { + self.editor.clear(); + } + return self.open_sessions(); + } if matches!( parse(self.editor.text(), !self.auth_methods.is_empty()), Parsed::Branch ) { return self.open_branch_chooser(); } + if !self.session_connected + && !matches!( + parse(self.editor.text(), !self.auth_methods.is_empty()), + Parsed::Resume { .. } | Parsed::Close | Parsed::Agents + ) + { + self.require_session_connection(); + return Action::None; + } if self.editor.is_empty() { return Action::None; } @@ -3868,15 +4098,7 @@ impl App { self.toast("usage: /resume "); Action::None } - Parsed::Sessions => { - if self.session_catalog_pending { - self.toast("session catalog scan is already in progress"); - Action::None - } else { - self.session_catalog_pending = true; - Action::ListSessions - } - } + Parsed::Sessions => self.open_sessions(), Parsed::Branch => self.open_branch_chooser(), Parsed::Transcript => { self.open_navigation(); @@ -4058,7 +4280,8 @@ impl App { } pub fn handle_mouse(&mut self, mouse: MouseEvent) -> Action { - if self.model_switch.is_some() + if self.session_dialog.is_some() + || self.model_switch.is_some() || self.navigation.dialog.is_some() || self.branch_chooser.is_some() || self.editing_branch() @@ -4361,6 +4584,12 @@ mod test_support { } impl App { + pub(in crate::tui) fn session_rename_active(&self) -> bool { + self.session_dialog + .as_ref() + .is_some_and(|dialog| dialog.rename.is_some()) + } + pub(super) fn agents(&self) -> Vec<&AgentRow> { self.agent_tree_rows() .into_iter() @@ -7565,16 +7794,428 @@ mod tests { ] } + fn explorer_entries() -> Vec { + ["source", "sibling"] + .into_iter() + .enumerate() + .map(|(index, id)| crate::session::CatalogEntry { + id: id.into(), + title: Some(format!("{id} title")), + preview: Some(format!("latest {id}")), + is_subagent: false, + lineage: crate::session::CatalogLineage::Root, + branch_point: None, + updated_at: index as u64, + }) + .collect() + } + + fn open_explorer(app: &mut App) { + let Action::ListSessions { epoch } = app.handle_key(press(KeyCode::F(4))) else { + panic!("catalog requested") + }; + app.apply(Update::SessionCatalog { + epoch, + result: Ok(explorer_entries()), + }); + } + #[test] - fn sessions_command_defers_catalog_work_to_the_event_loop() { + fn loading_explorer_owns_keys_and_paste_without_submitting_or_steering_draft() { + for phase in [Phase::Idle, Phase::Working] { + let mut app = app(); + app.start_session("source".into()); + app.paste("parked draft"); + app.attach( + PathBuf::from("/tmp/image.png"), + "image/png", + AttachmentKind::Image, + 12, + ); + app.editor.move_left(); + let draft = app.editor.text().to_owned(); + let cursor = app.editor.cursor(); + let attachments = app.attachments.clone(); + app.phase = phase; + app.can_steer = true; + let Action::ListSessions { epoch } = app.handle_key(press(KeyCode::F(4))) else { + panic!("catalog requested") + }; + assert!(app.session_dialog.is_some()); + assert!(matches!( + app.handle_key(press(KeyCode::Char('s'))), + Action::None + )); + app.paste("ource"); + // The response is deliberately not delivered yet. A deliberate + // Enter must be harmless even outside the paste timing window. + std::thread::sleep(super::PASTE_GAP * 2); + assert!(matches!( + app.handle_key(press(KeyCode::Enter)), + Action::None + )); + assert_eq!(app.editor.text(), draft); + assert_eq!(app.editor.cursor(), cursor); + assert_eq!(app.attachments, attachments); + assert_eq!(app.session_dialog.as_ref().unwrap().query, "source"); + app.apply(Update::SessionCatalog { + epoch, + result: Ok(explorer_entries()), + }); + assert_eq!(app.session_dialog.as_ref().unwrap().query, "source"); + assert_eq!(app.session_matches.len(), 1); + assert_eq!( + app.session_forest.rows[app.session_matches[0].index] + .entry + .id, + "source" + ); + assert_eq!(app.editor.text(), draft); + assert_eq!(app.attachments, attachments); + } + } + + #[test] + fn loading_explorer_escape_reopen_retains_new_query_and_ignores_old_result() { let mut app = app(); - app.editor.insert_str("/sessions"); + app.paste("parked draft"); + let Action::ListSessions { epoch: old } = app.handle_key(press(KeyCode::F(4))) else { + panic!() + }; + app.paste("obsolete"); + app.handle_key(press(KeyCode::Esc)); + assert!(app.session_dialog.is_none()); + let Action::ListSessions { epoch: new } = app.handle_key(press(KeyCode::F(4))) else { + panic!() + }; + app.paste("source"); + app.apply(Update::SessionCatalog { + epoch: old, + result: Ok(explorer_entries()), + }); + assert!(app.session_matches.is_empty()); + assert_eq!(app.session_dialog.as_ref().unwrap().query, "source"); + std::thread::sleep(super::PASTE_GAP * 2); + assert!(matches!( + app.handle_key(press(KeyCode::Enter)), + Action::None + )); + app.apply(Update::SessionCatalog { + epoch: new, + result: Ok(explorer_entries()), + }); + assert_eq!(app.session_matches.len(), 1); + assert_eq!(app.editor.text(), "parked draft"); + } + + #[test] + fn disconnected_editor_and_checkout_preserve_draft_until_explicit_reload() { + let mut app = app(); + app.start_session("source".into()); + app.paste("parked draft"); + app.attach( + PathBuf::from("/tmp/image.png"), + "image/png", + AttachmentKind::Image, + 12, + ); + let draft = app.editor.text().to_owned(); + let attachments = app.attachments.clone(); + app.session_connected = false; + assert!(matches!( + app.handle_key(press(KeyCode::Enter)), + Action::None + )); + app.handle_key(press(KeyCode::F(3))); + std::thread::sleep(super::PASTE_GAP * 2); + assert!(matches!(app.handle_key(press(KeyCode::F(2))), Action::None)); + assert!(app.branch_chooser.is_none()); + assert_eq!(app.editor.text(), draft); + assert_eq!(app.attachments, attachments); + app.handle_key(press(KeyCode::Esc)); + open_explorer(&mut app); + app.paste("source"); + app.handle_key(press(KeyCode::Enter)); // Finish search, never submit. + std::thread::sleep(super::PASTE_GAP * 2); + assert!( + matches!(app.handle_key(press(KeyCode::Enter)), Action::Resume(id) if id == "source") + ); + assert!(!app.session_connected); // Only successful load can reconnect. + assert_eq!(app.editor.text(), draft); + assert_eq!(app.attachments, attachments); + } + + #[test] + fn explorer_request_epochs_reject_dismissed_and_reopened_catalogs() { + let mut app = app(); + let Action::ListSessions { epoch: old } = app.open_sessions() else { + panic!() + }; + app.handle_key(press(KeyCode::Esc)); + app.apply(Update::SessionCatalog { + epoch: old, + result: Ok(explorer_entries()), + }); + assert!(app.session_dialog.is_none()); + let Action::ListSessions { epoch: new } = app.open_sessions() else { + panic!() + }; + assert_ne!(old, new); + app.apply(Update::SessionCatalog { + epoch: old, + result: Ok(explorer_entries()), + }); + assert!(app.session_catalog_pending); + assert!(app.session_dialog.is_some()); + assert!(app.session_matches.is_empty()); + app.apply(Update::SessionCatalog { + epoch: new, + result: Ok(explorer_entries()), + }); + app.handle_key(press(KeyCode::Esc)); + app.apply(Update::SessionCatalog { + epoch: new, + result: Ok(explorer_entries()), + }); + assert!(app.session_dialog.is_none()); + } + + #[test] + fn explorer_epochs_remain_unique_across_provisional_view_and_source_restoration() { + let mut app = app(); + let Action::ListSessions { + epoch: source_request, + } = app.open_sessions() + else { + panic!() + }; + prepare_branch(&mut app); + let Action::ListSessions { + epoch: child_request, + } = app.open_sessions() + else { + panic!() + }; + assert_ne!(source_request, child_request); + app.apply(Update::SessionCatalog { + epoch: source_request, + result: Ok(explorer_entries()), + }); + assert!(app.session_dialog.is_some()); + assert!(app.session_matches.is_empty()); + app.abandon_branch(); + let Action::ListSessions { + epoch: restored_request, + } = app.open_sessions() + else { + panic!() + }; + assert_ne!(restored_request, source_request); + assert_ne!(restored_request, child_request); + app.apply(Update::SessionCatalog { + epoch: child_request, + result: Ok(explorer_entries()), + }); + assert!(app.session_dialog.is_some()); + assert!(app.session_matches.is_empty()); + app.apply(Update::SessionCatalog { + epoch: restored_request, + result: Ok(explorer_entries()), + }); + assert!(app.session_dialog.is_some()); + } + + #[test] + fn explorer_preserves_draft_and_attachments_and_rejects_dirty_cross_session_switch() { + let mut app = app(); + app.start_session("source".into()); + app.paste("unsent draft"); + app.attach( + PathBuf::from("/tmp/image.png"), + "image/png", + AttachmentKind::Image, + 12, + ); + app.editor.move_left(); + let draft = app.editor.text().to_owned(); + let cursor = app.editor.cursor(); + let attachments = app.attachments.clone(); + open_explorer(&mut app); + app.last_key = None; assert!(matches!( app.handle_key(press(KeyCode::Enter)), - Action::ListSessions + Action::None )); + assert!(app.toast.as_ref().unwrap().0.contains("unsent draft")); + assert!(app.session_dialog.is_some()); + app.phase = Phase::Working; + app.paste("source"); + app.last_key = None; + app.handle_key(press(KeyCode::Enter)); // leave search + app.last_key = None; + assert!(matches!( + app.handle_key(press(KeyCode::Enter)), + Action::None + )); + app.phase = Phase::Idle; + app.last_key = None; + assert!( + matches!(app.handle_key(press(KeyCode::Enter)), Action::Resume(id) if id == "source") + ); + app.start_session("source".into()); + assert_eq!(app.editor.text(), draft); + assert_eq!(app.editor.cursor(), cursor); + assert_eq!(app.attachments, attachments); assert!(app.session_dialog.is_none()); + } + + #[test] + fn explicitly_cleared_attachment_draft_no_longer_blocks_cross_session_switch() { + let mut app = app(); + app.start_session("source".into()); + app.attach( + PathBuf::from("/tmp/image.png"), + "image/png", + AttachmentKind::Image, + 12, + ); + assert!(!app.session_switch_allowed("sibling")); + app.editor.clear(); // Explicitly remove the prompt and its attachment placeholder. + assert!(app.session_switch_allowed("sibling")); + assert!(app.attachments.is_empty()); + } + + #[test] + fn explorer_filter_rename_targets_visible_id_and_preserves_selection() { + let mut app = app(); + open_explorer(&mut app); + app.paste("source"); + assert_eq!(app.session_matches.len(), 1); + app.handle_key(press(KeyCode::Enter)); + app.last_key = None; + app.handle_key(press(KeyCode::Char('r'))); + app.paste("Renamed source"); + assert!( + matches!(app.handle_key(press(KeyCode::Enter)), Action::RenameSession { session_id, .. } if session_id == "source") + ); + app.apply(Update::SessionRenamed { + session_id: "source".into(), + display_name: Some("Renamed source".into()), + result: Ok(Some("Renamed source".into())), + }); + let selected = app.session_dialog.as_ref().unwrap().selected; + assert_eq!( + app.session_forest.rows[app.session_matches[selected].index] + .entry + .id, + "source" + ); + assert_eq!( + app.session_forest.rows[app.session_matches[selected].index] + .entry + .title + .as_deref(), + Some("Renamed source") + ); + } + + #[test] + fn explorer_unicode_paste_is_bounded_and_search_uses_receipt_time() { + let mut app = app(); + app.paste("parked"); + open_explorer(&mut app); + app.paste(&format!("\n\t{}", "界👩‍💻".repeat(2000))); + let query = &app.session_dialog.as_ref().unwrap().query; + assert!(query.len() <= 4096); + assert!(!query.contains('\n')); + assert!(app.session_matches.is_empty()); + assert_eq!(app.editor.text(), "parked"); + let received_at = Instant::now() - super::PASTE_GAP * 4; + app.last_key = Some(received_at); + app.handle_key_at(press(KeyCode::Enter), received_at + super::PASTE_GAP / 2); + assert!(app.session_dialog.as_ref().unwrap().searching); + app.session_dialog.as_mut().unwrap().query.clear(); + app.filter_sessions(None); + assert_eq!(app.session_matches.len(), 2); + app.last_key = Some(received_at); + app.handle_key_at(press(KeyCode::Enter), received_at + super::PASTE_GAP * 2); + assert!(!app.session_dialog.as_ref().unwrap().searching); + } + + #[test] + fn explorer_switch_gates_cover_steer_edit_mutation_handoff_and_all_busy_phases() { + let mut app = queued_app(); + app.paste("parked draft"); + begin_steer_edit(&mut app); + let pending_edit = app.editor.text().to_owned(); + app.phase = Phase::Idle; + open_explorer(&mut app); + assert!(!app.session_switch_allowed("source")); + app.paste("source"); + app.handle_key(press(KeyCode::Esc)); // leave search + app.handle_key(press(KeyCode::Esc)); // close explorer + assert_eq!(app.editor.text(), pending_edit); + assert!(app.editing_steer()); + app.cancel_steer_edit(); + app.pending_steers.clear(); + app.editor.clear(); + app.queue_handoff = true; + assert!(!app.session_switch_allowed("source")); + app.queue_handoff = false; + for phase in [Phase::Working, Phase::Blocked, Phase::Cancelling] { + app.phase = phase; + assert!(!app.session_switch_allowed("source")); + } + app.phase = Phase::Idle; + assert!(app.session_switch_allowed("source")); + } + + #[test] + fn explorer_browses_provisional_checkout_but_cannot_switch() { + let mut app = app(); + prepare_branch(&mut app); + let text = app.editor.text().to_owned(); + open_explorer(&mut app); + app.last_key = None; + assert!(matches!( + app.handle_key(press(KeyCode::Enter)), + Action::None + )); + assert!(app.editing_branch()); + app.handle_key(press(KeyCode::Esc)); + assert_eq!(app.editor.text(), text); + assert!(app.editing_branch()); + } + + #[test] + fn transcript_footer_action_opens_authoritative_chooser_with_idle_guard() { + let mut app = app(); + app.paste("parked draft"); + app.open_navigation(); + app.phase = Phase::Working; + assert!(matches!( + app.handle_navigation_key(press(KeyCode::F(2)), false), + Action::None + )); + assert!(app.branch_chooser.is_none()); + app.phase = Phase::Idle; + assert!(matches!( + app.handle_navigation_key(press(KeyCode::F(2)), false), + Action::ListPromptBranches { .. } + )); + assert_eq!(app.editor.text(), "parked draft"); + } + + #[test] + fn sessions_command_defers_catalog_work_to_the_event_loop() { + let mut app = app(); app.editor.insert_str("/sessions"); + assert!(matches!( + app.handle_key(press(KeyCode::Enter)), + Action::ListSessions { .. } + )); + assert!(app.session_dialog.is_some()); + assert!(app.editor.is_empty()); assert!(matches!( app.handle_key(press(KeyCode::Enter)), Action::None @@ -7586,22 +8227,29 @@ mod tests { } #[test] - fn session_catalog_update_opens_the_dialog() { + fn session_catalog_update_populates_the_immediately_opened_dialog() { let mut app = app(); assert!(matches!( app.handle_key(press(KeyCode::Char('@'))), Action::SearchFiles { .. } )); - app.apply(Update::SessionCatalog(Ok(vec![ - crate::session::CatalogEntry { + let Action::ListSessions { epoch } = app.handle_key(press(KeyCode::F(4))) else { + panic!("catalog requested") + }; + assert!(app.session_dialog.is_some()); + app.apply(Update::SessionCatalog { + epoch, + result: Ok(vec![crate::session::CatalogEntry { id: "saved".into(), title: Some("Saved".into()), preview: None, is_subagent: false, + lineage: crate::session::CatalogLineage::Root, + branch_point: None, updated_at: 0, - }, - ]))); - assert_eq!(app.session_choices[0].id, "saved"); + }]), + }); + assert_eq!(app.session_forest.rows[0].entry.id, "saved"); assert!(app.session_dialog.is_some()); assert!(app.file_picker.is_none()); } @@ -7609,22 +8257,29 @@ mod tests { #[test] fn session_dialog_selects_a_catalog_entry_for_existing_resume_flow() { let mut app = app(); - app.session_choices = ["newer", "older"] - .into_iter() - .map(|id| crate::session::CatalogEntry { - id: id.into(), - title: Some(format!("{id} title")), - preview: None, - is_subagent: false, - updated_at: 0, - }) - .collect(); + app.set_session_choices( + ["newer", "older"] + .into_iter() + .map(|id| crate::session::CatalogEntry { + id: id.into(), + title: Some(format!("{id} title")), + preview: None, + is_subagent: false, + lineage: crate::session::CatalogLineage::Root, + branch_point: None, + updated_at: u64::from(id == "newer"), + }) + .collect(), + ); app.session_dialog = Some(super::SessionDialog { + query: String::new(), + searching: false, selected: 0, rename: None, }); assert!(matches!(app.handle_key(press(KeyCode::Down)), Action::None)); + app.last_key = None; // A deliberate activation, not an unbracketed paste burst. assert!(matches!( app.handle_key(press(KeyCode::Enter)), Action::Resume(id) if id == "older" @@ -7635,17 +8290,23 @@ mod tests { #[test] fn session_dialog_renames_in_place_and_preserves_selection() { let mut app = app(); - app.session_choices = ["newer", "older"] - .into_iter() - .map(|id| crate::session::CatalogEntry { - id: id.into(), - title: Some(format!("{id} title")), - preview: None, - is_subagent: false, - updated_at: 0, - }) - .collect(); + app.set_session_choices( + ["newer", "older"] + .into_iter() + .map(|id| crate::session::CatalogEntry { + id: id.into(), + title: Some(format!("{id} title")), + preview: None, + is_subagent: false, + lineage: crate::session::CatalogLineage::Root, + branch_point: None, + updated_at: u64::from(id == "newer"), + }) + .collect(), + ); app.session_dialog = Some(super::SessionDialog { + query: String::new(), + searching: false, selected: 1, rename: None, }); @@ -7676,7 +8337,10 @@ mod tests { result: Ok(Some("OAuth bug".into())), }); assert_eq!(app.session_dialog.as_ref().unwrap().selected, 1); - assert_eq!(app.session_choices[1].title.as_deref(), Some("OAuth bug")); + assert_eq!( + app.session_forest.rows[1].entry.title.as_deref(), + Some("OAuth bug") + ); app.session_dialog.as_mut().unwrap().rename = Some(super::SessionRename::Saving); app.apply(Update::SessionRenamed { @@ -7693,14 +8357,18 @@ mod tests { #[test] fn session_dialog_confirms_before_clearing_a_name() { let mut app = app(); - app.session_choices = vec![crate::session::CatalogEntry { + app.set_session_choices(vec![crate::session::CatalogEntry { id: "saved".into(), title: Some("Generated".into()), preview: None, is_subagent: false, + lineage: crate::session::CatalogLineage::Root, + branch_point: None, updated_at: 0, - }]; + }]); app.session_dialog = Some(super::SessionDialog { + query: String::new(), + searching: false, selected: 0, rename: Some(super::SessionRename::Editing(String::new())), }); @@ -7733,6 +8401,8 @@ mod tests { fn session_rename_backspace_removes_a_complete_grapheme() { let mut app = app(); app.session_dialog = Some(super::SessionDialog { + query: String::new(), + searching: false, selected: 0, rename: Some(super::SessionRename::Editing("e\u{301} 👨‍👩‍👧".into())), }); @@ -7751,6 +8421,8 @@ mod tests { fn session_rename_clear_confirmation_ignores_enter_from_a_paste_burst() { let mut app = app(); app.session_dialog = Some(super::SessionDialog { + query: String::new(), + searching: false, selected: 0, rename: Some(super::SessionRename::ConfirmClear), }); @@ -7771,6 +8443,8 @@ mod tests { for input in [String::new(), "Pasted name".into()] { let mut app = app(); app.session_dialog = Some(super::SessionDialog { + query: String::new(), + searching: false, selected: 0, rename: Some(super::SessionRename::Editing(input.clone())), }); diff --git a/src/tui/branches.rs b/src/tui/branches.rs new file mode 100644 index 0000000..82569b5 --- /dev/null +++ b/src/tui/branches.rs @@ -0,0 +1,501 @@ +//! Pure, iterative branch navigation over one read-only catalog snapshot. +//! No transcript loading, persistence, or runtime activation belongs here. + +use std::collections::HashMap; + +use crate::session::{CatalogEntry, CatalogLineage}; + +/// Deep chains keep their real ancestry but cannot consume the terminal width. +pub(super) const MAX_DISPLAY_DEPTH: usize = 16; +const MAX_QUERY_BYTES: usize = 4096; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub(super) struct BranchRow { + pub entry: CatalogEntry, + /// Index in `BranchForest::rows`, always preceding this row. + pub parent: Option, + pub depth: usize, + /// Invalid metadata, missing parents, and broken cycles are visible roots. + pub warning: Option, +} + +impl BranchRow { + pub fn display_depth(&self) -> usize { + self.depth.min(MAX_DISPLAY_DEPTH) + } + + /// Render alongside the row: indentation is capped, never silently lost. + pub fn depth_note(&self) -> Option { + (self.depth > MAX_DISPLAY_DEPTH).then(|| { + format!( + "{} ancestry levels not indented", + self.depth - MAX_DISPLAY_DEPTH + ) + }) + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(super) struct BranchMatch { + pub index: usize, + /// False means this row is retained solely to explain a matching descendant. + pub direct_match: bool, +} + +#[derive(Default, Debug)] +pub(super) struct BranchForest { + /// Preorder, with stable newest-activity then descending-ID sibling ordering. + /// Every input row occurs exactly once, including malformed/duplicate rows. + pub rows: Vec, + search_text: Vec<[String; 3]>, +} + +impl BranchForest { + pub fn new(mut entries: Vec) -> Self { + entries.sort_by(|left, right| { + right + .updated_at + .cmp(&left.updated_at) + .then_with(|| right.id.cmp(&left.id)) + // Catalog IDs are unique. Tie-break malformed synthetic inputs + // too, without making hash-map iteration part of the ordering. + .then_with(|| left.title.cmp(&right.title)) + .then_with(|| left.preview.cmp(&right.preview)) + .then_with(|| left.branch_point.cmp(&right.branch_point)) + .then_with(|| left.is_subagent.cmp(&right.is_subagent)) + .then_with(|| left.lineage.cmp(&right.lineage)) + }); + let count = entries.len(); + let mut by_id = HashMap::with_capacity(count); + for (index, entry) in entries.iter().enumerate() { + by_id + .entry(entry.id.as_str()) + .and_modify(|value| *value = None) + .or_insert(Some(index)); + } + let mut parents = vec![None; count]; + let mut warnings = vec![None; count]; + for (index, entry) in entries.iter().enumerate() { + if by_id[entry.id.as_str()].is_none() { + warnings[index] = Some("Duplicate session ID; ancestry is ambiguous".into()); + continue; + } + match &entry.lineage { + CatalogLineage::Root => {} + CatalogLineage::Warning(warning) => warnings[index] = Some(warning.clone()), + CatalogLineage::Branch(metadata) => { + match by_id.get(metadata.parent_session_id.as_str()) { + Some(Some(parent)) => parents[index] = Some(*parent), + Some(None) => { + warnings[index] = Some(format!( + "Ambiguous parent {}; shown as an orphan root", + metadata.parent_session_id + )); + } + None => { + warnings[index] = Some(format!( + "Missing parent {}; shown as an orphan root", + metadata.parent_session_id + )); + } + } + } + } + } + // A functional parent graph admits linear cycle detection. Each node + // is walked once; no recursive DFS or per-row ancestor walk is needed. + let mut colors = vec![0_u8; count]; + let mut positions = vec![0; count]; + let mut path: Vec = Vec::new(); + for start in 0..count { + if colors[start] != 0 { + continue; + } + path.clear(); + let mut next = Some(start); + while let Some(index) = next { + if colors[index] == 2 { + break; + } + if colors[index] == 1 { + // Cut the lexicographically smallest cycle member's edge, + // independent of input order, activity, or descendant order. + let cut = *path[positions[index]..] + .iter() + .min_by_key(|&&member| &entries[member].id) + .expect("a gray node is on this nonempty path"); + let parent = parents[cut].take().expect("cycle member has a parent"); + warnings[cut] = Some(format!( + "Cycle edge to {} omitted; shown as a root", + entries[parent].id + )); + break; + } + colors[index] = 1; + positions[index] = path.len(); + path.push(index); + next = parents[index]; + } + for &index in &path { + colors[index] = 2; + } + } + let mut children = vec![Vec::new(); count]; + let mut roots = Vec::new(); + for (index, parent) in parents.iter().enumerate() { + if let Some(parent) = parent { + children[*parent].push(index); + } else { + roots.push(index); + } + } + let mut stack: Vec<_> = roots + .into_iter() + .rev() + .map(|index| (index, None, 0)) + .collect(); + let mut entries: Vec<_> = entries.into_iter().map(Some).collect(); + let mut rows = Vec::with_capacity(count); + while let Some((index, parent, depth)) = stack.pop() { + let row_index = rows.len(); + rows.push(BranchRow { + entry: entries[index].take().expect("each row is visited once"), + parent, + depth, + warning: warnings[index].take(), + }); + stack.extend( + children[index] + .iter() + .rev() + .map(|&child| (child, Some(row_index), depth + 1)), + ); + } + let search_text = rows + .iter() + .map(|row| { + [ + row.entry.id.to_lowercase(), + row.entry + .title + .as_deref() + .unwrap_or_default() + .to_lowercase(), + row.entry + .preview + .as_deref() + .unwrap_or_default() + .to_lowercase(), + ] + }) + .collect(); + Self { rows, search_text } + } + + /// Case-insensitive name/ID/preview search. Retain all matching ancestors in + /// preorder and label them via `direct_match`; unrelated descendants stay out. + /// Work and storage are linear even for a matching leaf in a very deep chain. + pub fn search(&self, query: &str) -> Vec { + let query = query.trim(); + let mut end = query.len().min(MAX_QUERY_BYTES); + while !query.is_char_boundary(end) { + end -= 1; + } + let query = query[..end].to_lowercase(); + let direct: Vec<_> = self + .search_text + .iter() + .map(|fields| query.is_empty() || fields.iter().any(|field| field.contains(&query))) + .collect(); + let mut included = direct.clone(); + for index in (0..self.rows.len()).rev() { + if included[index] + && let Some(parent) = self.rows[index].parent + { + included[parent] = true; + } + } + included + .into_iter() + .enumerate() + .filter(|(_, include)| *include) + .map(|(index, _)| BranchMatch { + index, + direct_match: direct[index], + }) + .collect() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::session::{ + BranchBoundary, BranchCompletion, BranchProvenance, BranchRequest, BranchSelection, + }; + + fn entry(id: &str, parent: Option<&str>, updated_at: u64) -> CatalogEntry { + CatalogEntry { + id: id.into(), + title: None, + preview: None, + branch_point: None, + is_subagent: false, + updated_at, + lineage: parent.map_or(CatalogLineage::Root, |parent| { + CatalogLineage::Branch(Box::new(BranchProvenance { + version: 1, + parent_session_id: parent.into(), + boundary: BranchBoundary { + state_index: 0, + prefix_len: 1, + prefix_hash: "a".repeat(64), + }, + checkout_id: format!("checkout-{id}"), + request: BranchRequest { + id: format!("request-{id}"), + selection: BranchSelection { + provider: "openrouter".into(), + model: "model".into(), + reasoning: "default".into(), + }, + }, + completion: BranchCompletion { + prefix_len: 1, + prefix_hash: "b".repeat(64), + prompt_hash: "c".repeat(64), + }, + })) + }), + } + } + + fn ids(forest: &BranchForest) -> Vec<&str> { + forest + .rows + .iter() + .map(|row| row.entry.id.as_str()) + .collect() + } + + fn assert_forest(forest: &BranchForest, count: usize) { + assert_eq!(forest.rows.len(), count); + for (index, row) in forest.rows.iter().enumerate() { + if let Some(parent) = row.parent { + assert!(parent < index); + assert_eq!(row.depth, forest.rows[parent].depth + 1); + } else { + assert_eq!(row.depth, 0); + } + } + } + + #[test] + fn empty_roots_nested_and_siblings_use_stable_activity_then_id_order() { + let empty = BranchForest::new(Vec::new()); + assert!(empty.rows.is_empty()); + assert!(empty.search("").is_empty()); + let entries = vec![ + entry("root-a", None, 10), + entry("root-b", None, 10), + entry("old", Some("root-a"), 1), + entry("new-a", Some("root-a"), 30), + entry("new-b", Some("root-a"), 30), + entry("grandchild", Some("old"), 100), + ]; + let expected = BranchForest::new(entries.clone()); + assert_eq!( + ids(&expected), + ["root-b", "root-a", "new-b", "new-a", "old", "grandchild"] + ); + assert_forest(&expected, entries.len()); + for shift in 0..entries.len() { + let mut shuffled = entries.clone(); + shuffled.rotate_left(shift); + shuffled.reverse(); + assert_eq!(BranchForest::new(shuffled).rows, expected.rows); + } + } + + #[test] + fn missing_parents_and_metadata_warnings_stay_visible_as_roots() { + let mut malformed = entry("malformed", None, 50); + malformed.lineage = CatalogLineage::Warning("incomplete prompt checkout".into()); + let forest = BranchForest::new(vec![ + malformed, + entry("orphan", Some("outside-workspace"), 10), + entry("nested-orphan", Some("orphan"), 20), + entry("legacy", None, 0), + ]); + assert_forest(&forest, 4); + assert_eq!( + ids(&forest), + ["malformed", "orphan", "nested-orphan", "legacy"] + ); + assert!( + forest.rows[0] + .warning + .as_deref() + .unwrap() + .contains("incomplete") + ); + assert!( + forest.rows[1] + .warning + .as_deref() + .unwrap() + .contains("outside-workspace") + ); + assert_eq!(forest.rows[2].parent, Some(1)); + assert!(forest.rows[3].warning.is_none()); + } + + #[test] + fn cycles_cut_smallest_id_edge_and_preserve_every_row_once() { + let entries = vec![ + entry("a", Some("b"), 2), + entry("b", Some("c"), 20), + entry("c", Some("a"), 30), + entry("tail", Some("b"), 99), + entry("self", Some("self"), 10), + entry("unrelated", None, 0), + ]; + let expected = BranchForest::new(entries.clone()); + assert_forest(&expected, entries.len()); + let mut unique = ids(&expected); + unique.sort_unstable(); + unique.dedup(); + assert_eq!(unique.len(), entries.len()); + let cut = expected + .rows + .iter() + .find(|row| row.entry.id == "a") + .unwrap(); + assert!(cut.parent.is_none()); + assert!(cut.warning.as_deref().unwrap().contains("Cycle edge to b")); + let self_cycle = expected + .rows + .iter() + .find(|row| row.entry.id == "self") + .unwrap(); + assert!(self_cycle.parent.is_none()); + assert!(self_cycle.warning.is_some()); + for shift in 0..entries.len() { + let mut shuffled = entries.clone(); + shuffled.rotate_left(shift); + shuffled.reverse(); + assert_eq!(BranchForest::new(shuffled).rows, expected.rows); + } + let mut changed_activity = entries; + changed_activity[0].updated_at = 1000; + let forest = BranchForest::new(changed_activity); + assert!( + forest + .rows + .iter() + .find(|row| row.entry.id == "a") + .unwrap() + .parent + .is_none() + ); + } + + #[test] + fn search_matches_name_id_preview_and_retains_only_required_ancestors() { + let mut child = entry("child", Some("root"), 5); + child.title = Some("Café RENAMED".into()); + let mut leaf = entry("leaf-id", Some("child"), 10); + leaf.preview = Some("Needle in the preview".into()); + let forest = BranchForest::new(vec![ + entry("root", None, 0), + child, + leaf, + entry("other", Some("root"), 8), + ]); + for query in ["LEAF-ID", " needle "] { + let matches = forest.search(query); + assert_eq!( + matches + .iter() + .map(|found| forest.rows[found.index].entry.id.as_str()) + .collect::>(), + ["root", "child", "leaf-id"] + ); + assert_eq!( + matches + .iter() + .map(|found| found.direct_match) + .collect::>(), + [false, false, true] + ); + } + assert_eq!(forest.search("cAFÉ renamed").len(), 2); + assert_eq!(forest.search("root").len(), 1); + assert_eq!(forest.search(" ").len(), 4); + assert!(forest.search("absent").is_empty()); + assert!(forest.search(&"界".repeat(MAX_QUERY_BYTES)).is_empty()); + } + + #[test] + fn deep_and_wide_forests_are_iterative_and_indentation_is_bounded() { + const COUNT: usize = 20_000; + let mut chain = Vec::with_capacity(COUNT); + for index in 0..COUNT { + let parent = index.checked_sub(1).map(|index| format!("node-{index}")); + chain.push(entry( + &format!("node-{index}"), + parent.as_deref(), + index as u64, + )); + } + let forest = BranchForest::new(chain); + assert_forest(&forest, COUNT); + let leaf = forest.rows.last().unwrap(); + assert_eq!(leaf.depth, COUNT - 1); + assert_eq!(leaf.display_depth(), MAX_DISPLAY_DEPTH); + assert!( + leaf.depth_note() + .unwrap() + .contains("ancestry levels not indented") + ); + assert!(forest.rows[0].depth_note().is_none()); + let matches = forest.search(&format!("node-{}", COUNT - 1)); + assert_eq!(matches.len(), COUNT); + assert_eq!(matches.iter().filter(|found| found.direct_match).count(), 1); + drop(forest); // Flat ownership also makes destruction stack-safe. + let mut wide = vec![entry("root", None, 0)]; + wide.extend( + (1..COUNT).map(|index| entry(&format!("child-{index}"), Some("root"), index as u64)), + ); + let forest = BranchForest::new(wide); + assert_forest(&forest, COUNT); + assert!( + forest.rows[1..] + .iter() + .all(|row| row.parent == Some(0) && row.depth == 1) + ); + } + + #[test] + fn duplicate_ids_do_not_drop_rows_or_invent_an_ambiguous_parent() { + let mut duplicate = entry("duplicate", None, 1); + duplicate.title = Some("renamed duplicate".into()); + let entries = vec![ + entry("duplicate", None, 1), + duplicate, + entry("duplicate", Some("different-parent"), 1), + entry("child", Some("duplicate"), 10), + ]; + let forest = BranchForest::new(entries.clone()); + assert_forest(&forest, 4); + assert!( + forest + .rows + .iter() + .all(|row| row.parent.is_none() && row.warning.is_some()) + ); + let mut reversed = entries; + reversed.reverse(); + assert_eq!(BranchForest::new(reversed).rows, forest.rows); + } +} diff --git a/src/tui/command.rs b/src/tui/command.rs index 5d4fcc9..4c2c286 100644 --- a/src/tui/command.rs +++ b/src/tui/command.rs @@ -62,7 +62,7 @@ const LOCAL_COMMANDS: &[Spec] = &[ }, Spec { token: "/sessions", - description: "Browse saved sessions", + description: "Browse saved conversation branches (F4)", kind: Kind::Sessions, }, Spec { @@ -95,6 +95,11 @@ const LOCAL_COMMANDS: &[Spec] = &[ description: "Edit a previous text prompt in a new session", kind: Kind::Branch, }, + Spec { + token: "/branches", + description: "Browse saved conversation branches (F4)", + kind: Kind::Sessions, + }, Spec { token: "/login", description: "Authenticate with the agent", @@ -256,6 +261,8 @@ mod tests { } ); assert_eq!(parse("/sessions"), Parsed::Sessions); + assert_eq!(parse("/branches"), Parsed::Sessions); + assert_eq!(parse("/branch"), Parsed::Branch); assert_eq!(parse("/close"), Parsed::Close); assert_eq!(parse("/model"), Parsed::Model { query: None }); assert_eq!(parse("/effort"), Parsed::Effort { value: None }); @@ -423,6 +430,7 @@ mod tests { "/agents", "/transcript", "/branch", + "/branches", "/compact", ] ); diff --git a/src/tui/mod.rs b/src/tui/mod.rs index 7d4df26..0323442 100644 --- a/src/tui/mod.rs +++ b/src/tui/mod.rs @@ -7,6 +7,7 @@ //! ([`crate::events`]) that feeds the live graph of a running Runlet program. mod app; +mod branches; mod command; mod editor; mod image; @@ -174,6 +175,7 @@ const MAX_OUTPUT_LINES: usize = 5_000; struct ActiveSessionRoute { id: String, generation: u64, + runtime_activation: Option, } enum BranchResponse { @@ -194,6 +196,7 @@ struct BranchCompletion { struct BranchReplayBuffer { request: Option<(u64, u64)>, notifications: Vec, + diagnostics: Vec, } impl BranchCompletion { @@ -203,17 +206,42 @@ impl BranchCompletion { } impl BranchReplayBuffer { - fn finish(&mut self, generation: u64, epoch: u64) -> Vec { + fn finish( + &mut self, + generation: u64, + epoch: u64, + ) -> (Vec, Vec) { if self.request != Some((generation, epoch)) { - return Vec::new(); + return (Vec::new(), Vec::new()); } self.request = None; - std::mem::take(&mut self.notifications) + ( + std::mem::take(&mut self.notifications), + std::mem::take(&mut self.diagnostics), + ) + } + + fn route_diagnostic( + &mut self, + current: &ActiveSessionRoute, + update: QueuedUpdate, + ) -> Option { + if self.request.is_some() + && matches!(&update.update, Update::RoutedRuntime { session_id, .. } if session_id != ¤t.id) + { + // Child identity is still unverified. Retain, but do not apply it; + // the successful submit response will authorize exactly one epoch. + self.diagnostics.push(update); + None + } else { + Some(update) + } } } struct QueuedUpdate { generation: Option, + runtime_activation: Option, update: Update, } @@ -221,6 +249,7 @@ impl QueuedUpdate { fn global(update: Update) -> Self { Self { generation: None, + runtime_activation: None, update, } } @@ -228,28 +257,66 @@ impl QueuedUpdate { fn for_session(generation: u64, update: Update) -> Self { Self { generation: Some(generation), + runtime_activation: None, update, } } } -// Stderr and ACP are independent streams. Capture identity while reading stderr, -// not when applying its queued events after a possibly newer ACP activation. -// Markers are stream boundaries, never commands to switch the visible session. -fn runtime_diagnostic_update( +// Stderr and ACP are independent streams. The envelope records emission-time +// identity, and only a successful ACP response can authorize that identity. +// Never assign the current generation while reading delayed stderr bytes. +fn runtime_diagnostic_envelope_update( route: &mut Option, - event: events::RuntimeEvent, -) -> Option { + diagnostic: events::DiagnosticEvent, +) -> Option { + let events::DiagnosticEvent { event, activation } = diagnostic; match event { events::RuntimeEvent::SessionStarted { session_id } => { *route = Some(session_id); None } - event @ events::RuntimeEvent::StorageStatus { .. } => Some(Update::Runtime(event)), - event => route.as_ref().map(|session_id| Update::RoutedRuntime { - session_id: session_id.clone(), + event @ events::RuntimeEvent::StorageStatus { .. } => { + Some(QueuedUpdate::global(Update::Runtime(event))) + } + event => { + let session_id = activation + .as_ref() + .map(|activation| &activation.session_id) + .or(route.as_ref())? + .clone(); + Some(QueuedUpdate { + generation: None, + runtime_activation: activation, + update: Update::RoutedRuntime { session_id, event }, + }) + } + } +} + +// App-level replay tests also exercise historical unscoped runtime lines. +#[cfg(test)] +fn runtime_diagnostic_update( + route: &mut Option, + event: events::RuntimeEvent, +) -> Option { + runtime_diagnostic_envelope_update( + route, + events::DiagnosticEvent { event, - }), + activation: None, + }, + ) + .map(|queued| queued.update) +} + +fn response_activation(meta: Option<&serde_json::Map>) -> Option { + meta?.get(events::ACTIVATION_META_KEY)?.as_u64() +} + +fn verify_runtime_activation(route: &Arc>, epoch: Option) { + if let Ok(mut route) = route.lock() { + route.runtime_activation = epoch; } } @@ -257,6 +324,7 @@ fn transition_route(route: &Arc>, id: String) { if let Ok(mut route) = route.lock() { route.id = id; route.generation = route.generation.wrapping_add(1); + route.runtime_activation = None; } } @@ -264,6 +332,22 @@ fn accept_queued_update( route: &Arc>, queued: QueuedUpdate, ) -> Option { + // No marker or legacy line is activation authority. Historical child + // replay goes through ACP; StorageStatus and process updates remain global. + if let Update::RoutedRuntime { session_id, .. } = &queued.update { + let route = route.lock().ok()?; + let accepted = session_id == &route.id + && match &queued.runtime_activation { + Some(activation) => { + activation.session_id == route.id + && Some(activation.epoch) == route.runtime_activation + } + None => false, + }; + if !accepted { + return None; + } + } let accepted = queued.generation.is_none_or(|generation| { route .lock() @@ -634,6 +718,7 @@ async fn request_initial_session( connection: &agent_client_protocol::V2ConnectionTo, resume_id: Option<&str>, root: &Path, + route: &Arc>, ) -> Result<(wire::SessionId, Vec), agent_client_protocol::Error> { if let Some(resume_id) = resume_id { let response = connection @@ -643,12 +728,14 @@ async fn request_initial_session( ) .block_task() .await?; + verify_runtime_activation(route, response_activation(response.meta.as_ref())); Ok((wire::SessionId::new(resume_id), response.config_options)) } else { let response = connection .send_request(wire::NewSessionRequest::new(root.to_path_buf())) .block_task() .await?; + verify_runtime_activation(route, response_activation(response.meta.as_ref())); Ok((response.session_id, response.config_options)) } } @@ -875,6 +962,7 @@ pub async fn run_with_reasoning_effort_and_openrouter_key( let active_persisted_id = Arc::new(Mutex::new(ActiveSessionRoute { id: persisted_session_id.clone(), generation: 0, + runtime_activation: None, })); let root = root.to_path_buf(); let model = model.to_string(); @@ -929,6 +1017,8 @@ pub async fn run_with_reasoning_effort_and_openrouter_key( let (branch_tx, mut branch_rx) = mpsc::unbounded_channel::(); let branch_replay = Arc::new(Mutex::new(BranchReplayBuffer::default())); let notification_branch_replay = Arc::clone(&branch_replay); + let diagnostic_branch_replay = Arc::clone(&branch_replay); + let diagnostic_session = Arc::clone(&active_persisted_id); // The agent's own diagnostics are the only explanation of a failed start, // so they are kept aside as well as shown in the log pane. @@ -939,16 +1029,19 @@ pub async fn run_with_reasoning_effort_and_openrouter_key( let mut lines = BufReader::new(stderr).lines(); let mut runtime_route = None; while let Ok(Some(line)) = lines.next_line().await { - let update = match events::parse(&line) { + let update = match events::parse_diagnostic(&line) { Some(event) => { - let Some(update) = runtime_diagnostic_update(&mut runtime_route, event) + let Some(update) = + runtime_diagnostic_envelope_update(&mut runtime_route, event) else { continue; }; update } None if line.starts_with("A2A listening on ") => { - Update::A2aAddress(line.trim_start_matches("A2A listening on ").to_string()) + QueuedUpdate::global(Update::A2aAddress( + line.trim_start_matches("A2A listening on ").to_string(), + )) } None => { if let Ok(mut recent) = recorder.lock() { @@ -956,10 +1049,15 @@ pub async fn run_with_reasoning_effort_and_openrouter_key( let extra = recent.len().saturating_sub(FAILURE_LINES); recent.drain(..extra); } - Update::Log(line) + QueuedUpdate::global(Update::Log(line)) } }; - if diagnostics.send(QueuedUpdate::global(update)).is_err() { + let mut replay = diagnostic_branch_replay.lock().expect("branch replay lock"); + let current = diagnostic_session.lock().expect("session route lock"); + let Some(update) = replay.route_diagnostic(¤t, update) else { + continue; + }; + if diagnostics.send(update).is_err() { return; } } @@ -1090,6 +1188,7 @@ pub async fn run_with_reasoning_effort_and_openrouter_key( &connection, resume_session_id.as_deref(), &root, + &transition_session, ), &mut exit_rx, HANDSHAKE, @@ -1212,6 +1311,7 @@ pub async fn run_with_reasoning_effort_and_openrouter_key( &connection, resume_session_id.as_deref(), &root, + &transition_session, ), &mut exit_rx, stop.requested(), @@ -1327,6 +1427,9 @@ pub async fn run_with_reasoning_effort_and_openrouter_key( Ok(action) => action, Err(_) => return Ok(()), }; + if !session_action_allowed(&mut app, &action) { + continue; + } match action { Action::Quit => return Ok(()), Action::ListPromptBranches { epoch } => { @@ -1357,6 +1460,7 @@ pub async fn run_with_reasoning_effort_and_openrouter_key( let mut replay = branch_replay.lock().expect("branch replay lock"); replay.request = Some((generation, epoch)); replay.notifications.clear(); + replay.diagnostics.clear(); } let connection = connection.clone(); let session_id = session_id.clone(); @@ -1461,6 +1565,7 @@ pub async fn run_with_reasoning_effort_and_openrouter_key( .send_request(wire::NewSessionRequest::new(root.clone())) .block_task() .await?; + let runtime_activation = response_activation(session.meta.as_ref()); session_id = session.session_id; let config_options = if let Some(choice) = &saved_model_default { connection @@ -1479,6 +1584,7 @@ pub async fn run_with_reasoning_effort_and_openrouter_key( agent_client_protocol::Error::into_internal_error(std::io::Error::other(error)) })?; transition_route(&transition_session, persisted_id.clone()); + verify_runtime_activation(&transition_session, runtime_activation); images.clear(); app.start_session(persisted_id); app.activate_runtime_session(); @@ -1496,7 +1602,7 @@ pub async fn run_with_reasoning_effort_and_openrouter_key( } } } - Action::ListSessions => { + Action::ListSessions { epoch } => { let Ok(route) = transition_session.lock() else { app.note("could not start session catalog scan"); continue; @@ -1516,7 +1622,7 @@ pub async fn run_with_reasoning_effort_and_openrouter_key( .and_then(|result| result); let _ = updates.send(QueuedUpdate::for_session( generation, - Update::SessionCatalog(result), + Update::SessionCatalog { epoch, result }, )); }); } @@ -1564,22 +1670,17 @@ pub async fn run_with_reasoning_effort_and_openrouter_key( app.note(format!("invalid session id: {error}")); continue; } - let Some(previous_persisted_id) = - previous_session_for_resume(&session_id, &requested_id) - .map_err(|error| { - agent_client_protocol::Error::into_internal_error( - std::io::Error::other(error), - ) - })? - else { - app.note(format!("session {requested_id} is already active")); + if !app.session_switch_allowed(&requested_id) { continue; - }; - // Session switching has no `--force` flag. Reclaim only a - // lock that the OS proves no live Kit process still holds. + } + let previous_persisted_id = durable_session_id(&session_id) + .map_err(|error| agent_client_protocol::Error::into_internal_error(std::io::Error::other(error)))?; + let same_session = requested_id == previous_persisted_id; + // The current actor owns its live lock. Other destinations may + // reclaim only locks the OS proves no live Kit process holds. if let Err(error) = crate::session::load(&root, &requested_id) .and_then(|_| { - if preserved_branch_sources.contains(&requested_id) { + if same_session || preserved_branch_sources.contains(&requested_id) { Ok(()) } else { crate::session::remove_stale_lock(&root, &requested_id) @@ -1589,63 +1690,15 @@ pub async fn run_with_reasoning_effort_and_openrouter_key( app.note(format!("could not resume session: {error}")); continue; } - let previous_session_id = session_id.clone(); - if let Err(error) = connection - .send_request(CloseSessionRequest::new(session_id.clone())) - .block_task() - .await - { - app.note(format!( - "could not close the current session: {}", - error.message - )); - continue; - } - session_id = wire::SessionId::new(requested_id.clone()); - transition_route(&transition_session, requested_id.clone()); - images.clear(); - app.start_session(requested_id.clone()); - match request_resume(&connection, session_id.clone(), root.clone()).await { - Ok(response) => { - app.activate_runtime_session(); - refresh_config_state(&mut app, Some(&response.config_options)); - }, - Err(error) => { - session_id = previous_session_id; - transition_route( - &transition_session, - previous_persisted_id.clone(), - ); - images.clear(); - app.start_session(previous_persisted_id); - let restored = request_resume( - &connection, - session_id.clone(), - root.clone(), - ) - .await; - match restored { - Ok(response) => { - app.activate_runtime_session(); - refresh_config_state( - &mut app, - Some(&response.config_options), - ); - app.note(format!( - "could not resume {requested_id}: {}", - error.message - )); - } - Err(restore_error) => { - return Err(agent_client_protocol::Error::into_internal_error( - std::io::Error::other(format!( - "could not resume {requested_id}: {}; could not restore the previous session: {}", - error.message, restore_error.message - )), - )); - } - } - } + if resume_with_restoration( + &mut app, + &connection, + &mut session_id, + &transition_session, + &root, + requested_id, + ).await { + images.clear(); } } Action::Close => return Ok(()), @@ -1894,9 +1947,9 @@ pub async fn run_with_reasoning_effort_and_openrouter_key( let mut replay = branch_replay.lock().expect("branch replay lock"); let current = completion.is_current(&transition_session.lock().expect("session route lock"), &app); let submitted = matches!(completion.response, BranchResponse::Submitted(_)); - let buffered = if submitted { + let (buffered, buffered_diagnostics) = if submitted { replay.finish(completion.generation, completion.epoch) - } else { Vec::new() }; + } else { (Vec::new(), Vec::new()) }; if !current { continue; } match completion.response { BranchResponse::Listed(result) => app.branch_listed(completion.epoch, result.map(|response| response.boundaries)), @@ -1911,6 +1964,7 @@ pub async fn run_with_reasoning_effort_and_openrouter_key( preserved_branch_sources.insert(session_id.to_string()); session_id = response.session_id; transition_route(&transition_session, child_id.clone()); + verify_runtime_activation(&transition_session, response.runtime_activation); images.clear(); refresh_config_state(&mut app, Some(&response.config_options)); // Replay precedes activation, just as with resume. No ordinary @@ -1923,6 +1977,11 @@ pub async fn run_with_reasoning_effort_and_openrouter_key( app.apply(update); } } + for queued in buffered_diagnostics { + if let Some(update) = accept_queued_update(&transition_session, queued) { + app.apply(update); + } + } } } } @@ -1956,7 +2015,7 @@ pub async fn run_with_reasoning_effort_and_openrouter_key( .await; result?; if let Some(closed) = closed { - closed?; + recovery_close_result(closed)?; } Ok(()) }) @@ -2236,17 +2295,18 @@ fn handle(app: &mut App, received: ReceivedEvent) -> Action { if app.model_switch.is_some() { return Action::None; } - if app.navigation.dialog.is_some() + if app.session_dialog.is_some() + || app.navigation.dialog.is_some() || app.branch_chooser.is_some() || app.editing_branch() { app.paste(&text); return Action::None; } - if app.queue_focused && !app.session_rename_active() { + if app.queue_focused { return Action::None; } - if app.session_rename_active() || app.editing_steer() { + if app.editing_steer() { app.paste(&text); } else if let Some(attachments) = attachments_from_paste(&app.root, &text) { app.prune_attachments(); @@ -2544,6 +2604,100 @@ fn leave(terminal: &mut DefaultTerminal) { TERMINAL_ACTIVE.store(false, Ordering::Relaxed); } +// A failed prior recovery (or a retired actor) can leave no live map entry. +// Retrying the same durable ID must still reach load, not get stuck at close. +fn recovery_close_result( + result: Result, +) -> Result<(), agent_client_protocol::Error> { + match result { + Ok(_) => Ok(()), + Err(error) + if i32::from(error.code) == i32::from(wire::Error::resource_not_found(None).code) => + { + Ok(()) + } + Err(error) => Err(error), + } +} + +// Keep this allowlist local-only: disconnected sessions can still be inspected +// and explicitly retried, but must not dispatch work to an absent actor. +fn session_action_allowed(app: &mut App, action: &Action) -> bool { + matches!( + action, + Action::None + | Action::Redraw + | Action::ListSessions { .. } + | Action::RenameSession { .. } + | Action::Resume(_) + | Action::Close + | Action::Quit + | Action::Copy(_) + ) || app.require_session_connection() +} + +/// Close the actor, load the destination, and make at most one restoration +/// attempt. The visible session is replaced only after ACP acknowledges a load. +/// A failed restoration leaves the editor usable and Resume available for retry. +async fn resume_with_restoration( + app: &mut App, + connection: &agent_client_protocol::V2ConnectionTo, + session_id: &mut wire::SessionId, + route: &Arc>, + root: &std::path::Path, + requested_id: String, +) -> bool { + let previous_id = session_id.to_string(); + if let Err(error) = recovery_close_result( + connection + .send_request(CloseSessionRequest::new(session_id.clone())) + .block_task() + .await, + ) { + app.note(format!( + "could not close the current session: {}", + error.message + )); + return false; + } + app.session_connected = false; + transition_route(route, String::new()); + let mut load_error = None; + for target in [&requested_id, &previous_id] { + *session_id = wire::SessionId::new(target.clone()); + transition_route(route, target.clone()); + match request_resume(connection, session_id.clone(), root.to_path_buf()).await { + Ok(response) => { + verify_runtime_activation(route, response_activation(response.meta.as_ref())); + app.start_session(target.clone()); + app.activate_runtime_session(); + refresh_config_state(app, Some(&response.config_options)); + if let Some(error) = load_error { + app.note(format!( + "could not resume {requested_id}: {error}; previous session restored" + )); + } + return true; + } + Err(error) => { + if let Some(load_error) = load_error { + // Invalidate replay and completions from both failed attempts. + // Keep the durable ID for a subsequent explicit Resume retry, + // but do not advertise it as an active backend route. + transition_route(route, String::new()); + app.note(format!( + "could not resume {requested_id}: {load_error}; could not restore the previous session: {}; disconnected — use Resume to retry", + error.message + )); + return false; + } + load_error = Some(error.message); + } + } + } + unreachable!("the two load attempts return on success or restoration failure") +} + async fn request_resume( connection: &agent_client_protocol::V2ConnectionTo, session_id: wire::SessionId, @@ -2564,14 +2718,6 @@ fn durable_session_id(session_id: &wire::SessionId) -> Result { Ok(session_id) } -fn previous_session_for_resume( - current: &wire::SessionId, - requested: &str, -) -> Result, String> { - let current = durable_session_id(current)?; - Ok((current != requested).then_some(current)) -} - /// Maps one ACP session notification onto client updates. fn translate(notification: UpdateSessionNotification) -> (String, Vec) { let session_id = notification.session_id.to_string(); @@ -2939,11 +3085,10 @@ mod tests { apply_pending_updates, attachments_from_paste, authentication_required, client_capabilities, command, credential_storage_for_launch, current_model_choice, detach_from_controlling_terminal, durable_session_id, effort_state, error_detail, - message_of, osc52, previous_session_for_resume, prompt_blocks, readable, - refresh_config_state, refresh_session_after_auth, save_effort_default_to, - save_model_defaults_to, terminal_auth_command, transition_route, translate, - translate_for_session, usable_terminal_auth_methods, user_message_of, - wait_for_connected_authentication, wire, + message_of, osc52, prompt_blocks, readable, refresh_config_state, + refresh_session_after_auth, save_effort_default_to, save_model_defaults_to, + terminal_auth_command, transition_route, translate, translate_for_session, + usable_terminal_auth_methods, user_message_of, wait_for_connected_authentication, wire, }; use crate::{ tools::mcp::CredentialStorage, @@ -3328,21 +3473,6 @@ mod tests { child.wait().await.unwrap(); } - #[test] - fn resuming_the_active_session_is_a_noop() { - let current = wire::SessionId::new("current"); - assert_eq!( - previous_session_for_resume(¤t, "current").unwrap(), - None - ); - assert_eq!( - previous_session_for_resume(¤t, "other") - .unwrap() - .as_deref(), - Some("current") - ); - } - #[test] fn model_switch_request_rejects_invalid_actions_and_missing_warnings() { use crate::protocols::acp::model_switch::{Confirmation, Decision, META, Warning}; @@ -3425,6 +3555,7 @@ mod tests { let route = Arc::new(Mutex::new(ActiveSessionRoute { id: "first".into(), generation: 0, + runtime_activation: None, })); let choice = ModelChoice { id: "openrouter:target".into(), @@ -3468,6 +3599,7 @@ mod tests { let route = Arc::new(Mutex::new(ActiveSessionRoute { id: "first".into(), generation: 0, + runtime_activation: None, })); let operation = app .begin_model_switch( @@ -3493,6 +3625,7 @@ mod tests { let route = Arc::new(Mutex::new(ActiveSessionRoute { id: "first".into(), generation: 0, + runtime_activation: None, })); let old = QueuedUpdate::for_session(0, Update::Log("old".into())); @@ -3532,10 +3665,11 @@ mod tests { ) { for _ in 0..4 { let line = lines.next_line().await.unwrap().unwrap(); - if let Some(update) = - super::runtime_diagnostic_update(ingress, crate::events::parse(&line).unwrap()) - { - tx.send(QueuedUpdate::global(update)).unwrap(); + if let Some(update) = super::runtime_diagnostic_envelope_update( + ingress, + crate::events::parse_diagnostic(&line).unwrap(), + ) { + tx.send(update).unwrap(); } } } @@ -3555,6 +3689,7 @@ mod tests { let route = Arc::new(Mutex::new(ActiveSessionRoute { id: "child".into(), generation: 0, + runtime_activation: Some(1), })); let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel(); let (mut writer, reader) = tokio::io::duplex(4096); @@ -3585,8 +3720,18 @@ mod tests { ] { writer .write_all( - format!("{EVENT_MARKER}{}\n", serde_json::to_string(&event).unwrap()) - .as_bytes(), + format!( + "{EVENT_MARKER}{}\n", + serde_json::to_string(&crate::events::DiagnosticEvent { + event, + activation: Some(crate::events::DiagnosticActivation { + session_id: session.into(), + epoch: if session == "child" { 1 } else { 2 } + }), + }) + .unwrap() + ) + .as_bytes(), ) .await .unwrap(); @@ -3596,6 +3741,7 @@ mod tests { ingest(&mut lines, &mut ingress, &tx).await; } transition_route(&route, "source".into()); + super::verify_runtime_activation(&route, Some(2)); app.start_session("source".into()); app.activate_runtime_session(); app.apply(Update::ToolStarted { @@ -3634,6 +3780,651 @@ mod tests { } } + #[tokio::test] + async fn diagnostic_epochs_reject_delayed_same_id_bytes_and_queued_events() { + use crate::events::{DiagnosticActivation, DiagnosticEvent, EVENT_MARKER, RuntimeEvent}; + use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader}; + + for via_other_session in [false, true] { + for read_before_activation in [false, true] { + let route = Arc::new(Mutex::new(ActiveSessionRoute { + id: "source".into(), + generation: 0, + runtime_activation: Some(10), + })); + let mut ingress = None; + let (mut writer, reader) = tokio::io::duplex(4096); + let mut lines = BufReader::new(reader).lines(); + for (epoch, event) in [ + ( + 10, + RuntimeEvent::SessionStarted { + session_id: "source".into(), + }, + ), + ( + 10, + RuntimeEvent::CompactionStarted { + reason: "stale".into(), + at: 1, + }, + ), + ( + 10, + RuntimeEvent::StorageStatus { + pending: true, + exhausted: true, + }, + ), + ( + 12, + RuntimeEvent::ChildStarted { + call: "compose:child".into(), + tool: "shell".into(), + summary: "current".into(), + at: 2, + }, + ), + ] { + let diagnostic = DiagnosticEvent { + event, + activation: Some(DiagnosticActivation { + session_id: "source".into(), + epoch, + }), + }; + writer + .write_all( + format!( + "{EVENT_MARKER}{}\n", + serde_json::to_string(&diagnostic).unwrap() + ) + .as_bytes(), + ) + .await + .unwrap(); + } + let mut queued = Vec::new(); + if read_before_activation { + for _ in 0..4 { + let line = lines.next_line().await.unwrap().unwrap(); + if let Some(update) = super::runtime_diagnostic_envelope_update( + &mut ingress, + crate::events::parse_diagnostic(&line).unwrap(), + ) { + queued.push(update); + } + } + } + if via_other_session { + transition_route(&route, "other".into()); + super::verify_runtime_activation(&route, Some(11)); + } + transition_route(&route, "source".into()); + super::verify_runtime_activation(&route, Some(12)); + if !read_before_activation { + for _ in 0..4 { + let line = lines.next_line().await.unwrap().unwrap(); + if let Some(update) = super::runtime_diagnostic_envelope_update( + &mut ingress, + crate::events::parse_diagnostic(&line).unwrap(), + ) { + queued.push(update); + } + } + } + assert_eq!(route.lock().unwrap().runtime_activation, Some(12)); + let mut queued = queued.into_iter(); + assert!(accept_queued_update(&route, queued.next().unwrap()).is_none()); + assert!(matches!( + accept_queued_update(&route, queued.next().unwrap()), + Some(Update::Runtime(RuntimeEvent::StorageStatus { + pending: true, + exhausted: true + })) + )); + assert!(matches!( + accept_queued_update(&route, queued.next().unwrap()), + Some(Update::RoutedRuntime { + event: RuntimeEvent::ChildStarted { .. }, + .. + }) + )); + + // Neither a bare same-ID marker nor a forged/unacknowledged epoch + // can authorize legacy or future diagnostics for this activation. + for activation in [ + None, + Some(DiagnosticActivation { + session_id: "source".into(), + epoch: 13, + }), + Some(DiagnosticActivation { + session_id: "other".into(), + epoch: 12, + }), + ] { + let update = super::runtime_diagnostic_envelope_update( + &mut ingress, + DiagnosticEvent { + event: RuntimeEvent::CompactionStarted { + reason: "unverified".into(), + at: 3, + }, + activation, + }, + ) + .unwrap(); + assert!(accept_queued_update(&route, update).is_none()); + } + } + } + } + + // The peer actually retires its actor on Close and faults Resume responses. + // No production pause/observer hooks are needed to exercise the runtime path. + #[test] + fn spawned_producer_diagnostics_are_rejected_when_read_after_reactivation() { + use crate::events::{self, RuntimeEvent}; + + const CHILD: &str = "KIT_TUI_PRODUCER_TEST"; + if let Ok(mode) = std::env::var(CHILD) { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap() + .block_on(async { + events::activate_diagnostics("source"); + let (entered_tx, entered) = tokio::sync::oneshot::channel(); + let (release, released) = tokio::sync::oneshot::channel(); + // Return the task handle: awaiting it inside this scope + // would deadlock before the parent releases its barrier. + #[allow(clippy::async_yields_async)] + let producer = events::scope_diagnostics("source", async { + tokio::spawn(events::inherit_diagnostics(async move { + entered_tx.send(()).unwrap(); + released.await.unwrap(); + events::emit(&RuntimeEvent::CompactionStarted { + reason: "old producer".into(), + at: 1, + }); + })) + }) + .await; + entered.await.unwrap(); + if mode == "roundtrip" { + events::activate_diagnostics("sibling"); + } + let verified_epoch = events::activate_diagnostics("source"); + // Transfer the actual activation return value, independently + // of stderr markers, just as the ACP response establishes it. + println!("\nverified-epoch:{verified_epoch}"); + release.send(()).unwrap(); + producer.await.unwrap(); + events::scope_diagnostics("source", async { + events::emit(&RuntimeEvent::CompactionStarted { + reason: "current producer".into(), + at: 2, + }); + }) + .await; + tokio::spawn(async { + events::emit(&RuntimeEvent::CompactionStarted { + reason: "unscoped producer".into(), + at: 3, + }); + events::emit(&RuntimeEvent::StorageStatus { + pending: true, + exhausted: false, + }); + }) + .await + .unwrap(); + }); + return; + } + for mode in ["same", "roundtrip"] { + let output = std::process::Command::new(std::env::current_exe().unwrap()) + .args(["--exact", "tui::tests::spawned_producer_diagnostics_are_rejected_when_read_after_reactivation", "--nocapture"]) + .env(CHILD, mode).env(events::EVENTS_ENV, "1").output().unwrap(); + assert!( + output.status.success(), + "{}", + String::from_utf8_lossy(&output.stderr) + ); + let epoch = String::from_utf8(output.stdout) + .unwrap() + .lines() + .find_map(|line| { + line.strip_prefix("verified-epoch:") + .and_then(|value| value.parse::().ok()) + }) + .unwrap(); + let route = Arc::new(Mutex::new(ActiveSessionRoute { + id: "source".into(), + generation: 1, + runtime_activation: Some(epoch), + })); + let mut ingress = None; + let mut observed = Vec::new(); + // Only read the real emitted bytes after reactivation and producer + // completion. No hand-stamped DiagnosticEvent fixtures are used. + for diagnostic in String::from_utf8(output.stderr) + .unwrap() + .lines() + .filter_map(events::parse_diagnostic) + { + if matches!(diagnostic.event, RuntimeEvent::SessionStarted { .. }) { + super::runtime_diagnostic_envelope_update(&mut ingress, diagnostic); + continue; + } + let label = match &diagnostic.event { + RuntimeEvent::CompactionStarted { reason, .. } => { + if reason == "old producer" { + assert!(diagnostic.activation.as_ref().unwrap().epoch < epoch); + } + reason.clone() + } + RuntimeEvent::StorageStatus { .. } => "global storage".into(), + _ => continue, + }; + let queued = + super::runtime_diagnostic_envelope_update(&mut ingress, diagnostic).unwrap(); + let accepted = accept_queued_update(&route, queued).is_some(); + assert_eq!( + accepted, + label == "current producer" || label == "global storage", + "{mode}/{label}" + ); + observed.push(label); + } + assert_eq!( + observed, + [ + "old producer", + "current producer", + "unscoped producer", + "global storage" + ] + ); + } + } + + async fn same_id_recovery_over_transport(failures: usize) { + use std::collections::VecDeque; + + struct Peer { + active: bool, + outcomes: VecDeque, + requests: Vec<&'static str>, + } + let peer = Arc::new(Mutex::new(Peer { + active: true, + outcomes: std::iter::repeat_n(false, failures).chain([true]).collect(), + requests: Vec::new(), + })); + let route = Arc::new(Mutex::new(ActiveSessionRoute { + id: "source".into(), + generation: 10, + runtime_activation: Some(41), + })); + let (updates, mut replay) = tokio::sync::mpsc::unbounded_channel(); + let notification_route = Arc::clone(&route); + let agent = agent_client_protocol::Agent + .v2() + .on_receive_request( + async move |_request: wire::InitializeRequest, responder, _cx| { + responder.respond( + wire::InitializeResponse::new( + ProtocolVersion::V2, + wire::Implementation::new("faulting-peer", "0"), + ) + .capabilities(agentkit_acp::v2::agent_capabilities()), + ) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + { + let peer = Arc::clone(&peer); + async move |request: wire::CloseSessionRequest, responder, _cx| { + assert_eq!(request.session_id.to_string(), "source"); + let mut peer = peer.lock().unwrap(); + peer.requests.push("close"); + if std::mem::replace(&mut peer.active, false) { + responder.respond(wire::CloseSessionResponse::new()) + } else { + responder.respond_with_error( + agent_client_protocol::Error::resource_not_found(None), + ) + } + } + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + { + let peer = Arc::clone(&peer); + async move |request: wire::ResumeSessionRequest, responder, cx| { + assert_eq!(request.session_id.to_string(), "source"); + assert!(matches!( + request.replay_from, + Some(wire::ReplayFrom::Start(_)) + )); + let mut peer = peer.lock().unwrap(); + peer.requests.push("resume"); + assert!(!peer.active, "close must retire the previous actor first"); + let succeeds = peer + .outcomes + .pop_front() + .expect("unexpected automatic retry"); + cx.send_notification(UpdateSessionNotification::new( + "source", + SessionUpdate::UserMessage(UserMessage::new("replayed").content(vec![ + ContentBlock::Text(TextContent::new(if succeeds { + "loaded replay" + } else { + "failed replay" + })), + ])), + ))?; + if succeeds { + peer.active = true; + responder.respond( + wire::ResumeSessionResponse::new() + .config_options(vec![SessionConfigOption::select( + "reasoning_effort", + "Reasoning effort", + "high", + vec![SessionConfigSelectGroup::new( + "reasoning-effort", + "Reasoning effort", + vec![SessionConfigSelectOption::new("high", "High")], + )], + )]) + .meta(serde_json::Map::from_iter([( + super::events::ACTIVATION_META_KEY.into(), + json!(42), + )])), + ) + } else { + responder.respond_with_error( + agent_client_protocol::Error::into_internal_error( + std::io::Error::other("injected resume failure"), + ), + ) + } + } + }, + agent_client_protocol::on_receive_request!(), + ); + let (client_transport, agent_transport) = Channel::duplex(); + let server = tokio::spawn(async move { agent.connect_to(agent_transport).await }); + let root = tempfile::tempdir().unwrap(); + let mut app = App::new( + root.path().into(), + "provider".into(), + "model".into(), + "a2a".into(), + ); + app.start_session("source".into()); + app.activate_runtime_session(); + app.apply(Update::UserMessage { + id: "visible".into(), + text: "keep the visible transcript until load succeeds".into(), + images: vec![], + append: false, + }); + app.paste("parked draft [Image #1]"); + app.attachments.push(crate::tui::app::Attachment { + path: root.path().join("draft.png"), + placeholder: "[Image #1]".into(), + mime_type: "image/png", + kind: crate::tui::app::AttachmentKind::Image, + size: 7, + }); + let draft = app.editor.text().to_owned(); + let attachments = app.attachments.clone(); + let workspace = root.path().to_path_buf(); + agent_client_protocol::Client + .v2() + .on_receive_notification( + async move |notification: UpdateSessionNotification, _cx| { + let route = notification_route.lock().unwrap().clone(); + for update in translate_for_session(notification, &route.id) { + updates + .send(QueuedUpdate::for_session(route.generation, update)) + .unwrap(); + } + Ok(()) + }, + agent_client_protocol::on_receive_notification!(), + ) + .connect_with(client_transport, async move |connection| { + connection + .send_request(wire::InitializeRequest::new( + ProtocolVersion::V2, + wire::Implementation::new("recovery-test", "0"), + )) + .block_task() + .await?; + let mut session_id = wire::SessionId::new("source"); + assert!(app.session_switch_allowed("source")); + let restored = super::resume_with_restoration( + &mut app, + &connection, + &mut session_id, + &route, + &workspace, + "source".into(), + ) + .await; + assert_eq!(restored, failures == 1); + assert_eq!(peer.lock().unwrap().requests, ["close", "resume", "resume"]); + assert_eq!(app.editor.text(), draft); + assert_eq!(app.attachments, attachments); + assert_eq!(app.session_id.as_deref(), Some("source")); + assert_eq!(session_id.to_string(), "source"); + for _ in 0..2 { + let update = replay.recv().await.unwrap(); + if let Some(update) = accept_queued_update(&route, update) { + app.apply(update); + } + } + if failures == 2 { + assert!(!app.session_connected); + assert!(!peer.lock().unwrap().active); + assert!( + matches!(&app.blocks[0], crate::tui::app::Block::User(message) + if message.text == "keep the visible transcript until load succeeds") + ); + assert!( + matches!(app.blocks.last(), Some(crate::tui::app::Block::Notice(message)) + if message.contains("disconnected") && message.contains("Resume to retry")) + ); + assert_eq!( + app.blocks.len(), + 2, + "failed replay must not replace the view" + ); + assert!(route.lock().unwrap().id.is_empty()); + assert_eq!(route.lock().unwrap().runtime_activation, None); + assert!(!super::session_action_allowed(&mut app, &Action::Cancel)); + assert!(super::session_action_allowed( + &mut app, + &Action::Copy("visible".into()) + )); + assert!(super::session_action_allowed( + &mut app, + &Action::ListSessions { epoch: 0 } + )); + assert!(super::session_action_allowed( + &mut app, + &Action::Resume("source".into()) + )); + assert!(app.session_switch_allowed("source")); + assert_eq!(app.editor.text(), draft); + assert_eq!(app.attachments, attachments); + // Explicit retry reaches Resume even though Close now reports + // that the already-retired actor is absent. + assert!( + super::resume_with_restoration( + &mut app, + &connection, + &mut session_id, + &route, + &workspace, + "source".into(), + ) + .await + ); + assert_eq!( + peer.lock().unwrap().requests, + ["close", "resume", "resume", "close", "resume"] + ); + let update = replay.recv().await.unwrap(); + app.apply( + accept_queued_update(&route, update).expect("retry replay is current"), + ); + } + assert!(app.session_connected); + assert!(peer.lock().unwrap().active); + assert!(peer.lock().unwrap().outcomes.is_empty()); + let users = app + .blocks + .iter() + .filter_map(|block| match block { + crate::tui::app::Block::User(message) => Some(message.text.as_str()), + _ => None, + }) + .collect::>(); + assert_eq!(users, ["loaded replay"]); + assert_eq!(route.lock().unwrap().id, "source"); + assert_eq!(route.lock().unwrap().runtime_activation, Some(42)); + assert_eq!(app.reasoning_effort, "high"); + assert_eq!(app.editor.text(), draft); + assert_eq!(app.attachments, attachments); + assert!(super::session_action_allowed(&mut app, &Action::Cancel)); + Ok(()) + }) + .await + .unwrap(); + server.abort(); + let _ = server.await; + } + + #[tokio::test] + async fn same_id_recovery_restores_after_first_resume_failure_over_transport() { + tokio::time::timeout( + std::time::Duration::from_secs(10), + same_id_recovery_over_transport(1), + ) + .await + .unwrap(); + } + + #[tokio::test] + async fn same_id_recovery_stays_disconnected_then_explicit_retry_succeeds_over_transport() { + tokio::time::timeout( + std::time::Duration::from_secs(10), + same_id_recovery_over_transport(2), + ) + .await + .unwrap(); + } + + #[test] + fn same_id_recovery_disconnected_gate_allows_local_actions_only() { + let mut app = App::new(".".into(), "provider".into(), "model".into(), "a2a".into()); + app.start_session("source".into()); + app.paste("untouched draft"); + app.session_connected = false; + let backend_actions = [ + Action::ListPromptBranches { epoch: 0 }, + Action::PreparePromptBranch { + epoch: 0, + address: "prompt".into(), + }, + Action::SubmitPromptBranch { + epoch: 0, + checkout_token: "token".into(), + text: "edit".into(), + }, + Action::Submit { + prompt: SubmittedPrompt { + text: "send".into(), + attachments: vec![], + }, + inject: false, + }, + Action::ReplaceSteer { + id: "steer".into(), + text: "replace".into(), + }, + Action::RevokeSteer { id: "steer".into() }, + Action::New(None), + Action::SelectModel { + choice: ModelChoice { + id: "provider:model".into(), + provider: "provider".into(), + model: "model".into(), + }, + save_defaults: false, + }, + Action::SelectEffort { + effort: "high".into(), + save_defaults: false, + }, + Action::Login(AuthMethodTerminal::new("provider", "Provider")), + Action::Cancel, + Action::DetachCompose("call".into()), + Action::CancelBackground("call".into()), + Action::SearchFiles { + query: "file".into(), + revision: 0, + activation: 0, + }, + ]; + for action in &backend_actions { + assert!(!super::session_action_allowed(&mut app, action)); + } + for action in [ + Action::None, + Action::Redraw, + Action::ListSessions { epoch: 0 }, + Action::RenameSession { + session_id: "source".into(), + display_name: Some("renamed".into()), + }, + Action::Resume("source".into()), + Action::Close, + Action::Quit, + Action::Copy("visible".into()), + ] { + assert!(super::session_action_allowed(&mut app, &action)); + } + assert_eq!(app.editor.text(), "untouched draft"); + app.session_connected = true; + for action in &backend_actions { + assert!(super::session_action_allowed(&mut app, action)); + } + } + + #[test] + fn same_id_recovery_can_retry_after_actor_is_already_absent() { + assert!(super::recovery_close_result(Ok(wire::CloseSessionResponse::new())).is_ok()); + assert!( + super::recovery_close_result(Err(agent_client_protocol::Error::resource_not_found( + None + ))) + .is_ok() + ); + assert!( + super::recovery_close_result(Err(agent_client_protocol::Error::internal_error())) + .is_err() + ); + } + #[test] fn queued_updates_are_applied_in_bounded_bursts() { let root = tempfile::tempdir().unwrap(); @@ -3646,6 +4437,7 @@ mod tests { let route = Arc::new(Mutex::new(ActiveSessionRoute { id: "session".into(), generation: 0, + runtime_activation: None, })); let (updates_tx, mut updates_rx) = tokio::sync::mpsc::unbounded_channel(); for index in 0..=MAX_BURST { @@ -3673,6 +4465,7 @@ mod tests { let route = Arc::new(Mutex::new(ActiveSessionRoute { id: "session".into(), generation: 0, + runtime_activation: None, })); let (_updates_tx, mut updates_rx) = tokio::sync::mpsc::unbounded_channel(); let (exit_tx, mut exit_rx) = tokio::sync::oneshot::channel(); @@ -3710,6 +4503,7 @@ mod tests { let route = Arc::new(Mutex::new(ActiveSessionRoute { id: "session".into(), generation: 0, + runtime_activation: None, })); let (_updates_tx, mut updates_rx) = tokio::sync::mpsc::unbounded_channel(); let (_exit_tx, mut exit_rx) = tokio::sync::oneshot::channel(); @@ -3762,6 +4556,7 @@ mod tests { let route = super::ActiveSessionRoute { id: "source".into(), generation: 7, + runtime_activation: None, }; for response in [ super::BranchResponse::Listed(Err("list".into())), @@ -3786,6 +4581,7 @@ mod tests { fn prompt_branch_replay_waits_for_child_route_and_stale_submit_does_not_drain_new_buffer() { let mut buffer = super::BranchReplayBuffer { request: Some((7, 3)), + diagnostics: Vec::new(), notifications: vec![UpdateSessionNotification::new( "child", SessionUpdate::UserMessage( @@ -3794,13 +4590,14 @@ mod tests { ), )], }; - assert!(buffer.finish(7, 2).is_empty()); + assert!(buffer.finish(7, 2).0.is_empty()); assert_eq!(buffer.notifications.len(), 1); let route = Arc::new(Mutex::new(super::ActiveSessionRoute { id: "source".into(), generation: 7, + runtime_activation: None, })); - let notifications = buffer.finish(7, 3); + let (notifications, _) = buffer.finish(7, 3); assert!(buffer.request.is_none()); transition_route(&route, "child".into()); let active = route.lock().unwrap(); @@ -3814,6 +4611,71 @@ mod tests { assert_eq!(active.generation, 8); } + #[test] + fn branch_diagnostics_wait_for_verified_child_epoch_without_blocking_storage() { + use crate::events::{DiagnosticActivation, DiagnosticEvent, RuntimeEvent}; + let route = Arc::new(Mutex::new(ActiveSessionRoute { + id: "source".into(), + generation: 7, + runtime_activation: Some(40), + })); + let mut replay = super::BranchReplayBuffer { + request: Some((7, 3)), + ..Default::default() + }; + for epoch in [41, 42] { + let update = super::runtime_diagnostic_envelope_update( + &mut None, + DiagnosticEvent { + event: RuntimeEvent::ChildStarted { + call: format!("compose:{epoch}"), + tool: "shell".into(), + summary: "early".into(), + at: 1, + }, + activation: Some(DiagnosticActivation { + session_id: "child".into(), + epoch, + }), + }, + ) + .unwrap(); + assert!( + replay + .route_diagnostic(&route.lock().unwrap(), update) + .is_none() + ); + } + let storage = super::runtime_diagnostic_envelope_update( + &mut None, + DiagnosticEvent { + event: RuntimeEvent::StorageStatus { + pending: true, + exhausted: false, + }, + activation: None, + }, + ) + .unwrap(); + assert!( + replay + .route_diagnostic(&route.lock().unwrap(), storage) + .is_some() + ); + assert!(replay.finish(7, 2).1.is_empty()); + assert_eq!(replay.diagnostics.len(), 2); + let (_, diagnostics) = replay.finish(7, 3); + transition_route(&route, "child".into()); + super::verify_runtime_activation(&route, Some(42)); + let accepted: Vec<_> = diagnostics + .into_iter() + .filter_map(|queued| accept_queued_update(&route, queued)) + .collect(); + assert!( + matches!(accepted.as_slice(), [Update::RoutedRuntime { event: RuntimeEvent::ChildStarted { call, .. }, .. }] if call == "compose:42") + ); + } + #[test] fn translates_accepted_user_messages_and_foreground_state() { let user = UpdateSessionNotification::new( @@ -4105,6 +4967,7 @@ mod tests { let route = Arc::new(Mutex::new(super::ActiveSessionRoute { id: "session".into(), generation: 1, + runtime_activation: None, })); let (updates, mut receiver) = tokio::sync::mpsc::unbounded_channel(); let (release, waiting) = tokio::sync::oneshot::channel(); @@ -4147,6 +5010,7 @@ mod tests { let route = Arc::new(Mutex::new(super::ActiveSessionRoute { id: "session".into(), generation: 1, + runtime_activation: None, })); let (updates, mut receiver) = tokio::sync::mpsc::unbounded_channel(); let (release, waiting) = tokio::sync::oneshot::channel(); @@ -4593,14 +5457,6 @@ mod tests { "model".into(), "a2a".into(), ); - app.paste("/sessions"); - assert!(matches!( - handle( - &mut app, - Event::Key(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)) - ), - Action::ListSessions - )); app.paste("parked draft"); if populated { app.apply(Update::SteerAccepted { @@ -4614,15 +5470,26 @@ mod tests { Event::Key(KeyEvent::new(KeyCode::F(2), KeyModifiers::NONE)), ); assert_eq!(app.queue_focused, populated); - app.apply(Update::SessionCatalog(Ok(vec![ - crate::session::CatalogEntry { + assert!(matches!( + handle( + &mut app, + Event::Key(KeyEvent::new(KeyCode::F(4), KeyModifiers::NONE)) + ), + Action::ListSessions { .. } + )); + app.apply(Update::SessionCatalog { + epoch: app.session_catalog_epoch, + result: Ok(vec![crate::session::CatalogEntry { id: "saved".into(), title: Some("Saved".into()), preview: None, is_subagent: false, + lineage: crate::session::CatalogLineage::Root, + branch_point: None, updated_at: 0, - }, - ]))); + }]), + }); + std::thread::sleep(std::time::Duration::from_millis(20)); handle( &mut app, Event::Key(KeyEvent::new(KeyCode::Char('r'), KeyModifiers::NONE)), @@ -4639,6 +5506,67 @@ mod tests { } } + #[test] + fn loading_session_dialog_owns_media_paste_and_enter_over_parked_queue() { + use crossterm::event::{KeyCode, KeyEvent, KeyModifiers}; + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("image.png"); + std::fs::write(&path, b"png").unwrap(); + let mut app = App::new( + directory.path().into(), + "provider".into(), + "model".into(), + "a2a".into(), + ); + app.paste("parked draft "); + handle(&mut app, Event::Paste(path.display().to_string())); + app.apply(Update::SteerAccepted { + id: "pending".into(), + text: "queued text".into(), + editable: true, + }); + handle( + &mut app, + Event::Key(KeyEvent::new(KeyCode::F(2), KeyModifiers::NONE)), + ); + let draft = app.editor.text().to_owned(); + let attachments = app.attachments.clone(); + let Action::ListSessions { epoch } = handle( + &mut app, + Event::Key(KeyEvent::new(KeyCode::F(4), KeyModifiers::NONE)), + ) else { + panic!("catalog requested") + }; + handle(&mut app, Event::Paste(path.display().to_string())); + std::thread::sleep(std::time::Duration::from_millis(20)); + assert!(matches!( + handle( + &mut app, + Event::Key(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)) + ), + Action::None + )); + assert_eq!( + app.session_dialog.as_ref().unwrap().query, + path.to_str().unwrap() + ); + assert_eq!(app.editor.text(), draft); + assert_eq!(app.attachments, attachments); + assert!(app.queue_focused); + handle( + &mut app, + Event::Key(KeyEvent::new(KeyCode::Esc, KeyModifiers::NONE)), + ); + app.apply(Update::SessionCatalog { + epoch, + result: Ok(Vec::new()), + }); + assert!(app.session_dialog.is_none()); + assert_eq!(app.editor.text(), draft); + assert_eq!(app.attachments, attachments); + assert!(app.queue_focused); + } + #[test] fn session_rename_paste_is_not_interpreted_as_an_attachment() { let directory = tempfile::tempdir().unwrap(); @@ -4651,6 +5579,8 @@ mod tests { "a2a".into(), ); app.session_dialog = Some(SessionDialog { + query: String::new(), + searching: false, selected: 0, rename: Some(SessionRename::Editing(String::new())), }); @@ -4664,6 +5594,36 @@ mod tests { )); } + #[test] + fn session_browsing_paste_precedes_queue_media_and_parked_editor() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("image.png"); + std::fs::write(&path, b"png").unwrap(); + for queue_focused in [false, true] { + let mut app = App::new( + directory.path().into(), + "provider".into(), + "model".into(), + "a2a".into(), + ); + app.paste("parked draft"); + app.queue_focused = queue_focused; + app.session_dialog = Some(SessionDialog { + query: String::new(), + searching: false, + selected: 0, + rename: None, + }); + handle(&mut app, Event::Paste(path.display().to_string())); + let dialog = app.session_dialog.as_ref().unwrap(); + assert!(dialog.searching); + assert_eq!(dialog.query, path.to_str().unwrap()); + assert_eq!(app.editor.text(), "parked draft"); + assert!(app.attachments.is_empty()); + assert_eq!(app.queue_focused, queue_focused); + } + } + #[test] fn session_rename_confirmation_and_saving_consume_paste() { let directory = tempfile::tempdir().unwrap(); @@ -4678,6 +5638,8 @@ mod tests { "a2a".into(), ); app.session_dialog = Some(SessionDialog { + query: String::new(), + searching: false, selected: 0, rename: Some(rename), }); diff --git a/src/tui/ui.rs b/src/tui/ui.rs index 8864ba8..4f1fd6f 100644 --- a/src/tui/ui.rs +++ b/src/tui/ui.rs @@ -70,6 +70,7 @@ pub fn draw(frame: &mut Frame<'_>, app: &mut App, images: &mut ImageRuntime) { .min(available_start_prompt_rows); let show_start = frame.area().width >= 20 && available_start_prompt_rows >= START_MIN_PROMPT_ROWS + && app.session_connected && app.blocks.is_empty() && app.pending_steers.is_empty() && !app.editing_steer() @@ -135,8 +136,6 @@ pub fn draw(frame: &mut Frame<'_>, app: &mut App, images: &mut ImageRuntime) { draw_model_switch_dialog(frame, pending); } else if app.file_picker.is_some() { draw_file_picker(frame, app, prompt_area, prompt_viewport, picker_below); - } else if app.session_dialog.is_some() { - draw_session_dialog(frame, app); } else if app.model_dialog.is_some() { draw_model_dialog(frame, app); } else if app.effort_dialog.is_some() { @@ -148,6 +147,9 @@ pub fn draw(frame: &mut Frame<'_>, app: &mut App, images: &mut ImageRuntime) { if app.branch_chooser.is_some() { draw_branch_chooser(frame, app); } + if app.session_dialog.is_some() { + draw_session_dialog(frame, app); + } // Durability stays visible on the start screen and over session pickers. // Pending data belongs to the process, not the currently selected session. if app.storage_pending || app.storage_exhausted { @@ -458,11 +460,13 @@ fn draw_session_dialog(frame: &mut Frame<'_>, app: &App) { .session_dialog .as_ref() .and_then(|dialog| dialog.rename.as_ref()); - let prompt_rows = u16::from(rename.is_some()); + let dialog = app.session_dialog.as_ref().expect("session dialog is open"); + let prompt_rows = u16::from(rename.is_some() || dialog.searching || !dialog.query.is_empty()); let height = outer.height.min( - (app.session_choices.len() as u16) - .saturating_add(3 + prompt_rows) - .min(23), + u16::try_from(app.session_matches.len().max(1)) + .unwrap_or(u16::MAX) + .saturating_add(5 + prompt_rows) + .min(25), ); let area = Rect::new( outer.x + outer.width.saturating_sub(width) / 2, @@ -470,15 +474,22 @@ fn draw_session_dialog(frame: &mut Frame<'_>, app: &App) { width, height, ); - let selected = app - .session_dialog - .as_ref() - .map_or(0, |dialog| dialog.selected); - let panel = Panel::bordered().title(" sessions "); + let selected = dialog.selected; + let panel = Panel::bordered().title(if app.session_connected { + " branches · sessions " + } else { + " branches · disconnected " + }); let inner = panel.inner(area); let footer_rows = u16::from(inner.height > prompt_rows.saturating_add(1)); - let [list, prompt, footer] = Layout::vertical([ + let detail_rows = if inner.height > prompt_rows + footer_rows + 2 { + 2 + } else { + 0 + }; + let [list, detail, prompt, footer] = Layout::vertical([ Constraint::Min(0), + Constraint::Length(detail_rows), Constraint::Length(prompt_rows), Constraint::Length(footer_rows), ]) @@ -486,36 +497,51 @@ fn draw_session_dialog(frame: &mut Frame<'_>, app: &App) { let visible = list.height as usize; let start = selected .saturating_sub(visible / 2) - .min(app.session_choices.len().saturating_sub(visible)); - let lines = app.session_choices[start..] + .min(app.session_matches.len().saturating_sub(visible)); + let lines = app.session_matches[start..] .iter() .take(visible) .enumerate() - .map(|(offset, entry)| { + .map(|(offset, matched)| { + let row = &app.session_forest.rows[matched.index]; + let entry = &row.entry; let is_selected = start + offset == selected; - let label = entry - .title - .as_deref() - .or(entry.preview.as_deref()) - .unwrap_or("untitled"); - let detail = entry - .preview - .as_deref() - .filter(|preview| *preview != label) - .map_or_else( - || label.to_string(), - |preview| format!("{label} — {preview}"), - ); - let updated = entry.updated_at_rfc3339(); - let text = format!( - "{}{} · {} · {}", - if is_selected { "› " } else { " " }, - &updated[..10], - entry.id, - detail + let current = app.session_id.as_deref() == Some(entry.id.as_str()); + let indent = row + .display_depth() + .min((list.width as usize).saturating_sub(24) / 2); + let short_id: String = entry.id.chars().take(12).collect(); + let title = truncate_to_width(entry.title.as_deref().unwrap_or("untitled"), 24); + let preview = truncate_to_width(entry.preview.as_deref().unwrap_or(""), 48); + let warning = row.warning.as_deref().map_or("", |warning| { + if warning.contains("orphan") { + " [orphan]" + } else if warning.to_lowercase().contains("cycle") { + " [cycle]" + } else if warning.contains("incomplete") { + " [incomplete]" + } else { + " [warning]" + } + }); + let label = format!( + "{}{}{}{}{}{}{} · {} · {}", + if is_selected { "›" } else { " " }, + if current { "●" } else { " " }, + " ".repeat(indent), + if row.parent.is_some() { "↳ " } else { "" }, + short_id, + warning, + if matched.direct_match { + "" + } else { + " [ancestor]" + }, + title, + preview ); Line::from(Span::styled( - truncate_to_width(&text, list.width as usize), + truncate_to_width(&label, list.width as usize), if is_selected { theme::accent() } else { @@ -526,18 +552,88 @@ fn draw_session_dialog(frame: &mut Frame<'_>, app: &App) { .collect::>(); frame.render_widget(Clear, area); frame.render_widget(panel, area); - frame.render_widget(Paragraph::new(lines), list); - + if app.session_catalog_pending { + frame.render_widget(Paragraph::new("Loading sessions… type to search"), list); + } else if lines.is_empty() { + frame.render_widget(Paragraph::new("No matching sessions"), list); + } else { + frame.render_widget(Paragraph::new(lines), list); + } + if let Some(matched) = app.session_matches.get(selected) { + let row = &app.session_forest.rows[matched.index]; + let point = row + .entry + .branch_point + .clone() + .unwrap_or_else(|| match &row.entry.lineage { + crate::session::CatalogLineage::Branch(metadata) => format!( + "parent {} · persisted prefix {} · {}", + metadata.parent_session_id, + metadata.boundary.prefix_len, + metadata + .boundary + .prefix_hash + .chars() + .take(12) + .collect::() + ), + crate::session::CatalogLineage::Root => { + "Independent session (no branch lineage)".into() + } + crate::session::CatalogLineage::Warning(warning) => warning.clone(), + }); + let point = format!("Branch point: {point}"); + let status = row + .warning + .clone() + .or_else(|| row.depth_note()) + .unwrap_or_else(|| { + if let crate::session::CatalogLineage::Branch(metadata) = &row.entry.lineage { + return format!( + "Parent {} · persisted prefix {} · {}", + metadata.parent_session_id, + metadata.boundary.prefix_len, + metadata + .boundary + .prefix_hash + .chars() + .take(12) + .collect::() + ); + } + if app.session_id.as_deref() == Some(row.entry.id.as_str()) { + "● current session · Enter closes/reloads committed history".into() + } else { + "● marks current session; › marks selection".into() + } + }); + frame.render_widget( + Paragraph::new(vec![ + Line::from(truncate_to_width(&point, detail.width as usize)), + Line::from(truncate_to_width(&status, detail.width as usize)), + ]) + .style(theme::dim()), + detail, + ); + } let footer_text = match rename { Some(SessionRename::Editing(_)) => "enter save · esc cancel", Some(SessionRename::ConfirmClear) => "enter clear name · esc cancel", Some(SessionRename::Saving) => "saving…", - None => "↑/↓ select · enter resume · r rename · esc close", + None if app.session_catalog_pending => { + "type search · esc/F4 close · activation waits for catalog" + } + None if dialog.searching => "type search · enter/esc finish search", + None => "↑/↓ select · enter resume · / search · r rename · esc/F4 close", }; frame.render_widget( Paragraph::new(Span::styled(footer_text, theme::dim())), footer, ); + if rename.is_none() && prompt_rows > 0 { + let value = visible_query_tail(&dialog.query, (prompt.width as usize).saturating_sub(9)); + frame.render_widget(Paragraph::new(format!("search: {value}")), prompt); + } match rename { Some(SessionRename::Editing(input)) if prompt.width > 0 => { @@ -869,7 +965,7 @@ fn draw_navigation(frame: &mut Frame<'_>, app: &App, matches: &[usize]) { frame.render_widget( Paragraph::new(vec![ Line::from(Span::styled( - "↑/↓ · Enter reveal · /branch ↵ edit prompt · Esc/F3 close", + "↑/↓ · Enter reveal · F2 branch chooser · Esc/F3 close", theme::dim(), )), Line::from(Span::styled( @@ -2527,6 +2623,14 @@ fn prompt_lines( } fn draw_status(frame: &mut Frame<'_>, app: &App, area: Rect) { + if !app.session_connected { + frame.render_widget( + Paragraph::new(" disconnected · F4: select current session and Enter to retry") + .style(Style::default().fg(theme::warn_color())), + area, + ); + return; + } let mut left = match app.phase { Phase::Idle => vec![Span::styled(" ready", theme::dim())], Phase::Cancelling => vec![ @@ -2579,7 +2683,7 @@ fn draw_status(frame: &mut Frame<'_>, app: &App, area: Rect) { } else if !app.pending_steers.is_empty() { "F2 queue ⏎ send ⇧⏎ newline " } else { - "F3 transcript ⏎ send ⇧⏎ newline ^l log ^c quit " + "F3 transcript F4 branches ⏎ send ⇧⏎ newline ^l log ^c quit " }; let used: usize = left.iter().map(|span| span.content.chars().count()).sum(); let gap = (area.width as usize) @@ -3886,6 +3990,72 @@ mod tests { assert!(!dismissed.contains(" commands "), "{dismissed}"); } + #[test] + fn explorer_renders_loading_and_keeps_query_visible_before_catalog_arrives() { + let mut app = panel_app(0); + app.paste("parked draft"); + app.handle_key(KeyEvent::new(KeyCode::F(4), KeyModifiers::NONE)); + app.paste("searching now"); + let output = render(&mut app, 100, 18); + assert!(output.contains("Loading sessions"), "{output}"); + assert!(output.contains("search: searching now"), "{output}"); + assert!(!output.contains("No matching sessions"), "{output}"); + assert_eq!(app.editor.text(), "parked draft"); + } + + #[test] + fn disconnected_status_never_claims_ready_even_with_empty_history() { + let mut app = panel_app(0); + app.session_connected = false; + let output = render(&mut app, 100, 18); + assert!(output.contains("disconnected"), "{output}"); + assert!(output.contains("Enter to retry"), "{output}"); + assert!(!output.contains(" ready"), "{output}"); + } + + #[test] + fn branch_explorer_renders_current_separately_from_selection_and_bounded_search() { + let mut app = panel_app(0); + app.session_id = Some("source".into()); + app.set_session_choices(vec![ + crate::session::CatalogEntry { + id: "source".into(), + title: Some("Current source".into()), + preview: Some("latest source".into()), + is_subagent: false, + lineage: crate::session::CatalogLineage::Root, + branch_point: None, + updated_at: 1, + }, + crate::session::CatalogEntry { + id: "orphan".into(), + title: Some("Orphan name".into()), + preview: Some("latest orphan".into()), + is_subagent: false, + lineage: crate::session::CatalogLineage::Warning( + "Missing parent; shown as an orphan root".into(), + ), + branch_point: None, + updated_at: 2, + }, + ]); + app.session_dialog = Some(SessionDialog { + query: String::new(), + searching: false, + selected: 0, + rename: None, + }); + let output = render(&mut app, 100, 18); + assert!(output.contains("› orphan [orphan]"), "{output}"); + assert!(output.contains("●source"), "{output}"); + assert!(output.contains("latest orphan"), "{output}"); + assert!(output.contains("Missing parent"), "{output}"); + app.paste(&"界👩‍💻".repeat(3000)); + let output = render(&mut app, 32, 8); + assert!(output.contains("No matching sessions"), "{output}"); + assert_eq!(app.editor.text(), ""); + } + #[test] fn session_dialog_renders_custom_names_and_inline_rename_states() { let mut app = App::new( @@ -3894,14 +4064,18 @@ mod tests { "model".into(), "127.0.0.1:7331".into(), ); - app.session_choices = vec![crate::session::CatalogEntry { + app.set_session_choices(vec![crate::session::CatalogEntry { id: "s-abc123".into(), title: Some("OAuth token bug".into()), preview: Some("Preview remains available".into()), is_subagent: false, + lineage: crate::session::CatalogLineage::Root, + branch_point: None, updated_at: 0, - }]; + }]); app.session_dialog = Some(SessionDialog { + query: String::new(), + searching: false, selected: 0, rename: None, }); @@ -3931,7 +4105,7 @@ mod tests { ); assert!(confirming.contains("enter clear name"), "{confirming}"); - app.session_choices[0].title = Some("界".repeat(100)); + app.session_forest.rows[0].entry.title = Some("界".repeat(100)); app.session_dialog.as_mut().unwrap().rename = None; let narrow = render(&mut app, 20, 6); assert!(narrow.contains('…'), "{narrow}"); From b2d80778a28fec8c992830102a037bb28b967754 Mon Sep 17 00:00:00 2001 From: daniel Date: Sun, 6 Sep 2026 00:25:56 +0100 Subject: [PATCH 2/7] feat(diagnostics): correlate nested child operation ownership --- docs/user/subagents-and-acp-harnesses.md | 8 + docs/user/tui-and-sessions.md | 4 +- src/acp_child.rs | 816 ++++++++++- src/events.rs | 425 +++++- src/protocols/acp.rs | 969 ++++++++---- src/protocols/acp/v2.rs | 1697 +++++++++++++++++----- src/runtime.rs | 17 +- src/runtime/diagnostics.rs | 277 +++- src/tools/subagent.rs | 9 +- src/tools/subagent/tests.rs | 14 +- src/tui/mod.rs | 10 +- 11 files changed, 3502 insertions(+), 744 deletions(-) diff --git a/docs/user/subagents-and-acp-harnesses.md b/docs/user/subagents-and-acp-harnesses.md index 2929218..8a2bed6 100644 --- a/docs/user/subagents-and-acp-harnesses.md +++ b/docs/user/subagents-and-acp-harnesses.md @@ -140,6 +140,14 @@ Built-in subagent transcripts are durable on disk, but their reusable parent-own Generic external child harnesses remain ACP v1: they must speak newline-delimited JSON-RPC over stdio and support `initialize`, `session/new`, and `session/prompt`. `session/fork` and `session/close` are optional capabilities. Keep stdout protocol-only; the agent may log to stderr. Kit runs the executable directly from the subagent's selected working directory, which defaults to Kit's working directory, and inherits the parent environment. It does not invoke a shell, so pipes, environment assignments, compound commands, and shell quoting in `command` or `args` do not work. +### Private nested diagnostic ownership + +Updated Kit parents and children negotiate `_meta.kitDiagnosticOwnership` with the exact value `{"version":1,"transport":"stderr"}` during initialization. The child must acknowledge support explicitly. Each invocation then carries `_meta.kitDiagnosticOperation`, an opaque connection-local token. The child captures it at admission and echoes it as the optional `operation` field of owned runtime diagnostic envelopes on stderr. Tokens are nonempty ASCII letters, digits, or `-_.:`, at most 128 bytes; they are ephemeral and are not saved in transcripts or fork context. Existing metadata keys are preserved. + +The parent resolves tokens against immutable caller scopes registered before requests are sent, retained until process exit and stderr drain. Child-local activation IDs remain separate. Each grandchild boundary resolves its own tokens and restores upstream ownership. In negotiated mode, missing, unknown, or malformed tokens cannot acquire launch or current-session attribution. Detached producer ownership survives responses and cancellation. Background continuations recover origins from existing tool-call IDs when their results actually enter the loop. An unrelated live task does not suppress a new prompt. If a continuation actually combines different origins, it has no single attributed owner; previously captured producers retain their individual ownership. + +The extension is optional for external agents. Missing or unknown negotiation versions keep legacy ACP operations usable but cannot supply absent diagnostic identity. See [live diagnostic compatibility](tui-and-sessions.md) for the older-agent boundary after session reactivation. + Configure trusted argv profiles in `~/.kit/config.toml`: ```toml diff --git a/docs/user/tui-and-sessions.md b/docs/user/tui-and-sessions.md index 36bf977..1e9dc02 100644 --- a/docs/user/tui-and-sessions.md +++ b/docs/user/tui-and-sessions.md @@ -103,7 +103,9 @@ Each durable workspace session appears once, using its existing name, short ID, Enter resumes the selected visible session only while **idle**, outside provisional checkout, submission, pending-message edits, or conflicting queued work. Browsing is still available while working. A different session cannot replace an unsent draft or attachments: send or clear them first. Selecting the current session explicitly closes and reloads committed history while retaining its unsent editor draft and attachments. Reload never reconstructs or submits cancelled input. If loading fails after the current actor closes, Kit makes one restoration attempt. If that also fails, the TUI shows **disconnected**, retains the draft and attachments, and blocks backend-dependent actions. Open F4 and select the current session again to retry; only a successful load reconnects it. -Live diagnostic compatibility: after a session reload or switch back, Kit rejects delayed diagnostics from earlier activations. A reused legacy ACP v1 child shares one stderr stream without per-operation identity, so its nested live stderr details remain tied to its original activation and are also rejected. Parent-owned subagent roster events, explicit closes, and tool results still update normally. Newly launched children report live details under the new activation. +Live diagnostic compatibility: after a session reload or switch back, Kit rejects delayed diagnostics from earlier activations. Updated Kit children explicitly negotiate private diagnostic ownership over ACP metadata. Each invocation receives its own operation token, captured by its producers and retained for the child process lifetime. New nested details from a reused child therefore appear under the new activation, while delayed old details remain tied to the old activation. This also applies to native-fork siblings, detached tools, and nested grandchildren; request completion does not reassign their ownership. + +Older external agents that do not negotiate this extension still support ordinary ACP requests, results, and parent-owned subagent roster events. Their process-wide stderr lacks per-operation identity, so it retains immutable launch provenance: after a parent reactivation, new nested stderr details cannot be recovered from that stream. Kit does not invent missing ownership or retag old output. Ordinary stderr remains log output, and storage status remains process-global. A mixed-origin autonomous continuation or unsolicited external event has no single originating invocation; its diagnostics remain unattributed rather than borrowing the latest prompt. ### Branch from a text checkpoint in a new session diff --git a/src/acp_child.rs b/src/acp_child.rs index c81839f..20d1a0e 100644 --- a/src/acp_child.rs +++ b/src/acp_child.rs @@ -22,14 +22,82 @@ use agentkit_core::TurnCancellation; use serde::Deserialize; use serde_json::Value; use tokio::{ - io::{AsyncBufReadExt, AsyncRead, BufReader}, + io::{AsyncBufReadExt, AsyncRead, AsyncReadExt, BufReader}, process::Command, sync::{mpsc, oneshot, watch}, task::JoinSet, }; use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt}; -use crate::tools::mcp::CredentialStorage; +use crate::{ + events::{ + DiagnosticOperation, DiagnosticScope, offer_diagnostic_ownership, ownership_supported, + set_diagnostic_operation, + }, + tools::mcp::CredentialStorage, +}; + +/// Immutable producer ownership, separate from request-lifetime output routes. +/// Both the process actor and stderr drain retain this registry. Entries are +/// never removed on response, cancellation, or session close. +struct DiagnosticOperationRoutes { + operations: Mutex<(u64, HashMap)>, + negotiated: AtomicBool, +} + +impl DiagnosticOperationRoutes { + fn new() -> Self { + Self { + operations: Mutex::new((0, HashMap::new())), + negotiated: AtomicBool::new(false), + } + } + + fn register(&self) -> Result { + let scope = DiagnosticScope::capture(); + let mut routes = self + .operations + .lock() + .map_err(|_| ChildError::Failed("diagnostic operation registry unavailable".into()))?; + routes.0 = routes + .0 + .checked_add(1) + .ok_or_else(|| ChildError::Failed("diagnostic operation sequence exhausted".into()))?; + let operation = DiagnosticOperation::new(format!("operation-{}", routes.0)) + .expect("generated diagnostic operation is valid"); + routes.1.insert(operation.clone(), scope); + Ok(operation) + } + + fn resolve(&self, operation: &DiagnosticOperation) -> Option { + self.operations.lock().ok()?.1.get(operation).cloned() + } +} + +fn owned_close( + session_id: SessionId, + operation: &DiagnosticOperation, + negotiated: bool, +) -> CloseSessionRequest { + let mut request = CloseSessionRequest::new(session_id); + if negotiated { + set_diagnostic_operation(&mut request.meta, operation); + } + request +} + +fn owned_model( + session_id: SessionId, + model: &str, + operation: &DiagnosticOperation, + negotiated: bool, +) -> SetSessionConfigOptionRequest { + let mut request = SetSessionConfigOptionRequest::new(session_id, "model", model); + if negotiated { + set_diagnostic_operation(&mut request.meta, operation); + } + request +} const HANDSHAKE: Duration = Duration::from_secs(30); const PRE_HANDSHAKE_EXIT_SETTLE: Duration = Duration::from_millis(250); @@ -413,6 +481,7 @@ impl std::fmt::Display for ChildError { struct Prompt { // The worker, not the waiting caller, owns serialization until settlement. serial: tokio::sync::OwnedMutexGuard<()>, + operation: DiagnosticOperation, session_id: SessionId, text: String, cancellation: TurnCancellation, @@ -420,6 +489,7 @@ struct Prompt { } struct Fork { serial: tokio::sync::OwnedMutexGuard<()>, + operation: DiagnosticOperation, session_id: SessionId, model: Option, parent: Option<(String, String)>, @@ -427,6 +497,7 @@ struct Fork { reply: oneshot::Sender>, } struct Close { + operation: DiagnosticOperation, session_id: SessionId, reply: oneshot::Sender>, } @@ -532,6 +603,7 @@ fn rendered_text_only(value: &Value) -> Option<&str> { /// A logical ACP session. Multiple forked sessions may share one child process. #[derive(Clone)] pub(crate) struct ChildSession { + diagnostics: Arc, tx: mpsc::Sender, session_id: SessionId, capabilities: agentkit_acp::AgentCapabilities, @@ -551,6 +623,10 @@ impl ChildSession { ) -> Result { let context = config.harnesses.launch_context(&harness); let actor_context = context.clone(); + let diagnostics = Arc::new(DiagnosticOperationRoutes::new()); + // Admission ownership must be captured before the actor is spawned. + let admission = diagnostics.register()?; + let actor_diagnostics = Arc::clone(&diagnostics); let (tx, mut rx) = mpsc::channel(1); let (ready_tx, mut ready_rx) = oneshot::channel(); let (closed_tx, closed_rx) = watch::channel(false); @@ -564,6 +640,8 @@ impl ChildSession { model, depth, context: actor_context, + diagnostics: actor_diagnostics, + admission, }, &mut rx, ready_tx, @@ -592,6 +670,7 @@ impl ChildSession { }; match result { Ok(ready) => Ok(Self { + diagnostics, tx: actor_tx, session_id: ready.session_id, capabilities: ready.capabilities, @@ -620,6 +699,7 @@ impl ChildSession { } pub async fn close(&self) -> Result<(), ChildError> { + let operation = self.diagnostics.register()?; if let Some(ancestor_id) = &self.descendant_parent { // Explicit close belongs to its caller, which may be a new turn // using a reused or native-forked child. Deferred cleanup callers @@ -640,6 +720,7 @@ impl ChildSession { let (reply, response) = oneshot::channel(); self.tx .send(Request::Close(Close { + operation, session_id: self.session_id.clone(), reply, })) @@ -660,6 +741,7 @@ impl ChildSession { parent: Option<(String, String)>, cancellation: &TurnCancellation, ) -> Result { + let operation = self.diagnostics.register()?; let serial = tokio::select! { serial = self.serial.clone().lock_owned() => serial, () = cancellation.cancelled() => return Err(ChildError::Cancelled), @@ -674,6 +756,7 @@ impl ChildSession { tokio::select! { sent = self.tx.send(Request::Fork(Fork { serial, + operation, session_id: self.session_id.clone(), model: model.map(str::to_owned), parent, @@ -686,6 +769,7 @@ impl ChildSession { ChildError::TerminalFailed("nested agent process exited without a fork response".into()) })??; Ok(Self { + diagnostics: Arc::clone(&self.diagnostics), tx: self.tx.clone(), session_id, capabilities: self.capabilities.clone(), @@ -700,6 +784,7 @@ impl ChildSession { text: String, cancellation: TurnCancellation, ) -> Result { + let operation = self.diagnostics.register()?; let serial = tokio::select! { serial = self.serial.clone().lock_owned() => serial, () = cancellation.cancelled() => return Err(ChildError::Cancelled), @@ -707,6 +792,7 @@ impl ChildSession { let (reply, response) = oneshot::channel(); let request = Request::Prompt(Prompt { serial, + operation, session_id: self.session_id.clone(), text, cancellation: cancellation.clone(), @@ -723,6 +809,8 @@ impl ChildSession { } struct RunConfig { + diagnostics: Arc, + admission: DiagnosticOperation, config: ChildConfig, harness: String, persisted: Option<(String, bool)>, @@ -752,6 +840,8 @@ async fn run( closed: watch::Sender, ) -> Result<(), String> { let RunConfig { + diagnostics, + admission, config, harness, persisted, @@ -788,7 +878,15 @@ async fn run( .take() .ok_or("could not open ACP harness stderr")?; let descendant_parent = config.parent_id.clone(); - spawn_forward_stderr(stderr, harness.clone(), config.parent_id.clone()); + let (negotiation_tx, negotiation_rx) = watch::channel(None); + let _stderr_task = spawn_forward_stderr( + stderr, + harness.clone(), + config.parent_id.clone(), + Some((Arc::clone(&diagnostics), negotiation_rx)), + ); + // Retain registrations until this actor exits, even if stderr closes first. + let connection_diagnostics = Arc::clone(&diagnostics); let transport = ByteStreams::new(stdin.compat_write(), stdout.compat()); let routes = Arc::new(Mutex::new( HashMap::>>::new(), @@ -826,10 +924,19 @@ async fn run( agent_client_protocol::on_receive_request!(), ) .connect_with(transport, async move |connection| { - let initialized = connection.send_request(agentkit_acp::InitializeRequest::new(ProtocolVersion::V1)).block_task().await?; + let mut initialize = agentkit_acp::InitializeRequest::new(ProtocolVersion::V1); + offer_diagnostic_ownership(&mut initialize.meta); + let initialized = connection.send_request(initialize).block_task().await?; + let negotiated = ownership_supported(initialized.meta.as_ref()); + connection_diagnostics.negotiated.store(negotiated, Ordering::Release); + let _ = negotiation_tx.send(Some(negotiated)); let capabilities = initialized.agent_capabilities; let supports_close = capabilities.session_capabilities.close.is_some(); - let session = connection.send_request(agentkit_acp::NewSessionRequest::new(root.clone())).block_task().await?; + let mut request = agentkit_acp::NewSessionRequest::new(root.clone()); + if negotiated { + set_diagnostic_operation(&mut request.meta, &admission); + } + let session = connection.send_request(request).block_task().await?; if let Some(model) = model { let selectable = session.config_options.as_deref().unwrap_or_default().iter().any(|option| { option.id.to_string() == "model" && matches!(option.kind, SessionConfigKind::Select(_)) @@ -839,13 +946,13 @@ async fn run( let _ = ready.send(Err(error)); return std::future::pending().await; } - if let Err(error) = connection.send_request(SetSessionConfigOptionRequest::new(session.session_id.clone(), "model", model.as_str())).block_task().await { + if let Err(error) = connection.send_request(owned_model(session.session_id.clone(), model.as_str(), &admission, negotiated)).block_task().await { let error = format!("ACP harness {harness:?} rejected model selection {model:?}: {error}"); let _ = ready.send(Err(error)); return std::future::pending().await; } } - let sessions = Arc::new(Mutex::new(vec![session.session_id.clone()])); + let sessions = Arc::new(Mutex::new(vec![(session.session_id.clone(), admission.clone())])); let (fatal_tx, mut fatal_rx) = mpsc::unbounded_channel(); let mut tasks = JoinSet::new(); ready_flag.store(true, Ordering::Release); @@ -874,22 +981,24 @@ async fn run( (FORK_PARENT_NAME_META.into(), Value::String(name)), ])); } + if negotiated { + set_diagnostic_operation(&mut request.meta, &fork.operation); + } let mut request = Box::pin(connection.send_request(request).block_task()); let result = tokio::select! { result = &mut request => match result { Ok(response) => { let session_id = response.session_id; if let Ok(mut sessions) = sessions.lock() { - sessions.push(session_id.clone()); + sessions.push((session_id.clone(), fork.operation.clone())); } if let Some(model) = fork.model { let selected = match tokio::time::timeout( HANDSHAKE, connection - .send_request(SetSessionConfigOptionRequest::new( - session_id.clone(), - "model", - model.as_str(), + .send_request(owned_model( + session_id.clone(), model.as_str(), + &fork.operation, negotiated, )) .block_task(), ) @@ -905,14 +1014,14 @@ async fn run( }; if selected.is_err() && supports_close { let close = connection - .send_request(CloseSessionRequest::new(session_id.clone())) + .send_request(owned_close(session_id.clone(), &fork.operation, negotiated)) .block_task(); if tokio::time::timeout(CANCEL_SETTLE, close) .await .is_ok_and(|result| result.is_ok()) && let Ok(mut sessions) = sessions.lock() { - sessions.retain(|id| id != &session_id); + sessions.retain(|(id, _)| id != &session_id); } } selected @@ -937,7 +1046,7 @@ async fn run( && tokio::time::timeout( CANCEL_SETTLE, cleanup_connection - .send_request(CloseSessionRequest::new(session_id.clone())) + .send_request(owned_close(session_id.clone(), &fork.operation, negotiated)) .block_task(), ) .await @@ -945,7 +1054,7 @@ async fn run( if !closed && let Ok(mut sessions) = cleanup_sessions.lock() { - sessions.push(session_id); + sessions.push((session_id, fork.operation.clone())); } }); Err(ChildError::Cancelled) @@ -965,7 +1074,7 @@ async fn run( && tokio::time::timeout( CANCEL_SETTLE, cleanup_connection - .send_request(CloseSessionRequest::new(session_id.clone())) + .send_request(owned_close(session_id.clone(), &fork.operation, negotiated)) .block_task(), ) .await @@ -973,7 +1082,7 @@ async fn run( if !closed && let Ok(mut sessions) = cleanup_sessions.lock() { - sessions.push(session_id); + sessions.push((session_id, fork.operation.clone())); } }); Err(ChildError::Failed(format!( @@ -990,7 +1099,7 @@ async fn run( let sessions = Arc::clone(&sessions); tasks.spawn(async move { let request = connection - .send_request(CloseSessionRequest::new(close.session_id.clone())) + .send_request(owned_close(close.session_id.clone(), &close.operation, negotiated)) .block_task(); let result = tokio::time::timeout(CANCEL_SETTLE, request) .await @@ -1005,7 +1114,7 @@ async fn run( if result.is_ok() && let Ok(mut sessions) = sessions.lock() { - sessions.retain(|id| id != &close.session_id); + sessions.retain(|(id, _)| id != &close.session_id); } let _ = close.reply.send(result); }); @@ -1019,15 +1128,23 @@ async fn run( let session_id = prompt.session_id.clone(); let output = Arc::new(Mutex::new(ChildOutput::default())); if let Ok(mut routes) = routes.lock() { routes.insert(session_id.clone(), Arc::clone(&output)); } - let request = connection.send_request(agentkit_acp::PromptRequest::new( + let mut request = agentkit_acp::PromptRequest::new( session_id.clone(), vec![ContentBlock::Text(agentkit_acp::TextContent::new(prompt.text))], - )).block_task(); + ); + if negotiated { + set_diagnostic_operation(&mut request.meta, &prompt.operation); + } + let request = connection.send_request(request).block_task(); tokio::pin!(request); let (response, cancelled) = tokio::select! { biased; result = &mut request => (result.map_err(|error| error.to_string()), false), () = prompt.cancellation.cancelled() => { - let _ = connection.send_notification(CancelNotification::new(session_id.clone())); + let mut cancel = CancelNotification::new(session_id.clone()); + if negotiated { + set_diagnostic_operation(&mut cancel.meta, &prompt.operation); + } + let _ = connection.send_notification(cancel); match tokio::time::timeout(CANCEL_SETTLE, &mut request).await { Ok(result) => (result.map_err(|error| error.to_string()), true), Err(_) => { @@ -1055,9 +1172,9 @@ async fn run( } let session_ids = sessions.lock().map(|sessions| sessions.clone()).unwrap_or_default(); if supports_close { - for session_id in session_ids { + for (session_id, operation) in session_ids { let close = connection - .send_request(CloseSessionRequest::new(session_id)) + .send_request(owned_close(session_id, &operation, negotiated)) .block_task(); if let Ok(result) = tokio::time::timeout(CANCEL_SETTLE, close).await { result?; @@ -1104,37 +1221,84 @@ async fn run( }) } -// Capture before spawn. ACP v1 stderr is one process-wide stream shared by -// reused sessions and native forks. Neither a prompt response nor a new request -// identifies the operation that produced buffered bytes; even child activation -// IDs belong to another process and have no parent-operation mapping here. -// Keep the launch scope for this stream and its EOF cleanup. After reactivation, -// the reader must reject these events, including otherwise-current nested -// diagnostics. This intentionally sacrifices nested live detail rather than -// retagging delayed old bytes. Parent-owned Subagents lifecycle events and -// explicit ChildSession::close use their invocation scope and remain current. +/// Translate only explicitly negotiated, known operation tokens. Legacy children +/// retain immutable launch provenance; neither reader timing nor child epochs +/// authorize retagging their diagnostics after parent reactivation. fn spawn_forward_stderr( stderr: impl AsyncRead + Unpin + Send + 'static, label: String, ancestor_id: Option, + ownership: Option<( + Arc, + watch::Receiver>, + )>, ) -> tokio::task::JoinHandle<()> { - let diagnostics = crate::events::DiagnosticScope::capture(); + let launch = DiagnosticScope::capture(); tokio::spawn(async move { - forward_stderr(stderr, &label, ancestor_id.as_deref(), |output| { - match output { - ForwardedStderr::RuntimeLine(line) => { - // Child epochs belong to another process. Re-envelope the - // event with this parent's captured activation at the IPC - // boundary; recursively forwarded roster events keep their - // own subagent/ancestor IDs in the unchanged payload. - if let Some(event) = crate::events::parse(&line) { - diagnostics.emit(&event); + let mut stderr = stderr; + let mut pending = Vec::new(); + let routes = if let Some((routes, mut negotiation)) = ownership { + // Drain while initialization is pending: waiting without reading can + // fill the stderr pipe and prevent the child from replying at all. + // Interpret these bytes only after agreement (or failed startup). + let mut buffer = [0; 8192]; + while negotiation.borrow().is_none() { + tokio::select! { + biased; + changed = negotiation.changed() => { + if changed.is_err() { break; } + } + read = stderr.read(&mut buffer) => match read { + Ok(0) | Err(_) => { + let _ = negotiation.wait_for(Option::is_some).await; + break; + } + Ok(count) => pending.extend_from_slice(&buffer[..count]), } } - ForwardedStderr::Diagnostic(line) => eprintln!("{line}"), - ForwardedStderr::Cleanup(event) => diagnostics.emit(&event), } - }) + Some(routes) + } else { + None + }; + let negotiated = routes + .as_ref() + .filter(|routes| routes.negotiated.load(Ordering::Acquire)); + forward_stderr( + pending.as_slice().chain(stderr), + &label, + ancestor_id.as_deref(), + negotiated.map(Arc::as_ref), + |output| { + match output { + ForwardedStderr::RuntimeLine(line) => { + if let Some(routes) = negotiated { + if let Some(envelope) = crate::events::parse_diagnostic(&line) + && let Some(scope) = envelope + .operation + .as_ref() + .and_then(|op| routes.resolve(op)) + { + // Restores both the parent's activation and its + // upstream operation at every descendant hop. + scope.emit(&envelope.event); + } + } else if let Some(event) = crate::events::parse(&line) { + launch.emit(&event); + } + } + ForwardedStderr::Diagnostic(line) => eprintln!("{line}"), + ForwardedStderr::Cleanup(event) => launch.emit(&event), + ForwardedStderr::OwnedCleanup(event, operation) => { + if let Some(scope) = + negotiated.and_then(|routes| routes.resolve(&operation)) + { + scope.emit(&event); + } + } + } + }, + ) .await; }) } @@ -1144,33 +1308,81 @@ enum ForwardedStderr { RuntimeLine(String), Diagnostic(String), Cleanup(crate::events::RuntimeEvent), + OwnedCleanup(crate::events::RuntimeEvent, DiagnosticOperation), } async fn forward_stderr( stderr: impl AsyncRead + Unpin, label: &str, ancestor_id: Option<&str>, + negotiated: Option<&DiagnosticOperationRoutes>, mut output: impl FnMut(ForwardedStderr), ) { - let mut ancestors = ancestor_id - .into_iter() - .map(str::to_owned) - .collect::>(); + // Legacy cleanup keeps its launch ancestor. Negotiated cleanup is derived + // solely from accepted roster provenance, not the most recent invocation. + let mut ancestors = if negotiated.is_none() { + ancestor_id + .into_iter() + .map(str::to_owned) + .collect::>() + } else { + BTreeSet::new() + }; + let mut owned_ancestors = + Vec::<(DiagnosticScope, DiagnosticOperation, BTreeSet)>::new(); let mut lines = BufReader::new(stderr).lines(); while let Ok(Some(line)) = lines.next_line().await { - if let Some(event) = crate::events::parse(&line) - && event.forward_from_child() + let envelope = if negotiated.is_some() { + crate::events::parse_diagnostic(&line) + } else { + crate::events::parse(&line).map(|event| crate::events::DiagnosticEvent { + event, + activation: None, + operation: None, + }) + }; + if let Some(envelope) = envelope + && envelope.event.forward_from_child() { - if let crate::events::RuntimeEvent::SubagentStateChanged { + if let Some(routes) = negotiated { + let Some(operation) = envelope.operation else { + continue; + }; + let Some(scope) = routes.resolve(&operation) else { + // Missing, malformed, or unknown ownership never falls back. + continue; + }; + let index = owned_ancestors + .iter() + .position(|(owner, _, _)| owner == &scope) + .unwrap_or_else(|| { + owned_ancestors.push((scope, operation, BTreeSet::new())); + owned_ancestors.len() - 1 + }); + let ancestors = &mut owned_ancestors[index].2; + match &envelope.event { + crate::events::RuntimeEvent::SubagentStateChanged { + parent_id: Some(parent_id), + .. + } => { + ancestors.insert(parent_id.clone()); + } + crate::events::RuntimeEvent::SubagentDescendantsRemoved { ancestor_id } => { + ancestors.remove(ancestor_id); + } + _ => {} + } + } else if let crate::events::RuntimeEvent::SubagentStateChanged { parent_id: Some(parent_id), .. - } = event + } = &envelope.event { - ancestors.insert(parent_id); + ancestors.insert(parent_id.clone()); } - // Preserve recursively forwarded private runtime events byte-for-byte. output(ForwardedStderr::RuntimeLine(line)); } else if let Some(line) = harness_diagnostic(label, &line) { + // StorageStatus and ordinary external stderr retain existing global + // diagnostic behavior, independent of operation negotiation. output(ForwardedStderr::Diagnostic(line)); } } @@ -1179,6 +1391,14 @@ async fn forward_stderr( crate::events::RuntimeEvent::SubagentDescendantsRemoved { ancestor_id }, )); } + for (_, operation, ancestors) in owned_ancestors { + for ancestor_id in ancestors { + output(ForwardedStderr::OwnedCleanup( + crate::events::RuntimeEvent::SubagentDescendantsRemoved { ancestor_id }, + operation.clone(), + )); + } + } } fn pre_handshake_exit(context: &LaunchContext, status: std::process::ExitStatus) -> String { @@ -1246,6 +1466,7 @@ mod test_support { serial: Arc::new(tokio::sync::Mutex::new(())), closed: watch::channel(false).1, descendant_parent: None, + diagnostics: Arc::new(DiagnosticOperationRoutes::new()), }, closed_rx, ) @@ -1261,6 +1482,7 @@ mod test_support { serial: Arc::new(tokio::sync::Mutex::new(())), closed: watch::channel(false).1, descendant_parent: None, + diagnostics: Arc::new(DiagnosticOperationRoutes::new()), } } } @@ -1289,6 +1511,7 @@ mod tests { serial: Arc::new(tokio::sync::Mutex::new(())), closed: watch::channel(false).1, descendant_parent: None, + diagnostics: Arc::new(DiagnosticOperationRoutes::new()), }; let caller = child.clone(); let task = tokio::spawn(async move { @@ -2181,6 +2404,462 @@ mod tests { ); } + /// Real stdio requests capture producer tokens; a delayed producer is released + /// only by the following request, after the old prompt has already replied. + #[test] + fn negotiated_stdio_ownership_survives_reuse_and_native_forks() { + use crate::events::{self, RuntimeEvent}; + const PROBE: &str = "KIT_TEST_NEGOTIATED_CHILD_PROBE"; + if std::env::var_os(PROBE).is_some() { + tokio::runtime::Builder::new_multi_thread() + .worker_threads(2).enable_all().build().unwrap().block_on(async { + let root = tempfile::tempdir().unwrap(); + let script = root.path().join("ownership.py"); + std::fs::write(&script, r#" +import json, sys, threading +mode = sys.argv[1] +marker = sys.argv[2] +lock = threading.Lock() +next_session = 0 +delayed = None +barrier = threading.Barrier(2) +def send(value): + with lock: + print(json.dumps(value), flush=True) +def response(request, result): + send({'jsonrpc': '2.0', 'id': request['id'], 'result': result}) +def emit(operation, label): + event = {'event': 'child_started', 'call': label, 'tool': label, 'summary': label, 'at': 1} + if operation is not None: event['operation'] = operation + with lock: + print(marker + json.dumps(event), file=sys.stderr, flush=True) +def prompt(request): + global delayed + params = request['params'] + operation = params.get('_meta', {}).get('kitDiagnosticOperation') + label = params['prompt'][0]['text'] + if label == 'old': + delayed = (operation, label) + else: + if label == 'new' and delayed is not None: + emit(*delayed) + if label.startswith('sibling'): barrier.wait(timeout=5) + emit(operation, label) + if label == 'new': + emit('unknown-token', 'unknown') + emit(None, 'missing') + emit(42, 'malformed') + response(request, {'stopReason': 'end_turn'}) +for line in sys.stdin: + request = json.loads(line) + method = request['method'] + params = request.get('params', {}) + if method == 'initialize': + assert params['_meta']['kitDiagnosticOwnership'] == {'version': 1, 'transport': 'stderr'} + result = {'protocolVersion': 1, 'agentCapabilities': {'sessionCapabilities': {'fork': {}, 'close': {}}}} + if mode != 'absent': + result['_meta'] = {'kitDiagnosticOwnership': {'version': 1 if mode == 'supported' else 99, 'transport': 'stderr'}} + if mode == 'supported': + for _ in range(2048): emit(None, 'before-agreement') + response(request, result) + continue + operation = params.get('_meta', {}).get('kitDiagnosticOperation') + assert (operation is not None) == (mode == 'supported'), (method, params) + if method in ('session/new', 'session/fork'): + if method == 'session/fork': + assert params['_meta']['kit.subagent.parent_id'].startswith('fork-') + assert params['_meta']['kit.subagent.parent_name'] == 'Fork owner' + next_session += 1 + emit(operation, method + '-' + str(next_session)) + response(request, {'sessionId': str(next_session)}) + elif method == 'session/prompt': + threading.Thread(target=prompt, args=(request,), daemon=True).start() + elif method == 'session/close': + emit(operation, 'close-' + params['sessionId']) + response(request, {}) +"#).unwrap(); + for mode in ["supported", "absent", "unknown-version"] { + let harnesses = AcpHarnesses::new(BTreeMap::from([("ownership".into(), AcpHarnessProfile { + command: "python3".into(), + args: vec![script.to_string_lossy().into_owned(), mode.into(), events::EVENT_MARKER.into()], + permissions: AcpPermissionPolicy::Deny, + })])).unwrap(); + let config = ChildConfig { + root: root.path().to_path_buf(), model: "unused".into(), + provider: Default::default(), reasoning_effort: None, openrouter_api_key: None, + configured_mcp_config: None, configured_mcp_config_inherited: false, + legacy_mcp_config: false, mcp_config: None, credential_storage: Default::default(), + telemetry: Default::default(), harnesses, default_harness: "acp.ownership".into(), + parent_id: None, parent_name: None, + }; + let base = owned_child_invocation(mode, "launch", ChildSession::start( + config, "acp.ownership".into(), None, None, 1, TurnCancellation::default(), + )).await.unwrap(); + assert_eq!(base.diagnostics.negotiated.load(Ordering::Acquire), mode == "supported"); + owned_child_invocation(mode, "old", base.prompt("old".into(), TurnCancellation::default())).await.unwrap(); + owned_child_invocation(mode, "new", base.prompt("new".into(), TurnCancellation::default())).await.unwrap(); + // A→B→A and A→A activate fresh epochs on the same process. + owned_child_invocation(mode, "again", base.prompt("again".into(), TurnCancellation::default())).await.unwrap(); + let first = owned_child_invocation(mode, "fork-first", base.fork(None, + Some(("fork-first".into(), "Fork owner".into())), &TurnCancellation::default())).await.unwrap(); + let second = owned_child_invocation(mode, "fork-second", base.fork(None, + Some(("fork-second".into(), "Fork owner".into())), &TurnCancellation::default())).await.unwrap(); + let (first_result, second_result) = tokio::join!( + owned_child_invocation(mode, "sibling-first", first.prompt("sibling-first".into(), TurnCancellation::default())), + owned_child_invocation(mode, "sibling-second", second.prompt("sibling-second".into(), TurnCancellation::default())), + ); + first_result.unwrap(); second_result.unwrap(); + owned_child_invocation(mode, "close-first", first.close()).await.unwrap(); + owned_child_invocation(mode, "close-second", second.close()).await.unwrap(); + owned_child_invocation(mode, "close-base", base.close()).await.unwrap(); + // Every completed request still has its immutable registration. + assert_eq!(base.diagnostics.operations.lock().unwrap().1.len(), 11); + let mut closed = base.closed_signal(); + let registry = Arc::downgrade(&base.diagnostics); + drop(first); drop(second); drop(base); + tokio::time::timeout(Duration::from_secs(5), closed.wait_for(|closed| *closed)).await.unwrap().unwrap(); + tokio::time::timeout(Duration::from_secs(5), async { + while registry.upgrade().is_some() { + tokio::task::yield_now().await; + } + }).await.expect("process exit and stderr drain release registrations"); + } + }); + return; + } + let output = std::process::Command::new(std::env::current_exe().unwrap()) + .args([ + "--exact", + "acp_child::tests::negotiated_stdio_ownership_survives_reuse_and_native_forks", + "--nocapture", + ]) + .env(PROBE, "1") + .env(events::EVENTS_ENV, "1") + .output() + .unwrap(); + assert!( + output.status.success(), + "{}", + String::from_utf8_lossy(&output.stderr) + ); + let stderr = String::from_utf8(output.stderr).unwrap(); + let diagnostics = stderr + .lines() + .filter_map(events::parse_diagnostic) + .collect::>(); + for (call, owner) in [ + ("session/new-1", "launch"), + ("old", "old"), + ("new", "new"), + ("again", "again"), + ("session/fork-2", "fork-first"), + ("session/fork-3", "fork-second"), + ("sibling-first", "sibling-first"), + ("sibling-second", "sibling-second"), + ("close-2", "close-first"), + ("close-3", "close-second"), + ("close-1", "close-base"), + ] { + let envelope = diagnostics.iter().find(|envelope| { + envelope.operation.as_ref().is_some_and(|op| op.as_str() == format!("supported:{owner}")) + && matches!(&envelope.event, RuntimeEvent::ChildStarted { call: actual, .. } if actual == call) + }).unwrap_or_else(|| panic!("missing owned event {call}: {stderr}")); + let admission = diagnostics + .iter() + .find(|event| { + event.operation == envelope.operation + && matches!(event.event, RuntimeEvent::SessionStarted { .. }) + }) + .expect("invocation activation marker"); + assert!(envelope.activation.is_some()); + assert_eq!( + envelope.activation, admission.activation, + "wrong epoch for {call}" + ); + } + let activations = ["old", "new", "again"].map(|label| { + diagnostics.iter().find(|envelope| { + envelope.operation.as_ref().is_some_and(|operation| operation.as_str() == format!("supported:{label}")) + && matches!(&envelope.event, RuntimeEvent::ChildStarted { call, .. } if call == label) + }).unwrap().activation.as_ref().unwrap() + }); + assert_eq!(activations[0].session_id, "supported:A"); + assert_eq!(activations[1].session_id, "supported:B"); + assert_eq!(activations[2].session_id, "supported:A"); + assert!(activations[0].epoch < activations[1].epoch); + assert!(activations[1].epoch < activations[2].epoch); + assert!(!diagnostics.iter().any(|envelope| { + envelope.operation.as_ref().is_some_and(|op| op.as_str().starts_with("supported:")) + && matches!(&envelope.event, RuntimeEvent::ChildStarted { call, .. } if ["unknown", "missing", "malformed", "before-agreement"].contains(&call.as_str())) + })); + for mode in ["absent", "unknown-version"] { + let forwarded = diagnostics + .iter() + .filter(|envelope| { + matches!(&envelope.event, RuntimeEvent::ChildStarted { .. }) + && envelope + .operation + .as_ref() + .is_some_and(|op| op.as_str().starts_with(mode)) + }) + .collect::>(); + assert!(!forwarded.is_empty()); + assert!( + forwarded + .iter() + .all(|envelope| envelope.operation.as_ref().unwrap().as_str() + == format!("{mode}:launch")) + ); + } + } + + async fn owned_child_invocation( + mode: &str, + label: &str, + future: impl std::future::Future, + ) -> T { + let session = format!("{mode}:{}", if label == "new" { "B" } else { "A" }); + DiagnosticScope::with_operation(DiagnosticOperation::new(format!("{mode}:{label}"))) + .scope(async { + crate::events::activate_diagnostics(&session); + crate::events::scope_diagnostics(&session, future).await + }) + .await + } + + #[test] + fn real_grandchild_emission_restores_ownership_at_both_boundaries() { + use crate::events::{self, RuntimeEvent}; + const MODE: &str = "KIT_TEST_GRANDCHILD_OWNERSHIP"; + const TOKEN: &str = "KIT_TEST_GRANDCHILD_TOKEN"; + let mode = std::env::var(MODE).unwrap_or_default(); + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + if mode == "grandchild" { + runtime.block_on( + DiagnosticScope::with_operation(DiagnosticOperation::new( + std::env::var(TOKEN).unwrap(), + )) + .scope(async { + events::emit(&RuntimeEvent::ChildStarted { + call: "grandchild-producer".into(), + tool: "shell".into(), + summary: "real producer".into(), + at: 1, + }); + }), + ); + return; + } + if mode == "child" { + runtime.block_on(async { + // One reused intermediate process serves two distinct upstream + // invocations. Each grandchild has its own connection-local map. + for upstream in ["operation-1", "operation-2"] { + DiagnosticScope::with_operation(DiagnosticOperation::new(upstream.into())) + .scope(async { + events::activate_diagnostics("child-local"); + events::scope_diagnostics("child-local", async { + let routes = Arc::new(DiagnosticOperationRoutes::new()); + routes.negotiated.store(true, Ordering::Release); + let _unrelated = routes.register().unwrap(); + let operation = routes.register().unwrap(); + let mut producer = tokio::process::Command::new(std::env::current_exe().unwrap()) + .args(["--exact", "acp_child::tests::real_grandchild_emission_restores_ownership_at_both_boundaries", "--nocapture"]) + .env(MODE, "grandchild").env(TOKEN, operation.as_str()) + .stdout(Stdio::null()).stderr(Stdio::piped()).spawn().unwrap(); + let (_ready, ready) = watch::channel(Some(true)); + let forwarding = spawn_forward_stderr(producer.stderr.take().unwrap(), + "grandchild".into(), None, Some((routes, ready))); + assert!(producer.wait().await.unwrap().success()); + forwarding.await.unwrap(); + }).await; + }).await; + } + }); + return; + } + let routes = DiagnosticOperationRoutes::new(); + let original = runtime.block_on(async { + let mut expected = Vec::new(); + for upstream in ["root-old", "root-new"] { + DiagnosticScope::with_operation(DiagnosticOperation::new(upstream.into())) + .scope(async { + events::activate_diagnostics("grandchild-parent"); + events::scope_diagnostics("grandchild-parent", async { + expected.push(DiagnosticScope::capture()); + routes.register().unwrap(); + }) + .await; + }) + .await; + } + expected + }); + let output = std::process::Command::new(std::env::current_exe().unwrap()) + .args([ + "--exact", + "acp_child::tests::real_grandchild_emission_restores_ownership_at_both_boundaries", + "--nocapture", + ]) + .env(MODE, "child") + .env(events::EVENTS_ENV, "1") + .output() + .unwrap(); + assert!( + output.status.success(), + "{}", + String::from_utf8_lossy(&output.stderr) + ); + let mut restored = Vec::new(); + // All bytes arrive after parent reactivation. Child-local epochs cannot + // authorize a parent route; only the second connection's exact map can. + runtime.block_on(forward_stderr( + output.stderr.as_slice(), + "child", + None, + Some(&routes), + |item| { + if let ForwardedStderr::RuntimeLine(line) = item { + let envelope = events::parse_diagnostic(&line).unwrap(); + assert!(matches!(envelope.event, RuntimeEvent::ChildStarted { .. })); + assert_eq!(envelope.activation.unwrap().session_id, "child-local"); + restored.push( + routes + .resolve(envelope.operation.as_ref().unwrap()) + .unwrap(), + ); + } + }, + )); + assert_eq!(restored, original); + assert_ne!(restored[0], restored[1]); + } + + #[tokio::test] + async fn negotiated_cleanup_uses_only_accepted_scopes_and_deduplicates() { + use crate::events::{DiagnosticEvent, RuntimeEvent, SubagentStatus}; + use tokio::io::AsyncWriteExt; + let routes = Arc::new(DiagnosticOperationRoutes::new()); + let scope = DiagnosticScope::with_operation(DiagnosticOperation::new("upstream-a".into())); + let first = scope.scope(async { routes.register().unwrap() }).await; + let duplicate = scope.scope(async { routes.register().unwrap() }).await; + let second = DiagnosticScope::with_operation(DiagnosticOperation::new("upstream-b".into())) + .scope(async { routes.register().unwrap() }) + .await; + let roster = RuntimeEvent::SubagentStateChanged { + id: "descendant".into(), + name: "Nested".into(), + status: SubagentStatus::Working, + outcome: None, + generation: 1, + task: "inspect".into(), + parent_id: Some("shared-ancestor".into()), + parent_name: None, + harness: BUILTIN_HARNESS.into(), + model: None, + created_at_unix_ms: 1, + generation_started_at_unix_ms: 1, + generation_finished_at_unix_ms: None, + }; + for abrupt in [false, true] { + let (mut writer, reader) = tokio::io::duplex(64); + let write_first = first.clone(); + let write_duplicate = duplicate.clone(); + let write_second = second.clone(); + let write_roster = roster.clone(); + let producer = tokio::spawn(async move { + for operation in [ + Some(write_first), + Some(write_duplicate), + Some(write_second), + None, + DiagnosticOperation::new("unknown".into()), + ] { + let line = format!( + "{}{}\n", + crate::events::EVENT_MARKER, + serde_json::to_string(&DiagnosticEvent { + event: write_roster.clone(), + activation: None, + operation, + }) + .unwrap() + ); + writer.write_all(line.as_bytes()).await.unwrap(); + } + writer.write_all(b"ordinary stderr\n").await.unwrap(); + let storage = format!( + "{}{}\n", + crate::events::EVENT_MARKER, + serde_json::to_string(&RuntimeEvent::StorageStatus { + pending: true, + exhausted: false + }) + .unwrap() + ); + writer.write_all(storage.as_bytes()).await.unwrap(); + if abrupt { + writer.write_all(&[0xff]).await.unwrap(); + } + }); + let mut output = Vec::new(); + forward_stderr( + reader, + BUILTIN_HARNESS, + Some("unrelated-launch"), + Some(&routes), + |item| output.push(item), + ) + .await; + producer.await.unwrap(); + assert_eq!( + output + .iter() + .filter(|item| matches!(item, ForwardedStderr::RuntimeLine(_))) + .count(), + 3 + ); + let cleanups = output + .iter() + .filter_map(|item| match item { + ForwardedStderr::OwnedCleanup( + RuntimeEvent::SubagentDescendantsRemoved { ancestor_id }, + operation, + ) => { + assert_eq!(ancestor_id, "shared-ancestor"); + Some(operation.clone()) + } + ForwardedStderr::Cleanup(_) => { + panic!("negotiated cleanup used launch provenance") + } + _ => None, + }) + .collect::>(); + assert_eq!(cleanups.len(), 2); + assert!(cleanups.contains(&first)); + assert!(cleanups.contains(&second)); + assert!(output.iter().any(|item| matches!(item, ForwardedStderr::Diagnostic(line) if line.ends_with("ordinary stderr")))); + assert!(output.iter().any(|item| matches!(item, ForwardedStderr::Diagnostic(line) if line.contains("storage_status")))); + // EOF and stream errors do not retire registrations while a process + // (or a retained logical session) could still own the connection. + assert!(routes.resolve(&first).is_some()); + } + } + + #[test] + fn diagnostic_operation_sequence_exhaustion_never_reuses_a_token() { + let routes = DiagnosticOperationRoutes::new(); + let first = routes.register().unwrap(); + routes.operations.lock().unwrap().0 = u64::MAX; + assert!(routes.register().is_err()); + assert!(routes.register().is_err()); + assert!(routes.resolve(&first).is_some()); + assert_eq!(routes.operations.lock().unwrap().1.len(), 1); + } + mod forwards_subagent_events { use super::*; @@ -2208,7 +2887,7 @@ mod tests { ); let input = format!("{line}\n"); let mut output = Vec::new(); - forward_stderr(input.as_bytes(), "acp.kit", Some("s-owner"), |item| { + forward_stderr(input.as_bytes(), "acp.kit", Some("s-owner"), None, |item| { output.push(item) }) .await; @@ -2247,7 +2926,7 @@ mod tests { // before awaiting the returned forwarder task. #[allow(clippy::async_yields_async)] let producer = events::scope_diagnostics(&session, async move { - spawn_forward_stderr(reader, "acp.kit".into(), Some(route.into())) + spawn_forward_stderr(reader, "acp.kit".into(), Some(route.into()), None) }); let (task, delayed) = if route == "construction" { (None, Some(producer)) @@ -2260,6 +2939,7 @@ mod tests { reader, "acp.kit".into(), Some(route.into()), + None, )), None, ) @@ -2282,6 +2962,7 @@ mod tests { ancestor_id: format!("forward-{route}"), }; let child_line = events::DiagnosticEvent { + operation: None, event, activation: Some(DiagnosticActivation { session_id: "child-process".into(), @@ -2441,7 +3122,10 @@ mod tests { let stderr = child.stderr.take().unwrap(); let mut output = Vec::new(); - forward_stderr(stderr, "acp.kit", Some("s-owner"), |item| output.push(item)).await; + forward_stderr(stderr, "acp.kit", Some("s-owner"), None, |item| { + output.push(item) + }) + .await; let status = child.wait().await.unwrap(); assert_eq!(status.code(), Some(23)); @@ -2449,7 +3133,9 @@ mod tests { .iter() .filter_map(|item| match item { ForwardedStderr::RuntimeLine(line) => Some(line.clone()), - ForwardedStderr::Diagnostic(_) | ForwardedStderr::Cleanup(_) => None, + ForwardedStderr::Diagnostic(_) + | ForwardedStderr::Cleanup(_) + | ForwardedStderr::OwnedCleanup(_, _) => None, }) .collect::>(); assert_eq!(forwarded, lines); @@ -2470,14 +3156,16 @@ mod tests { id, .. }) => id == "s-owner", - ForwardedStderr::Diagnostic(_) | ForwardedStderr::Cleanup(_) => false, + ForwardedStderr::Diagnostic(_) + | ForwardedStderr::Cleanup(_) + | ForwardedStderr::OwnedCleanup(_, _) => false, })); } #[tokio::test] async fn emits_cleanup_once_on_normal_eof() { let mut output = Vec::new(); - forward_stderr(&b""[..], "acp.kit", Some("s-owner"), |item| { + forward_stderr(&b""[..], "acp.kit", Some("s-owner"), None, |item| { output.push(item) }) .await; @@ -2487,7 +3175,7 @@ mod tests { #[tokio::test] async fn emits_cleanup_once_on_abrupt_stream_error() { let mut output = Vec::new(); - forward_stderr(&b"\xff"[..], "acp.kit", Some("s-owner"), |item| { + forward_stderr(&b"\xff"[..], "acp.kit", Some("s-owner"), None, |item| { output.push(item) }) .await; diff --git a/src/events.rs b/src/events.rs index 7f1ab66..fcb8783 100644 --- a/src/events.rs +++ b/src/events.rs @@ -14,6 +14,7 @@ //! ACP hosts never see the extra chatter. use std::{ + cell::RefCell, collections::HashMap, io::Write, sync::{Mutex, OnceLock}, @@ -145,6 +146,8 @@ pub struct DiagnosticEvent { pub event: RuntimeEvent, #[serde(default, skip_serializing_if = "Option::is_none")] pub activation: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub operation: Option, } #[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] @@ -153,6 +156,72 @@ pub struct DiagnosticActivation { pub epoch: u64, } +/// Private, connection-local diagnostic identity. This is not a persisted ID or +/// an ACP request ID. Reject malformed values instead of inventing ownership. +#[derive(Clone, Debug, Serialize, PartialEq, Eq, Hash)] +#[serde(transparent)] +pub struct DiagnosticOperation(String); + +impl DiagnosticOperation { + pub fn new(value: String) -> Option { + (!value.is_empty() + && value.len() <= 128 + && value + .bytes() + .all(|byte| byte.is_ascii_alphanumeric() || b"-_.:".contains(&byte))) + .then_some(Self(value)) + } + + pub fn as_str(&self) -> &str { + &self.0 + } +} + +impl<'de> Deserialize<'de> for DiagnosticOperation { + fn deserialize>(deserializer: D) -> Result { + let value = String::deserialize(deserializer)?; + Self::new(value) + .ok_or_else(|| serde::de::Error::custom("invalid diagnostic operation token")) + } +} + +pub(crate) const OWNERSHIP_META_KEY: &str = "kitDiagnosticOwnership"; +pub(crate) const OPERATION_META_KEY: &str = "kitDiagnosticOperation"; + +/// Both offer and acknowledgement must explicitly select our supported contract. +pub(crate) fn ownership_supported(meta: Option<&serde_json::Map>) -> bool { + meta.and_then(|meta| meta.get(OWNERSHIP_META_KEY)) + == Some(&serde_json::json!({"version": 1, "transport": "stderr"})) +} + +/// Merge our private extension; never replace unrelated ACP metadata. +pub(crate) fn offer_diagnostic_ownership(meta: &mut Option>) { + meta.get_or_insert_with(Default::default).insert( + OWNERSHIP_META_KEY.into(), + serde_json::json!({"version": 1, "transport": "stderr"}), + ); +} + +pub(crate) fn diagnostic_operation( + meta: Option<&serde_json::Map>, + negotiated: bool, +) -> Option { + if !negotiated { + return None; + } + serde_json::from_value(meta?.get(OPERATION_META_KEY)?.clone()).ok() +} + +pub(crate) fn set_diagnostic_operation( + meta: &mut Option>, + operation: &DiagnosticOperation, +) { + meta.get_or_insert_with(Default::default).insert( + OPERATION_META_KEY.into(), + Value::String(operation.as_str().to_owned()), + ); +} + pub const ACTIVATION_META_KEY: &str = "kitRuntimeActivation"; #[derive(Default)] @@ -185,16 +254,6 @@ pub fn activate_diagnostics(session_id: &str) -> u64 { epoch } -#[cfg(test)] -pub(crate) fn diagnostic_epoch(session_id: &str) -> Option { - diagnostic_routes() - .lock() - .expect("diagnostic routes lock") - .sessions - .get(session_id) - .copied() -} - /// Restore an actor's route without minting a new activation on each turn. pub fn restore_diagnostics(session_id: &str) { let mut routes = diagnostic_routes().lock().expect("diagnostic routes lock"); @@ -210,6 +269,10 @@ pub fn restore_diagnostics(session_id: &str) { tokio::task_local! { static DIAGNOSTIC_ACTIVATION: Option; + static DIAGNOSTIC_OPERATION: Option; + // Rebound only at actual loop input consumption, before admitting downstream + // producers. Captured DiagnosticScopes never retain this mutable cell. + static CONSUMED_DIAGNOSTIC_ORIGIN: RefCell>; } /// Preserve the turn's activation even if another actor restores its route, or @@ -229,7 +292,11 @@ pub fn scope_diagnostics( epoch: *epoch, }) }; - DiagnosticScope(activation).scope(future) + DiagnosticScope { + activation, + ..DiagnosticScope::capture() + } + .scope(future) } fn write_diagnostic_marker(routes: &DiagnosticRoutes) { @@ -243,6 +310,7 @@ fn write_diagnostic_marker(routes: &DiagnosticRoutes) { session_id: activation.session_id.clone(), }, activation: Some(activation.clone()), + operation: DiagnosticScope::capture().operation, }, ); } @@ -252,18 +320,67 @@ fn write_diagnostic_marker(routes: &DiagnosticRoutes) { /// when that task is first polled. Unscoped producers never borrow the active UI /// route, including after a durable session ID is reactivated. #[derive(Clone, Debug, Default, PartialEq, Eq)] -pub(crate) struct DiagnosticScope(Option); +pub(crate) struct DiagnosticScope { + activation: Option, + operation: Option, +} impl DiagnosticScope { pub(crate) fn capture() -> Self { - Self(DIAGNOSTIC_ACTIVATION.try_with(Clone::clone).unwrap_or(None)) + if let Ok(Some(origin)) = + CONSUMED_DIAGNOSTIC_ORIGIN.try_with(|origin| origin.borrow().clone()) + { + return origin; + } + Self { + activation: DIAGNOSTIC_ACTIVATION.try_with(Clone::clone).unwrap_or(None), + operation: DIAGNOSTIC_OPERATION.try_with(Clone::clone).unwrap_or(None), + } + } + + /// Incoming operations replace upstream ownership, never local activation. + pub(crate) fn with_operation(operation: Option) -> Self { + Self { + operation, + ..Self::capture() + } + } + + pub(crate) fn operation(&self) -> Option<&DiagnosticOperation> { + self.operation.as_ref() } pub(crate) fn scope( &self, future: F, ) -> impl std::future::Future + use { - DIAGNOSTIC_ACTIVATION.scope(self.0.clone(), future) + DIAGNOSTIC_ACTIVATION.scope( + self.activation.clone(), + DIAGNOSTIC_OPERATION.scope( + self.operation.clone(), + CONSUMED_DIAGNOSTIC_ORIGIN.scope(RefCell::new(None), future), + ), + ) + } + + /// Preserve identity across blocking filesystem work without an async runtime. + pub(crate) fn sync_scope(&self, work: impl FnOnce() -> R) -> R { + DIAGNOSTIC_ACTIVATION.sync_scope(self.activation.clone(), || { + DIAGNOSTIC_OPERATION.sync_scope(self.operation.clone(), || { + CONSUMED_DIAGNOSTIC_ORIGIN.sync_scope(RefCell::new(None), work) + }) + }) + } + + /// Bind the exact causes consumed by the loop before its next producer. + /// This is never called on request arrival. A previously captured scope is + /// immutable, and every async/blocking scope starts a separate binding. + /// Rebind both fields: root actors without upstream tokens must also retain + /// an old background producer's activation rather than the latest UI epoch. + pub(crate) fn bind_consumed_origin(&self) { + let _ = CONSUMED_DIAGNOSTIC_ORIGIN.try_with(|origin| { + *origin.borrow_mut() = Some(self.clone()); + }); } pub(crate) fn emit(&self, event: &RuntimeEvent) { @@ -277,7 +394,12 @@ impl DiagnosticScope { activation: if matches!(event, RuntimeEvent::StorageStatus { .. }) { None } else { - self.0.clone() + self.activation.clone() + }, + operation: if matches!(event, RuntimeEvent::StorageStatus { .. }) { + None + } else { + self.operation().cloned() }, }, ); @@ -304,17 +426,6 @@ fn write_diagnostic(writer: &mut impl Write, event: &DiagnosticEvent) { } } -#[cfg(test)] -fn write_event(writer: &mut impl Write, event: &RuntimeEvent) { - write_diagnostic( - writer, - &DiagnosticEvent { - event: event.clone(), - activation: None, - }, - ); -} - /// Decode both current activation-aware and historical unscoped stderr lines. #[must_use] pub fn parse_diagnostic(line: &str) -> Option { @@ -391,6 +502,20 @@ fn truncate(text: &str, limit: usize) -> String { format!("{kept}…") } +#[cfg(test)] +pub(crate) mod test_support { + use super::*; + + pub(crate) fn diagnostic_epoch(session_id: &str) -> Option { + diagnostic_routes() + .lock() + .expect("diagnostic routes lock") + .sessions + .get(session_id) + .copied() + } +} + #[cfg(test)] mod tests { use std::io::{self, Write}; @@ -401,9 +526,20 @@ mod tests { DIAGNOSTIC_ACTIVATION, DiagnosticActivation, DiagnosticEvent, EVENT_MARKER, GenerationOutcome, RuntimeEvent, SubagentStatus, activate_diagnostics, parse, parse_diagnostic, restore_diagnostics, scope_diagnostics, summarize_input, - summarize_output, write_diagnostic, write_event, + summarize_output, write_diagnostic, }; + fn write_event(writer: &mut impl Write, event: &RuntimeEvent) { + write_diagnostic( + writer, + &DiagnosticEvent { + event: event.clone(), + activation: None, + operation: None, + }, + ); + } + #[test] fn diagnostic_envelope_preserves_legacy_replay_and_rejects_malformed_identity() { let legacy = format!( @@ -413,6 +549,7 @@ mod tests { let event = parse(&legacy).unwrap(); assert_eq!(parse_diagnostic(&legacy).unwrap().activation, None); let current = DiagnosticEvent { + operation: None, event: event.clone(), activation: Some(DiagnosticActivation { session_id: "source".into(), @@ -438,6 +575,240 @@ mod tests { } } + #[test] + fn ownership_negotiation_is_explicit_and_metadata_is_merged() { + use super::*; + assert!(!ownership_supported(None)); + for value in [ + json!(null), + json!(true), + json!({"version": 2, "transport": "stderr"}), + json!({"version": "1", "transport": "stderr"}), + json!({"version": 1, "transport": "stdout"}), + ] { + let meta = serde_json::Map::from_iter([(OWNERSHIP_META_KEY.into(), value)]); + assert!(!ownership_supported(Some(&meta))); + } + let mut meta = Some(serde_json::Map::from_iter([ + ("kitForkParent".into(), json!("parent")), + (ACTIVATION_META_KEY.into(), json!(7)), + ])); + offer_diagnostic_ownership(&mut meta); + assert!(ownership_supported(meta.as_ref())); + let operation = DiagnosticOperation::new("op:1".into()).unwrap(); + set_diagnostic_operation(&mut meta, &operation); + assert_eq!(diagnostic_operation(meta.as_ref(), true), Some(operation)); + assert_eq!(diagnostic_operation(meta.as_ref(), false), None); + assert_eq!(meta.as_ref().unwrap()["kitForkParent"], "parent"); + assert_eq!(meta.as_ref().unwrap()[ACTIVATION_META_KEY], 7); + for bad in [ + json!(null), + json!(false), + json!(1), + json!(""), + json!("has space"), + json!("bad\nline"), + json!("é"), + json!("x".repeat(129)), + json!({"id":"op:1"}), + ] { + meta.as_mut() + .unwrap() + .insert(OPERATION_META_KEY.into(), bad); + assert_eq!(diagnostic_operation(meta.as_ref(), true), None); + } + } + + #[test] + fn optional_operation_envelope_is_additive_and_strict() { + use super::DiagnosticOperation; + let old = json!({"event":"session_started", "session_id":"child", + "activation":{"session_id":"child", "epoch":3}, "future_key":true}); + let decoded: DiagnosticEvent = serde_json::from_value(old.clone()).unwrap(); + assert_eq!(decoded.operation, None); + assert!( + serde_json::to_value(&decoded) + .unwrap() + .get("operation") + .is_none() + ); + let mut owned = old.clone(); + owned["operation"] = json!("op:123"); + let decoded: DiagnosticEvent = serde_json::from_value(owned).unwrap(); + assert_eq!(decoded.operation, DiagnosticOperation::new("op:123".into())); + assert_eq!( + serde_json::to_value(decoded).unwrap()["operation"], + "op:123" + ); + for bad in [ + json!(12), + json!(""), + json!("has space"), + json!("x".repeat(129)), + json!({"operation":"op:123"}), + json!([]), + ] { + let mut malformed = old.clone(); + malformed["operation"] = bad; + assert!(serde_json::from_value::(malformed.clone()).is_err()); + // Historical replay still sees exactly the frozen event payload. + assert!(serde_json::from_value::(malformed).is_ok()); + } + } + + #[tokio::test] + async fn operation_survives_local_activation_and_detached_producer() { + use super::{DiagnosticOperation, DiagnosticScope, inherit_diagnostics}; + let incoming = DiagnosticOperation::new("upstream:1".into()).unwrap(); + let scope = DiagnosticScope::with_operation(Some(incoming.clone())); + let (release, released) = tokio::sync::oneshot::channel(); + activate_diagnostics("owned-operation-test"); + #[allow(clippy::async_yields_async)] + let producer = scope + .scope(async { + scope_diagnostics("owned-operation-test", async { + let original = DiagnosticScope::capture(); + assert_eq!(original.operation(), Some(&incoming)); + tokio::spawn(inherit_diagnostics(async move { + released.await.unwrap(); + assert_eq!(DiagnosticScope::capture(), original); + })) + }) + .await + }) + .await; + activate_diagnostics("owned-operation-test"); + let next = DiagnosticScope::with_operation(DiagnosticOperation::new("upstream:2".into())); + next.scope(async { + scope_diagnostics("owned-operation-test", async { + assert_eq!(DiagnosticScope::capture().operation(), next.operation()); + release.send(()).unwrap(); + producer.await.unwrap(); + }) + .await; + }) + .await; + let blocking_origin = scope.clone(); + let actual = tokio::task::spawn_blocking(move || { + blocking_origin.sync_scope(DiagnosticScope::capture) + }) + .await + .unwrap(); + assert_eq!(actual, scope); + assert_eq!(DiagnosticScope::capture().operation(), None); + } + + #[tokio::test] + async fn consumed_origin_rebind_does_not_mutate_captured_producers() { + use super::{DiagnosticOperation, DiagnosticScope, inherit_diagnostics}; + let old = DiagnosticScope { + activation: Some(DiagnosticActivation { + session_id: "same".into(), + epoch: 1, + }), + operation: DiagnosticOperation::new("old-operation".into()), + }; + let new = DiagnosticScope { + activation: Some(DiagnosticActivation { + session_id: "same".into(), + epoch: 2, + }), + operation: DiagnosticOperation::new("new-operation".into()), + }; + old.scope(async { + let (release, released) = tokio::sync::oneshot::channel(); + let detached = tokio::spawn(inherit_diagnostics(async move { + released.await.unwrap(); + DiagnosticScope::capture() + })); + // A synchronous loop observer binds the actual consumed causes + // before the same poll starts a model/tool producer. + new.bind_consumed_origin(); + assert_eq!(DiagnosticScope::capture(), new); + release.send(()).unwrap(); + assert_eq!(detached.await.unwrap(), old); + DiagnosticScope::default().bind_consumed_origin(); + assert_eq!(DiagnosticScope::capture(), DiagnosticScope::default()); + // Explicit child/blocking scopes never share the mutable binding. + assert_eq!(old.scope(async { DiagnosticScope::capture() }).await, old); + assert_eq!(new.sync_scope(DiagnosticScope::capture), new); + assert_eq!(DiagnosticScope::capture(), DiagnosticScope::default()); + }) + .await; + assert_eq!(DiagnosticScope::capture(), DiagnosticScope::default()); + } + + #[test] + fn real_emission_keeps_operation_but_storage_status_is_global() { + const CHILD: &str = "KIT_EVENTS_OPERATION_EMISSION_TEST"; + if std::env::var_os(CHILD).is_some() { + let runtime = tokio::runtime::Builder::new_current_thread() + .build() + .unwrap(); + runtime.block_on(async { + activate_diagnostics("owned-emission"); + super::DiagnosticScope::with_operation(super::DiagnosticOperation::new( + "upstream:7".into(), + )) + .scope(async { + scope_diagnostics("owned-emission", async { + // Use the real async producer/emitter and stderr pipe. + tokio::spawn(super::inherit_diagnostics(async { + super::emit(&RuntimeEvent::CompactionStarted { + reason: "owned".into(), + at: 1, + }); + super::emit(&RuntimeEvent::StorageStatus { + pending: true, + exhausted: false, + }); + })) + .await + .unwrap(); + }) + .await; + }) + .await; + }); + return; + } + let output = std::process::Command::new(std::env::current_exe().unwrap()) + .args([ + "--exact", + "events::tests::real_emission_keeps_operation_but_storage_status_is_global", + "--nocapture", + ]) + .env(CHILD, "1") + .env(super::EVENTS_ENV, "1") + .output() + .unwrap(); + assert!( + output.status.success(), + "{}", + String::from_utf8_lossy(&output.stderr) + ); + let envelopes = String::from_utf8(output.stderr) + .unwrap() + .lines() + .filter_map(parse_diagnostic) + .collect::>(); + let owned = envelopes + .iter() + .find(|event| matches!(event.event, RuntimeEvent::CompactionStarted { .. })) + .unwrap(); + assert_eq!(owned.operation.as_ref().unwrap().as_str(), "upstream:7"); + assert_eq!( + owned.activation.as_ref().unwrap().session_id, + "owned-emission" + ); + let global = envelopes + .iter() + .find(|event| matches!(event.event, RuntimeEvent::StorageStatus { .. })) + .unwrap(); + assert_eq!(global.operation, None); + assert_eq!(global.activation, None); + } + #[tokio::test] async fn diagnostic_epochs_change_on_reactivation_not_turn_restoration() { let session_id = "event-epoch-test"; diff --git a/src/protocols/acp.rs b/src/protocols/acp.rs index 09909b3..f1348cf 100644 --- a/src/protocols/acp.rs +++ b/src/protocols/acp.rs @@ -5,7 +5,7 @@ use std::{ pin::Pin, sync::{ Arc, Mutex, Weak, - atomic::{AtomicU64, Ordering}, + atomic::{AtomicBool, AtomicU64, Ordering}, }, time::Duration, }; @@ -54,6 +54,10 @@ mod skill_catalog; pub mod v2; use crate::{ + events::{ + DiagnosticScope, diagnostic_operation, inherit_diagnostics, offer_diagnostic_ownership, + ownership_supported, + }, provider::{ ModelGroup, ModelSelection, ProviderKind, ReasoningEffort, SelectableAdapter, authentication_method_id, @@ -458,21 +462,25 @@ fn list_sessions_error(error: ListSessionsError) -> agent_client_protocol::Error enum Command { Prompt { + diagnostics: DiagnosticScope, request: PromptRequest, reply: oneshot::Sender>, }, Cancel, SetConfig { + diagnostics: DiagnosticScope, request: SetSessionConfigOptionRequest, cancellation_generation: u64, reply: oneshot::Sender>, }, Fork { + diagnostics: DiagnosticScope, parent_context: Option<(String, String)>, reply: oneshot::Sender>, }, Close { + diagnostics: DiagnosticScope, reply: oneshot::Sender<()>, }, } @@ -766,23 +774,30 @@ impl SessionRegistry { let mut closing = JoinSet::new(); for mut session in sessions.iter().cloned() { - closing.spawn(async move { + closing.spawn(inherit_diagnostics(async move { cancel_background_jobs(&session.tasks, &session.background_jobs).await; if let Some(commands) = session.commands.upgrade() { let (reply, acknowledged) = oneshot::channel(); - if commands.send(Command::Close { reply }).await.is_ok() { + if commands + .send(Command::Close { + diagnostics: DiagnosticScope::capture(), + reply, + }) + .await + .is_ok() + { let _ = acknowledged.await; } } session.completed.wait_for(|done| *done).await.is_ok() - }); + })); } for mut session in v2_sessions.iter().cloned() { - closing.spawn(async move { + closing.spawn(inherit_diagnostics(async move { (session.close)().await; session.completed.wait_for(|done| *done).await.is_ok() - }); + })); } let graceful = matches!( @@ -878,7 +893,9 @@ async fn bounded_authentication_reset( where T: Send + 'static, { - let task = tokio::spawn(async move { registry.reset_authentication_with(reset).await }); + let task = tokio::spawn(inherit_diagnostics(async move { + registry.reset_authentication_with(reset).await + })); match timeout(limit, task).await { Ok(Ok((complete, output))) => AuthenticationReset::Completed(complete, output), Ok(Err(error)) => AuthenticationReset::Failed(error), @@ -897,7 +914,12 @@ async fn logout_authentication( )); } let logout = async move { - match tokio::task::spawn_blocking(move || runtime.logout_authentication()).await { + let diagnostics = DiagnosticScope::capture(); + match tokio::task::spawn_blocking(move || { + diagnostics.sync_scope(|| runtime.logout_authentication()) + }) + .await + { Ok(result) => result.map_err(|error| error.to_string()), Err(error) => Err(format!("authentication logout task failed: {error}")), } @@ -1109,6 +1131,7 @@ impl Drop for SessionPublicationRollback<'_> { } struct Server { + diagnostic_ownership: AtomicBool, runtime: Arc, integration: Arc, registry: SessionRegistry, @@ -1117,8 +1140,19 @@ struct Server { } impl Server { + fn diagnostic_scope( + &self, + meta: Option<&serde_json::Map>, + ) -> DiagnosticScope { + DiagnosticScope::with_operation(diagnostic_operation( + meta, + self.diagnostic_ownership.load(Ordering::Acquire), + )) + } + fn new(runtime: Arc, integration: AcpIntegration, registry: SessionRegistry) -> Self { Self { + diagnostic_ownership: AtomicBool::new(false), runtime, integration: Arc::new(integration), registry, @@ -1223,16 +1257,24 @@ impl Server { async fn initialize(&self, request: InitializeRequest) -> InitializeResponse { let logout_authentication = self.runtime.supports_logout_authentication(); - InitializeResponse::new(agent_client_protocol::schema::ProtocolVersion::V1) - .agent_capabilities(capabilities(logout_authentication)) - .auth_methods(terminal_auth_methods( - &request.client_capabilities, - |provider| self.runtime.supports_terminal_authentication(provider), - )) - .agent_info(agentkit_acp::Implementation::new( - self.integration.name().to_string(), - self.integration.version().to_string(), - )) + let negotiated = ownership_supported(request.meta.as_ref()); + self.diagnostic_ownership + .store(negotiated, Ordering::Release); + let mut response = + InitializeResponse::new(agent_client_protocol::schema::ProtocolVersion::V1) + .agent_capabilities(capabilities(logout_authentication)) + .auth_methods(terminal_auth_methods( + &request.client_capabilities, + |provider| self.runtime.supports_terminal_authentication(provider), + )) + .agent_info(agentkit_acp::Implementation::new( + self.integration.name().to_string(), + self.integration.version().to_string(), + )); + if negotiated { + offer_diagnostic_ownership(&mut response.meta); + } + response } async fn new_session( @@ -1240,69 +1282,80 @@ impl Server { request: NewSessionRequest, connection: ConnectionTo, ) -> Result { - // session/new retains the configured-or-generated selection semantics. - let claim = self.runtime.claim_session()?; - let attached = self - .attach_session( - request.cwd, - request.additional_directories, - connection, - claim, - None, - ) - .await?; - let AttachedSession { - session_id, - config_options, - activation, - .. - } = attached; - let _ = activation.send(()); - Ok(NewSessionResponse::new(session_id).config_options(Some(config_options))) + self.diagnostic_scope(request.meta.as_ref()) + .scope(async { + // session/new retains the configured-or-generated selection semantics. + let claim = self.runtime.claim_session()?; + let attached = self + .attach_session( + request.cwd, + request.additional_directories, + connection, + claim, + None, + ) + .await?; + let AttachedSession { + session_id, + config_options, + activation, + .. + } = attached; + let _ = activation.send(()); + Ok(NewSessionResponse::new(session_id).config_options(Some(config_options))) + }) + .await } async fn list_sessions( &self, request: ListSessionsRequest, ) -> Result { - let cwd = self.runtime.root().to_path_buf(); - let offset = request - .cursor - .as_deref() - .map(parse_session_list_cursor) - .transpose()? - .unwrap_or(0); - if request - .cwd - .as_ref() - .is_some_and(|requested| requested != &cwd) - { - return Ok(ListSessionsResponse::new(Vec::new())); - } - let root = self.runtime.root().to_path_buf(); - let catalog = tokio::task::spawn_blocking(move || crate::session::catalog(&root)) - .await - .map_err(|error| { - ListSessionsError::Runtime(AcpRuntimeError::Loop(format!( - "session catalog worker failed: {error}" - ))) - })? - .map_err(|error| ListSessionsError::Runtime(AcpRuntimeError::Loop(error)))?; - if offset > catalog.len() { - return Err(ListSessionsError::InvalidCursor); - } - let end = catalog.len().min(offset + SESSION_LIST_PAGE_SIZE); - let sessions = catalog[offset..end] - .iter() - .map(|entry| { - SessionInfo::new(entry.id.clone(), cwd.clone()) - .title(entry.title.as_deref().map(str::to_owned)) - .updated_at(entry.updated_at_rfc3339()) + self.diagnostic_scope(request.meta.as_ref()) + .scope(async { + let cwd = self.runtime.root().to_path_buf(); + let offset = request + .cursor + .as_deref() + .map(parse_session_list_cursor) + .transpose()? + .unwrap_or(0); + if request + .cwd + .as_ref() + .is_some_and(|requested| requested != &cwd) + { + return Ok(ListSessionsResponse::new(Vec::new())); + } + let root = self.runtime.root().to_path_buf(); + let diagnostics = DiagnosticScope::capture(); + let catalog = tokio::task::spawn_blocking(move || { + diagnostics.sync_scope(|| crate::session::catalog(&root)) + }) + .await + .map_err(|error| { + ListSessionsError::Runtime(AcpRuntimeError::Loop(format!( + "session catalog worker failed: {error}" + ))) + })? + .map_err(|error| ListSessionsError::Runtime(AcpRuntimeError::Loop(error)))?; + if offset > catalog.len() { + return Err(ListSessionsError::InvalidCursor); + } + let end = catalog.len().min(offset + SESSION_LIST_PAGE_SIZE); + let sessions = catalog[offset..end] + .iter() + .map(|entry| { + SessionInfo::new(entry.id.clone(), cwd.clone()) + .title(entry.title.as_deref().map(str::to_owned)) + .updated_at(entry.updated_at_rfc3339()) + }) + .collect(); + let mut response = ListSessionsResponse::new(sessions); + response.next_cursor = (end < catalog.len()).then(|| format!("offset:{end}")); + Ok(response) }) - .collect(); - let mut response = ListSessionsResponse::new(sessions); - response.next_cursor = (end < catalog.len()).then(|| format!("offset:{end}")); - Ok(response) + .await } async fn load_session( @@ -1310,24 +1363,30 @@ impl Server { request: LoadSessionRequest, connection: ConnectionTo, ) -> Result { - let claim = self - .runtime - .claim_session_load(&request.session_id.to_string())?; - let attached = self - .attach_session( - request.cwd, - request.additional_directories, - connection, - claim, - None, - ) - .await?; - let replay = transcript_replay(&attached.session_id, &attached.canonical_transcript); - Ok(PreparedLoad { - response: LoadSessionResponse::new().config_options(Some(attached.config_options)), - replay, - activation: attached.activation, - }) + self.diagnostic_scope(request.meta.as_ref()) + .scope(async { + let claim = self + .runtime + .claim_session_load(&request.session_id.to_string())?; + let attached = self + .attach_session( + request.cwd, + request.additional_directories, + connection, + claim, + None, + ) + .await?; + let replay = + transcript_replay(&attached.session_id, &attached.canonical_transcript); + Ok(PreparedLoad { + response: LoadSessionResponse::new() + .config_options(Some(attached.config_options)), + replay, + activation: attached.activation, + }) + }) + .await } async fn fork_session( @@ -1335,42 +1394,49 @@ impl Server { request: ForkSessionRequest, connection: ConnectionTo, ) -> Result { - if !request.mcp_servers.is_empty() { - return Err(AcpRuntimeError::Loop( - "Kit does not accept per-session MCP servers".into(), - )); - } - let parent_context = fork_parent_context(request.meta.as_ref()); - let sender = self.sender(&request.session_id).await?; - let (tx, rx) = oneshot::channel(); - sender - .send(Command::Fork { - parent_context, - reply: tx, + self.diagnostic_scope(request.meta.as_ref()) + .scope(async { + if !request.mcp_servers.is_empty() { + return Err(AcpRuntimeError::Loop( + "Kit does not accept per-session MCP servers".into(), + )); + } + let parent_context = fork_parent_context(request.meta.as_ref()); + let sender = self.sender(&request.session_id).await?; + let (tx, rx) = oneshot::channel(); + sender + .send(Command::Fork { + diagnostics: DiagnosticScope::capture(), + parent_context, + reply: tx, + }) + .await + .map_err(|_| { + AcpRuntimeError::SessionNotFound(request.session_id.to_string()) + })?; + let forked = rx.await.map_err(|_| { + AcpRuntimeError::SessionNotFound(request.session_id.to_string()) + })??; + let claim = self.runtime.claim_session_fork()?; + let attached = self + .attach_session( + request.cwd, + request.additional_directories, + connection, + claim, + Some(forked), + ) + .await?; + Ok(PreparedFork { + response: ForkSessionResponse::new(attached.session_id) + .config_options(Some(attached.config_options)), + activation: attached.activation, + creation: attached + .pending_fork_creation + .expect("fork attachment must defer transcript creation"), + }) }) .await - .map_err(|_| AcpRuntimeError::SessionNotFound(request.session_id.to_string()))?; - let forked = rx - .await - .map_err(|_| AcpRuntimeError::SessionNotFound(request.session_id.to_string()))??; - let claim = self.runtime.claim_session_fork()?; - let attached = self - .attach_session( - request.cwd, - request.additional_directories, - connection, - claim, - Some(forked), - ) - .await?; - Ok(PreparedFork { - response: ForkSessionResponse::new(attached.session_id) - .config_options(Some(attached.config_options)), - activation: attached.activation, - creation: attached - .pending_fork_creation - .expect("fork attachment must defer transcript creation"), - }) } async fn attach_session( @@ -1471,6 +1537,7 @@ impl Server { .map_err(|error| record_acp_runtime_failure(&session_id, "skill_catalog", error))?; let (tx, rx) = mpsc::channel(8); let actor = SessionActor { + cleanup_diagnostics: DiagnosticScope::capture(), session_id: session_id.clone(), runtime: Arc::clone(&self.runtime), integration: Arc::clone(&self.integration), @@ -1495,13 +1562,15 @@ impl Server { token, completed, }; - let actor_task = tokio::spawn(async move { + // Never inherit attachment ownership into this long-lived actor. Each + // command and task-origin continuation carries its own immutable scope. + let actor_task = tokio::spawn(DiagnosticScope::default().scope(async move { let _shutdown_bridge = shutdown_bridge; let _guard = guard; if activated.await.is_ok() { session_actor(actor).await; } - }); + })); let registered = RegisteredSession { token, session_id: session_id.clone(), @@ -1562,90 +1631,118 @@ impl Server { } async fn prompt(&self, request: PromptRequest) -> Result { - let sender = self.sender(&request.session_id).await?; - let (tx, rx) = oneshot::channel(); - sender - .send(Command::Prompt { request, reply: tx }) + self.diagnostic_scope(request.meta.as_ref()) + .scope(async { + let sender = self.sender(&request.session_id).await?; + let (tx, rx) = oneshot::channel(); + sender + .send(Command::Prompt { + diagnostics: DiagnosticScope::capture(), + request, + reply: tx, + }) + .await + .map_err(|_| AcpRuntimeError::ClientClosed)?; + rx.await.map_err(|_| AcpRuntimeError::ClientClosed)? + }) .await - .map_err(|_| AcpRuntimeError::ClientClosed)?; - rx.await.map_err(|_| AcpRuntimeError::ClientClosed)? } async fn set_config( &self, request: SetSessionConfigOptionRequest, ) -> Result { - let sender = self.sender(&request.session_id).await.map_err(sdk_error)?; - let cancellation_generation = self - .integration - .cancellation_handle(&request.session_id) - .map_err(sdk_error)? - .generation(); - let (tx, rx) = oneshot::channel(); - sender - .send(Command::SetConfig { - request, - reply: tx, - cancellation_generation, + self.diagnostic_scope(request.meta.as_ref()) + .scope(async { + let sender = self.sender(&request.session_id).await.map_err(sdk_error)?; + let cancellation_generation = self + .integration + .cancellation_handle(&request.session_id) + .map_err(sdk_error)? + .generation(); + let (tx, rx) = oneshot::channel(); + sender + .send(Command::SetConfig { + diagnostics: DiagnosticScope::capture(), + request, + reply: tx, + cancellation_generation, + }) + .await + .map_err(|_| sdk_error(AcpRuntimeError::ClientClosed))?; + rx.await + .map_err(|_| sdk_error(AcpRuntimeError::ClientClosed))? }) .await - .map_err(|_| sdk_error(AcpRuntimeError::ClientClosed))?; - rx.await - .map_err(|_| sdk_error(AcpRuntimeError::ClientClosed))? } async fn cancel(&self, notification: CancelNotification) -> Result<(), AcpRuntimeError> { - // Interrupt out of band because the actor may currently be inside - // `driver.next()`. The queued marker preserves command ordering once - // that call settles. - let (sender, background_jobs, tasks, structured_completion) = self - .sessions - .lock() - .map_err(|_| AcpRuntimeError::ClientClosed)? - .get(¬ification.session_id) - .map(|session| { - ( - session.commands.clone(), - session.background_jobs.clone(), - session.tasks.clone(), - session.structured_completion, - ) + self.diagnostic_scope(notification.meta.as_ref()) + .scope(async { + // Interrupt out of band because the actor may currently be inside + // `driver.next()`. The queued marker preserves command ordering once + // that call settles. + let (sender, background_jobs, tasks, structured_completion) = self + .sessions + .lock() + .map_err(|_| AcpRuntimeError::ClientClosed)? + .get(¬ification.session_id) + .map(|session| { + ( + session.commands.clone(), + session.background_jobs.clone(), + session.tasks.clone(), + session.structured_completion, + ) + }) + .ok_or_else(|| { + AcpRuntimeError::SessionNotFound(notification.session_id.to_string()) + })?; + if structured_completion { + cancel_background_jobs(&tasks, &background_jobs).await; + } + self.integration + .interrupt_session(¬ification.session_id)?; + sender + .send(Command::Cancel) + .await + .map_err(|_| AcpRuntimeError::ClientClosed) }) - .ok_or_else(|| AcpRuntimeError::SessionNotFound(notification.session_id.to_string()))?; - if structured_completion { - cancel_background_jobs(&tasks, &background_jobs).await; - } - self.integration - .interrupt_session(¬ification.session_id)?; - sender - .send(Command::Cancel) .await - .map_err(|_| AcpRuntimeError::ClientClosed) } async fn close( &self, request: CloseSessionRequest, ) -> Result { - // Closing uses the same out-of-band interrupt, then waits for the - // actor to reach and acknowledge the serialized close boundary. - let session = self - .sessions - .lock() - .map_err(|_| AcpRuntimeError::ClientClosed)? - .remove(&request.session_id) - .ok_or_else(|| AcpRuntimeError::SessionNotFound(request.session_id.to_string()))?; - cancel_background_jobs(&session.tasks, &session.background_jobs).await; - self.integration.interrupt_session(&request.session_id)?; - let (tx, rx) = oneshot::channel(); - session - .commands - .send(Command::Close { reply: tx }) + self.diagnostic_scope(request.meta.as_ref()) + .scope(async { + // Closing uses the same out-of-band interrupt, then waits for the + // actor to reach and acknowledge the serialized close boundary. + let session = self + .sessions + .lock() + .map_err(|_| AcpRuntimeError::ClientClosed)? + .remove(&request.session_id) + .ok_or_else(|| { + AcpRuntimeError::SessionNotFound(request.session_id.to_string()) + })?; + cancel_background_jobs(&session.tasks, &session.background_jobs).await; + self.integration.interrupt_session(&request.session_id)?; + let (tx, rx) = oneshot::channel(); + session + .commands + .send(Command::Close { + diagnostics: DiagnosticScope::capture(), + reply: tx, + }) + .await + .map_err(|_| AcpRuntimeError::ClientClosed)?; + rx.await.map_err(|_| AcpRuntimeError::ClientClosed)?; + self.registry.remove(session.token); + Ok(CloseSessionResponse::new()) + }) .await - .map_err(|_| AcpRuntimeError::ClientClosed)?; - rx.await.map_err(|_| AcpRuntimeError::ClientClosed)?; - self.registry.remove(session.token); - Ok(CloseSessionResponse::new()) } async fn sender( @@ -1721,6 +1818,7 @@ pub(super) async fn detach_compose_call( } struct SessionActor { + cleanup_diagnostics: DiagnosticScope, session_id: agentkit_acp::SessionId, runtime: Arc, integration: Arc, @@ -1739,6 +1837,7 @@ struct SessionActor { async fn session_actor(actor: SessionActor) { let SessionActor { + cleanup_diagnostics, session_id, runtime, integration, @@ -1762,8 +1861,9 @@ async fn session_actor(actor: SessionActor) { // completion, preventing autonomous progress past that boundary. biased; command = commands.recv() => match command { - Some(Command::Prompt { request, reply }) => { - let result = activity.execute(activity::ExecutionOrigin::Prompt, drive_runtime_prompt( + Some(Command::Prompt { diagnostics, request, reply }) => { + let consumption = crate::runtime::diagnostics::ConsumptionScope::new(Some(diagnostics)); + let result = consumption.scope(activity.execute(activity::ExecutionOrigin::Prompt, drive_runtime_prompt( &session_id, &runtime, &integration, @@ -1773,7 +1873,7 @@ async fn session_actor(actor: SessionActor) { &tasks, &background_jobs, structured_completion, - ), |reason| Some(reason.clone())) + ), |reason| Some(reason.clone()))) .instrument(crate::telemetry::error_spans::operation("acp")) .await; let response = result.and_then(|reason| { @@ -1784,8 +1884,9 @@ async fn session_actor(actor: SessionActor) { // The server already interrupted the shared controller; this // marker only establishes its serialized actor position. Some(Command::Cancel) => {} - Some(Command::SetConfig { request, reply, cancellation_generation }) => { - let result = async { + Some(Command::SetConfig { diagnostics, request, reply, cancellation_generation }) => { + let consumption = crate::runtime::diagnostics::ConsumptionScope::new(Some(diagnostics)); + let result = consumption.scope(async { let cancellation = integration.cancellation_handle(&session_id).map_err(sdk_error)?; if cancellation.is_cancelled_since(cancellation_generation) { return Err(model_switch::error("model change cancelled")); } if request.config_id.to_string() == MODEL_CONFIG_ID { @@ -1812,14 +1913,15 @@ async fn session_actor(actor: SessionActor) { } if cancellation.is_cancelled_since(cancellation_generation) { return Err(model_switch::error("model change cancelled")); } set_config(&adapter, &catalog, request).map_err(sdk_error) - }.await; + }).await; let _ = reply.send(result); } Some(Command::Fork { + diagnostics, parent_context, reply, }) => { - let result = (|| { + let result = diagnostics.scope(async { let mut transcript = driver.snapshot().transcript; crate::transcript::sanitize_forked_transcript(&mut transcript); Ok(AcpForkState { @@ -1830,19 +1932,20 @@ async fn session_actor(actor: SessionActor) { .map_err(AcpRuntimeError::Loop)?, parent_context, }) - })(); + }).await; let _ = reply.send(result); } - Some(Command::Close { reply }) => { - clean_up_session(&session_id, &mut driver, &tasks, &background_jobs).await; - // A close acknowledgement means the actor-owned binding is - // already gone, so callers can immediately reuse the id. - drop(binding.take()); + Some(Command::Close { diagnostics, reply }) => { + diagnostics.scope(async { + clean_up_session(&session_id, &mut driver, &tasks, &background_jobs).await; + // Release the binding before acknowledging the close. + drop(binding.take()); + }).await; let _ = reply.send(()); break; } None => { - clean_up_session(&session_id, &mut driver, &tasks, &background_jobs).await; + cleanup_diagnostics.scope(clean_up_session(&session_id, &mut driver, &tasks, &background_jobs)).await; break; }, }, @@ -1852,12 +1955,12 @@ async fn session_actor(actor: SessionActor) { .submit_input(vec![Item::notification(event.message)]) .map_err(|error| AcpRuntimeError::Loop(error.to_string())); let result = match result { - Ok(()) => drive_unsolicited( + Ok(()) => crate::runtime::diagnostics::ConsumptionScope::new(Some(DiagnosticScope::default())).scope(drive_unsolicited( &session_id, &integration, &mut driver, &activity, - ).await, + )).await, Err(error) => Err(error), }; if let Err(error) = result { @@ -1871,12 +1974,15 @@ async fn session_actor(actor: SessionActor) { Some(TaskEvent::Completed(snapshot, _)) => { background_jobs.acknowledge_terminal(&snapshot.call_id); if snapshot.kind == agentkit_task_manager::TaskKind::Background { - let result = drive_unsolicited( + // Only resolutions actually appended to the loop are + // causes. A terminal wake or a still-live task is not. + let consumption = crate::runtime::diagnostics::ConsumptionScope::new(None); + let result = consumption.scope(drive_unsolicited( &session_id, &integration, &mut driver, &activity, - ).await; + )).await; if let Err(error) = result { eprintln!("autonomous ACP continuation failed for {session_id}: {error}"); } @@ -2396,11 +2502,12 @@ fn component( .on_receive_request( { let state = Arc::clone(&state); - async move |_request: LogoutRequest, responder, cx| { + async move |request: LogoutRequest, responder, cx| { + let diagnostics = state.diagnostic_scope(request.meta.as_ref()); let state = Arc::clone(&state); - cx.spawn(async move { + cx.spawn(diagnostics.scope(async move { responder.respond_with_result(state.logout().await.map_err(sdk_error)) - })?; + }))?; Ok(()) } }, @@ -2412,7 +2519,8 @@ fn component( async move |request: NewSessionRequest, responder, cx| { let state = Arc::clone(&state); let connection = cx.clone(); - cx.spawn(async move { + let diagnostics = state.diagnostic_scope(request.meta.as_ref()); + cx.spawn(diagnostics.scope(async move { let result = state.new_session(request, connection.clone()).await; let notification = result .as_ref() @@ -2423,7 +2531,7 @@ fn component( connection.send_notification(notification)?; } Ok(()) - })?; + }))?; Ok(()) } }, @@ -2436,7 +2544,8 @@ fn component( let state = Arc::clone(&state); let connection = cx.clone(); let session_id = request.session_id.clone(); - cx.spawn(async move { + let diagnostics = state.diagnostic_scope(request.meta.as_ref()); + cx.spawn(diagnostics.scope(async move { match state.load_session(request, connection.clone()).await { Ok(prepared) => { for notification in prepared.replay { @@ -2450,7 +2559,7 @@ fn component( } Err(error) => responder.respond_with_result(Err(sdk_error(error))), } - })?; + }))?; Ok(()) } }, @@ -2462,7 +2571,8 @@ fn component( async move |request: ForkSessionRequest, responder, cx| { let state = Arc::clone(&state); let connection = cx.clone(); - cx.spawn(async move { + let diagnostics = state.diagnostic_scope(request.meta.as_ref()); + cx.spawn(diagnostics.scope(async move { match state.fork_session(request, connection.clone()).await { Ok(prepared) => { if let Some(session_id) = prepared @@ -2475,7 +2585,7 @@ fn component( } Err(error) => responder.respond_with_result(Err(sdk_error(error))), } - })?; + }))?; Ok(()) } }, @@ -2486,14 +2596,15 @@ fn component( let state = Arc::clone(&state); async move |request: ListSessionsRequest, responder, cx| { let state = Arc::clone(&state); - cx.spawn(async move { + let diagnostics = state.diagnostic_scope(request.meta.as_ref()); + cx.spawn(diagnostics.scope(async move { responder.respond_with_result( state .list_sessions(request) .await .map_err(list_sessions_error), ) - })?; + }))?; Ok(()) } }, @@ -2504,10 +2615,11 @@ fn component( let state = Arc::clone(&state); async move |request: PromptRequest, responder, cx| { let state = Arc::clone(&state); - cx.spawn(async move { + let diagnostics = state.diagnostic_scope(request.meta.as_ref()); + cx.spawn(diagnostics.scope(async move { responder .respond_with_result(state.prompt(request).await.map_err(sdk_error)) - })?; + }))?; Ok(()) } }, @@ -2518,9 +2630,10 @@ fn component( let state = Arc::clone(&state); async move |request: SetSessionConfigOptionRequest, responder, cx| { let state = Arc::clone(&state); - cx.spawn(async move { + let diagnostics = state.diagnostic_scope(request.meta.as_ref()); + cx.spawn(diagnostics.scope(async move { responder.respond_with_result(state.set_config(request).await) - })?; + }))?; Ok(()) } }, @@ -2588,9 +2701,10 @@ fn component( let state = Arc::clone(&state); async move |request: CloseSessionRequest, responder, cx| { let state = Arc::clone(&state); - cx.spawn(async move { + let diagnostics = state.diagnostic_scope(request.meta.as_ref()); + cx.spawn(diagnostics.scope(async move { responder.respond_with_result(state.close(request).await.map_err(sdk_error)) - })?; + }))?; Ok(()) } }, @@ -3340,7 +3454,7 @@ pub(super) mod tests { let actor_closed = Arc::clone(&closed); let actor = tokio::spawn(async move { let _completion = CompletionOnDrop(completed); - if let Some(Command::Close { reply }) = received.recv().await { + if let Some(Command::Close { reply, .. }) = received.recv().await { actor_closed.store(true, Ordering::SeqCst); let _ = reply.send(()); } @@ -3405,7 +3519,7 @@ pub(super) mod tests { let actor_closed = Arc::clone(&closed); let actor = tokio::spawn(async move { let _completion = CompletionOnDrop(completed); - if let Some(Command::Close { reply }) = received.recv().await { + if let Some(Command::Close { reply, .. }) = received.recv().await { actor_closed.store(true, Ordering::SeqCst); let _ = reply.send(()); } @@ -4524,6 +4638,35 @@ pub(super) mod tests { } } + struct OwnershipTool { + inner: BlockingTool, + observations: Arc>>, + } + + #[async_trait] + impl Tool for OwnershipTool { + fn spec(&self) -> &ToolSpec { + self.inner.spec() + } + + async fn invoke( + &self, + request: ToolRequest, + ctx: &mut ToolContext<'_>, + ) -> Result { + self.observations + .lock() + .unwrap() + .push(DiagnosticScope::capture()); + let result = self.inner.invoke(request, ctx).await; + self.observations + .lock() + .unwrap() + .push(DiagnosticScope::capture()); + result + } + } + struct FinishBackgroundOnDrop { jobs: BackgroundJobs, call_id: String, @@ -5323,31 +5466,51 @@ pub(super) mod tests { } }); - let task_manager = AsyncTaskManager::new() - .routing(|_request: &agentkit_tools_core::ToolRequest| RoutingDecision::Foreground); - let tasks = task_manager.handle(); let background_jobs = BackgroundJobs::default(); - let tools = ToolRegistry::new().with(BlockingTool { - spec: ToolSpec { - name: ToolName::new(agentkit_tool_compose::COMPOSE_TOOL_NAME), - description: "controlled compose tool".into(), - input_schema: json!({"type": "object", "additionalProperties": false}), - output_schema: None, - annotations: ToolAnnotations::default(), - metadata: MetadataMap::new(), + let task_manager = crate::runtime::diagnostics::DiagnosticTaskManager::new( + AsyncTaskManager::new() + .routing(|_request: &agentkit_tools_core::ToolRequest| RoutingDecision::Foreground), + background_jobs.task_origins.clone(), + ); + let tasks = task_manager.handle(); + let tool_scopes = Arc::new(Mutex::new(Vec::new())); + let turn_scopes = Arc::new(Mutex::new(Vec::new())); + let original = DiagnosticScope::with_operation(crate::events::DiagnosticOperation::new( + "original".into(), + )); + let newer = DiagnosticScope::with_operation(crate::events::DiagnosticOperation::new( + "newer".into(), + )); + let tools = ToolRegistry::new().with(OwnershipTool { + observations: tool_scopes.clone(), + inner: BlockingTool { + spec: ToolSpec { + name: ToolName::new(agentkit_tool_compose::COMPOSE_TOOL_NAME), + description: "controlled compose tool".into(), + input_schema: json!({"type": "object", "additionalProperties": false}), + output_schema: None, + annotations: ToolAnnotations::default(), + metadata: MetadataMap::new(), + }, + entered: Arc::clone(&entered), + release: Arc::clone(&release), }, - entered: Arc::clone(&entered), - release: Arc::clone(&release), }); let (turn_states_tx, mut turn_states_rx) = mpsc::unbounded_channel(); let activity = test_activity(acp_session_id.clone(), turn_states_tx); let driver = Agent::builder() - .model(ScriptAdapter { - turns: Arc::clone(&turns), - user_items_seen: Arc::clone(&user_items_seen), - notification_items_seen: Arc::clone(¬ification_items_seen), + .model(OwnershipAdapter { + inner: ScriptAdapter { + turns: Arc::clone(&turns), + user_items_seen: Arc::clone(&user_items_seen), + notification_items_seen: Arc::clone(¬ification_items_seen), + }, + observations: turn_scopes.clone(), }) .add_tool_source(tools) + .observer(crate::runtime::diagnostics::ConsumptionObserver::new( + background_jobs.task_origins.clone(), + )) .task_manager(task_manager) .observer(integration.as_ref().clone()) .observer(activity.clone()) @@ -5364,7 +5527,26 @@ pub(super) mod tests { let root = tempfile::tempdir().unwrap(); let runtime = Runtime::new(root.path(), "gpt-5.4").unwrap(); let skills = runtime.current_skills().await.unwrap(); + let server = Arc::new(Server::new( + runtime.clone(), + integration.as_ref().clone(), + SessionRegistry::default(), + )); + let mut initialize = InitializeRequest::new(ProtocolVersion::V1); + offer_diagnostic_ownership(&mut initialize.meta); + server.initialize(initialize).await; + server.sessions.lock().unwrap().insert( + acp_session_id.clone(), + SessionHandle { + token: 1, + commands: commands_tx, + background_jobs: background_jobs.clone(), + structured_completion: false, + tasks: tasks.clone(), + }, + ); let actor = tokio::spawn(session_actor(SessionActor { + cleanup_diagnostics: DiagnosticScope::capture(), session_id: acp_session_id.clone(), runtime, integration: Arc::clone(&integration), @@ -5382,19 +5564,15 @@ pub(super) mod tests { mcp_events, })); - let (reply_tx, reply_rx) = oneshot::channel(); - commands_tx - .send(Command::Prompt { - request: PromptRequest::new( - acp_session_id.clone(), - vec![agentkit_acp::ContentBlock::Text( - agentkit_acp::TextContent::new("start one background call"), - )], - ), - reply: reply_tx, - }) - .await - .unwrap(); + let mut request = PromptRequest::new( + acp_session_id.clone(), + vec![ContentBlock::Text(TextContent::new( + "start one background call", + ))], + ); + crate::events::set_diagnostic_operation(&mut request.meta, original.operation().unwrap()); + let prompting = server.clone(); + let reply_rx = tokio::spawn(async move { prompting.prompt(request).await }); while !entered.load(Ordering::SeqCst) { tokio::task::yield_now().await; } @@ -5408,6 +5586,21 @@ pub(super) mod tests { .unwrap(); assert_eq!(turns.load(Ordering::SeqCst), 2); + // Admit a distinct prompt while the original detached producer is live. + // Its eventual tool result and autonomous continuation must stay original. + let mut cancel = CancelNotification::new(acp_session_id.clone()); + crate::events::set_diagnostic_operation(&mut cancel.meta, newer.operation().unwrap()); + server.cancel(cancel).await.unwrap(); + let mut request = PromptRequest::new( + acp_session_id.clone(), + vec![ContentBlock::Text(TextContent::new("new invocation"))], + ); + crate::events::set_diagnostic_operation(&mut request.meta, newer.operation().unwrap()); + timeout(Duration::from_secs(1), server.prompt(request)) + .await + .unwrap() + .unwrap(); + assert_eq!(turns.load(Ordering::SeqCst), 3); release.notify_one(); let completed = timeout(Duration::from_secs(1), async { loop { @@ -5445,10 +5638,23 @@ pub(super) mod tests { assert!(!ended.active); assert_eq!(started.turn_id, ended.turn_id); assert_eq!(started.session_id, acp_session_id); - assert_eq!(turns.load(Ordering::SeqCst), 3); + assert_eq!(turns.load(Ordering::SeqCst), 4); + assert_eq!( + *tool_scopes.lock().unwrap(), + vec![original.clone(), original.clone()] + ); + assert_eq!( + *turn_scopes.lock().unwrap(), + vec![ + original.clone(), + original.clone(), + newer.clone(), + original.clone() + ] + ); assert_eq!( user_items_seen.load(Ordering::SeqCst), - 1, + 2, "autonomous progress must not insert synthetic user content" ); @@ -5462,16 +5668,17 @@ pub(super) mod tests { let mcp_ended = turn_states_rx.recv().await.expect("missing MCP turn end"); assert!(mcp_started.active); assert!(!mcp_ended.active); - assert_eq!(turns.load(Ordering::SeqCst), 4); + assert_eq!(turns.load(Ordering::SeqCst), 5); assert_eq!(notification_items_seen.load(Ordering::SeqCst), 2); - assert_eq!(user_items_seen.load(Ordering::SeqCst), 1); + assert_eq!(user_items_seen.load(Ordering::SeqCst), 2); + assert_eq!( + turn_scopes.lock().unwrap().last(), + Some(&DiagnosticScope::default()) + ); - let (close_tx, close_rx) = oneshot::channel(); - commands_tx - .send(Command::Close { reply: close_tx }) - .await - .unwrap(); - timeout(Duration::from_secs(1), close_rx) + let mut close = CloseSessionRequest::new(acp_session_id.clone()); + crate::events::set_diagnostic_operation(&mut close.meta, newer.operation().unwrap()); + timeout(Duration::from_secs(1), server.close(close)) .await .unwrap() .unwrap(); @@ -5959,6 +6166,232 @@ pub(super) mod tests { } } + #[tokio::test] + async fn ownership_negotiation_requires_exact_offer_and_preserves_legacy_operations() { + for (offer, supported) in [ + (None, false), + (Some(json!({"version": 1, "transport": "stderr"})), true), + (Some(json!({"version": 2, "transport": "stderr"})), false), + (Some(json!({"version": 1, "transport": "stdout"})), false), + ] { + let root = tempfile::tempdir().unwrap(); + let runtime = Runtime::new(root.path(), "gpt-5.4").unwrap(); + let (client_transport, agent_transport) = Channel::duplex(); + let server = tokio::spawn(serve_transport(runtime, agent_transport)); + agent_client_protocol::Client + .builder() + .connect_with(client_transport, async move |connection| { + let mut request = InitializeRequest::new(ProtocolVersion::V1); + request.meta = Some(serde_json::Map::from_iter([( + "unrelated".into(), + json!(true), + )])); + if let Some(offer) = offer { + request + .meta + .as_mut() + .unwrap() + .insert(crate::events::OWNERSHIP_META_KEY.into(), offer); + } + let response = connection.send_request(request).block_task().await?; + assert_eq!(response.protocol_version, ProtocolVersion::V1); + assert_eq!(ownership_supported(response.meta.as_ref()), supported); + let listed = connection + .send_request(ListSessionsRequest::new()) + .block_task() + .await?; + assert!(listed.sessions.is_empty()); + Ok(()) + }) + .await + .unwrap(); + server.abort(); + let _ = server.await; + } + } + + struct OwnershipAdapter { + inner: ScriptAdapter, + observations: Arc>>, + } + + struct OwnershipSession { + inner: ScriptSession, + observations: Arc>>, + } + + #[async_trait] + impl ModelAdapter for OwnershipAdapter { + type Session = OwnershipSession; + + async fn start_session(&self, config: SessionConfig) -> Result { + Ok(OwnershipSession { + inner: self.inner.start_session(config).await?, + observations: self.observations.clone(), + }) + } + } + + #[async_trait] + impl ModelSession for OwnershipSession { + type Turn = ScriptTurn; + + async fn begin_turn( + &mut self, + request: TurnRequest, + cancellation: Option, + ) -> Result { + // TurnStarted precedes actual background-result consumption. The + // model boundary proves the downstream producer, not an early wake. + self.observations + .lock() + .unwrap() + .push(DiagnosticScope::capture()); + self.inner.begin_turn(request, cancellation).await + } + } + + #[derive(Clone)] + struct OwnershipObserver(Arc>>); + + impl LoopObserver for OwnershipObserver { + fn handle_event(&self, event: ObservedEvent) { + if matches!(event.event, AgentEvent::TurnStarted { .. }) { + self.0.lock().unwrap().push(DiagnosticScope::capture()); + } + } + } + + #[tokio::test] + async fn prompt_producers_use_negotiated_request_scope_across_actor_handoff() { + let root = tempfile::tempdir().unwrap(); + let runtime = Runtime::new(root.path(), "gpt-5.4").unwrap(); + let integration = AcpIntegration::builder() + .name("ownership-test") + .approval_resolver(AutoDenyResolver) + .build() + .unwrap(); + let server = Arc::new(Server::new( + runtime.clone(), + integration, + SessionRegistry::default(), + )); + let session_id = agentkit_acp::SessionId::new("s-ownership-actor"); + let (client, mut messages) = AcpClientHandle::channel(); + let drain = tokio::spawn(async move { + while let Some(message) = messages.recv().await { + if let AcpClientMessage::Flush { response } = message { + let _ = response.send(()); + } + } + }); + let cancellation = CancellationController::new(); + server + .integration + .bind_session( + AcpSessionBinding::new( + session_id.clone(), + AgentkitSessionId::new(session_id.to_string()), + client, + ) + .cancellation(cancellation.clone()), + ) + .unwrap(); + let observations = Arc::new(Mutex::new(Vec::new())); + let manager = AsyncTaskManager::new(); + let tasks = manager.handle(); + let background_jobs = BackgroundJobs::default(); + let driver = Agent::builder() + .model(ScriptAdapter { + // Start after the scripted tool turns: each prompt produces one model turn. + turns: Arc::new(AtomicUsize::new(2)), + user_items_seen: Arc::new(AtomicUsize::new(0)), + notification_items_seen: Arc::new(AtomicUsize::new(0)), + }) + .task_manager(manager) + .observer(server.integration.as_ref().clone()) + .observer(OwnershipObserver(observations.clone())) + .cancellation(cancellation.handle()) + .build() + .unwrap() + .start( + SessionConfig::new(AgentkitSessionId::new(session_id.to_string())).without_cache(), + ) + .await + .unwrap(); + let (commands, received) = mpsc::channel(8); + server.sessions.lock().unwrap().insert( + session_id.clone(), + SessionHandle { + token: 1, + commands, + background_jobs: background_jobs.clone(), + structured_completion: false, + tasks: tasks.clone(), + }, + ); + let skills = runtime.current_skills().await.unwrap(); + let mcp = crate::tools::mcp::empty(); + let actor = tokio::spawn(session_actor(SessionActor { + cleanup_diagnostics: DiagnosticScope::default(), + session_id: session_id.clone(), + runtime, + integration: server.integration.clone(), + binding: SessionBindingGuard::new(server.integration.clone(), session_id.clone()), + driver, + tasks, + background_jobs, + structured_completion: false, + skill_catalog: skill_catalog::SkillCatalogMonitor::new(&skills.skills).unwrap(), + adapter: SelectableAdapter::new(crate::ProviderKind::OpenAiSubscription, "gpt-5.4") + .unwrap(), + catalog: Vec::new(), + commands: received, + activity: legacy_activity(session_id.clone(), |_| Ok(())), + mcp_events: mcp.subscribe(session_id.to_string()), + })); + for negotiated in [true, false] { + let mut initialize = InitializeRequest::new(ProtocolVersion::V1); + if negotiated { + offer_diagnostic_ownership(&mut initialize.meta); + } + server.initialize(initialize).await; + // Reuse the actor across A→B→A and A→A; malformed/missing metadata + // must clear an ambient owner rather than silently inheriting it. + for token in [ + json!("A"), + json!("B"), + json!("A"), + json!("A"), + json!(42), + json!(null), + ] { + let mut request = PromptRequest::new( + session_id.clone(), + vec![ContentBlock::Text(TextContent::new("next"))], + ); + request.meta = Some(serde_json::Map::from_iter([( + crate::events::OPERATION_META_KEY.into(), + token, + )])); + let expected = server.diagnostic_scope(request.meta.as_ref()); + let ambient = DiagnosticScope::with_operation( + crate::events::DiagnosticOperation::new("ambient".into()), + ); + let result = ambient.scope(server.prompt(request)).await.unwrap(); + assert_eq!(result.stop_reason, StopReason::EndTurn); + assert_eq!(observations.lock().unwrap().pop(), Some(expected)); + assert!(observations.lock().unwrap().is_empty()); + } + } + server + .close(CloseSessionRequest::new(session_id)) + .await + .unwrap(); + actor.await.unwrap(); + drain.abort(); + } + #[tokio::test] async fn kit_server_advertises_supported_session_discovery_restoration_and_forking() { let root = tempfile::tempdir().unwrap(); diff --git a/src/protocols/acp/v2.rs b/src/protocols/acp/v2.rs index 218d61b..c78425d 100644 --- a/src/protocols/acp/v2.rs +++ b/src/protocols/acp/v2.rs @@ -29,8 +29,9 @@ use tokio::sync::{Notify, mpsc, oneshot, watch}; use tracing::Instrument as _; use crate::{ + events::{self, DiagnosticScope}, provider::{ProviderKind, SelectableAdapter, authentication_method_id}, - runtime::{AcpDriverContext, BackgroundJobs, InputSettlingDriver as LoopDriver, Runtime}, + runtime::{AcpDriverContext, BackgroundJobs, InputSettlingDriver as LoopDriver, Runtime, diagnostics::ConsumptionScope}, }; use super::activity::{ExecutionOrigin, SessionActivity}; @@ -180,10 +181,28 @@ struct ConnectionSink(V2ConnectionTo, Arc>); struct InjectionWork { admitting: usize, pending: HashSet, + origins: HashMap, branch_reserved: bool, admission_released: Arc, } +impl InjectionWork { + fn forget(&mut self, id: &wire::MessageId) { + self.pending.remove(id); + self.origins.remove(id); + } + + fn delivered(&mut self, notification: &wire::UpdateSessionNotification) { + if let wire::SessionUpdate::UserMessage(message) = ¬ification.update + && self.pending.remove(&message.message_id) + { + // The SDK has submitted this exact queued input and cannot begin + // its next model step until this acknowledgement returns. + ConsumptionScope::consumed(self.origins.remove(&message.message_id)); + } + } +} + struct InjectionAdmission(Arc>); impl Drop for InjectionAdmission { @@ -204,8 +223,7 @@ impl Drop for TrackedInjection { self.work .lock() .expect("injection work poisoned") - .pending - .remove(&self.id); + .forget(&self.id); } } } @@ -248,20 +266,13 @@ impl Drop for BranchAdmission { #[async_trait] impl AcpSessionUpdateSink for ConnectionSink { fn update(&self, notification: wire::UpdateSessionNotification) -> Result<(), AcpRuntimeError> { - let delivered = match ¬ification.update { - wire::SessionUpdate::UserMessage(message) => Some(message.message_id.clone()), - _ => None, - }; self.0 - .send_notification(notification) + .send_notification(notification.clone()) .map_err(|error| AcpRuntimeError::Sdk(error.to_string()))?; - if let Some(id) = delivered { - self.1 - .lock() - .expect("injection work poisoned") - .pending - .remove(&id); - } + self.1 + .lock() + .expect("injection work poisoned") + .delivered(¬ification); Ok(()) } @@ -662,6 +673,10 @@ struct PromptCommand { } enum Command { + Scoped { + origin: DiagnosticScope, + command: Box, + }, Snapshot { reply: oneshot::Sender>, }, @@ -691,6 +706,22 @@ enum Command { }, } +impl Command { + fn scoped(self) -> Self { + Self::Scoped { + origin: DiagnosticScope::capture(), + command: Box::new(self), + } + } + + fn into_scoped(self) -> (DiagnosticScope, Self) { + match self { + Self::Scoped { origin, command } => (origin, *command), + command => (DiagnosticScope::default(), command), + } + } +} + struct SessionHandle { injections: Arc>, token: u64, @@ -785,6 +816,7 @@ struct PendingSessionPublication { } struct Server { + diagnostic_ownership: AtomicBool, runtime: Arc, integration: Arc, registry: SessionRegistry, @@ -795,6 +827,7 @@ struct Server { impl Server { fn new(runtime: Arc, registry: SessionRegistry) -> Self { Self { + diagnostic_ownership: AtomicBool::new(false), runtime, integration: Arc::new(AcpIntegration::default()), registry, @@ -912,7 +945,10 @@ impl Server { "ACP v2 requires protocol version 2 or newer".into(), )); } - Ok(wire::InitializeResponse::new( + let negotiated = events::ownership_supported(request.meta.as_ref()); + self.diagnostic_ownership + .store(negotiated, Ordering::Release); + let mut response = wire::InitializeResponse::new( wire::ProtocolVersion::V2, wire::Implementation::new("kit", env!("CARGO_PKG_VERSION")), ) @@ -923,7 +959,21 @@ impl Server { }) } else { Vec::new() - })) + }); + if negotiated { + events::offer_diagnostic_ownership(&mut response.meta); + } + Ok(response) + } + + fn diagnostic_scope( + &self, + meta: Option<&serde_json::Map>, + ) -> DiagnosticScope { + DiagnosticScope::with_operation(events::diagnostic_operation( + meta, + self.diagnostic_ownership.load(Ordering::Acquire), + )) } async fn new_session( @@ -1187,7 +1237,7 @@ impl Server { token, completed, }; - let actor_task = tokio::spawn(async move { + let actor_task = tokio::spawn(events::inherit_diagnostics(async move { let _guard = guard; if activated.await.is_ok() { session_actor(actor).await; @@ -1195,7 +1245,7 @@ impl Server { // Abandoning response/replay retires the queued input as well. actor.busy.store(false, Ordering::Release); } - }); + })); let interrupt_handle = handle.clone(); let interrupt_background_jobs = background_jobs.clone(); let interrupt = Arc::new(move || { @@ -1211,15 +1261,19 @@ impl Server { let close_background_jobs = close_background_jobs.clone(); let close_tasks = close_tasks.clone(); let weak = weak.clone(); - Box::pin(async move { + Box::pin(events::inherit_diagnostics(async move { super::cancel_background_jobs(&close_tasks, &close_background_jobs).await; if let Some(commands) = weak.upgrade() { let (reply, acknowledged) = oneshot::channel(); - if commands.send(Command::Close { reply }).await.is_ok() { + if commands + .send(Command::Close { reply }.scoped()) + .await + .is_ok() + { let _ = acknowledged.await; } } - }) as std::pin::Pin + Send>> + })) as std::pin::Pin + Send>> }); // Registration, durable identity commit, and local publication are one // critical section so logout cannot reopen admission around a dead actor. @@ -1282,7 +1336,7 @@ impl Server { }; let (reply, response) = oneshot::channel(); sender - .send(Command::Snapshot { reply }) + .send(Command::Snapshot { reply }.scoped()) .await .map_err(|_| AcpRuntimeError::ClientClosed)?; let SessionSnapshot { @@ -1323,7 +1377,7 @@ impl Server { let (sender, _, _) = self.prompt_route(&request.session_id)?; let (reply, response) = oneshot::channel(); sender - .send(Command::ListPromptBranches { reply }) + .send(Command::ListPromptBranches { reply }.scoped()) .await .map_err(|_| AcpRuntimeError::ClientClosed)?; response.await.map_err(|_| AcpRuntimeError::ClientClosed)? @@ -1336,11 +1390,14 @@ impl Server { let (sender, admission) = self.branch_admission(&request.session_id)?; let (reply, response) = oneshot::channel(); sender - .send(Command::PreparePromptBranch { - address: request.address, - admission, - reply, - }) + .send( + Command::PreparePromptBranch { + address: request.address, + admission, + reply, + } + .scoped(), + ) .await .map_err(|_| AcpRuntimeError::ClientClosed)?; response.await.map_err(|_| AcpRuntimeError::ClientClosed)? @@ -1400,11 +1457,14 @@ impl Server { let (sender, admission) = self.branch_admission(&request.session_id)?; let (reply, response) = oneshot::channel(); sender - .send(Command::ReservePromptBranch { - checkout_token: request.checkout_token, - admission, - reply, - }) + .send( + Command::ReservePromptBranch { + checkout_token: request.checkout_token, + admission, + reply, + } + .scoped(), + ) .await .map_err(|_| AcpRuntimeError::ClientClosed)?; let (checkout, release) = response @@ -1454,6 +1514,37 @@ impl Server { Ok(InjectionAdmission(session.injections.clone())) } + async fn replace_injection( + &self, + request: wire::ReplaceInjectSessionRequest, + ) -> Result { + let tracker = self + .sessions + .lock() + .expect("ACP v2 session map poisoned") + .get(&request.session_id) + .map(|session| session.injections.clone()); + let Some(tracker) = tracker else { + return self.integration.replace_inject(request).await; + }; + let id = request.message_id.clone(); + let origin = DiagnosticScope::capture(); + let replacement = self.integration.replace_inject(request); + tokio::pin!(replacement); + std::future::poll_fn(|cx| { + // Pinned SDK replacement never calls the update sink. Its successful + // transition and our scope update must be atomic with that sink; + // release this guard after EACH poll, including Pending delivery. + let mut work = tracker.lock().expect("injection work poisoned"); + let result = std::future::Future::poll(replacement.as_mut(), cx); + if matches!(&result, std::task::Poll::Ready(Ok(_))) && work.pending.contains(&id) { + work.origins.insert(id.clone(), origin.clone()); + } + result + }) + .await + } + async fn prepare_prompt( &self, request: wire::PromptRequest, @@ -1475,11 +1566,14 @@ impl Server { handle.prepare_injection_turn(); let cancellation_generation = handle.cancellation_handle().generation(); let (reply, response) = oneshot::channel(); - permit.send(Command::Prompt(PromptCommand { - request, - cancellation_generation, - reply, - })); + permit.send( + Command::Prompt(PromptCommand { + request, + cancellation_generation, + reply, + }) + .scoped(), + ); match response.await { Ok(response) => response, Err(_) => { @@ -1530,11 +1624,14 @@ impl Server { }; let (reply, response) = oneshot::channel(); sender - .send(Command::SetConfig { - request, - reply, - cancellation_generation, - }) + .send( + Command::SetConfig { + request, + reply, + cancellation_generation, + } + .scoped(), + ) .await .map_err(|_| sdk_error(AcpRuntimeError::ClientClosed))?; response @@ -1546,27 +1643,31 @@ impl Server { &self, notification: wire::CancelSessionNotification, ) -> Result<(), AcpRuntimeError> { - let session = self - .sessions - .lock() - .map_err(|_| AcpRuntimeError::ClientClosed)? - .get(¬ification.session_id) - .map(|session| { - ( - session.integration.clone(), - session.background_jobs.clone(), - session.tasks.clone(), - session.structured_completion, - ) - }); - if let Some((handle, background_jobs, tasks, structured_completion)) = session { - // Interrupt admission before waiting for asynchronous task cleanup. - handle.interrupt(); - if structured_completion { - super::cancel_background_jobs(&tasks, &background_jobs).await; - } - } - Ok(()) + self.diagnostic_scope(notification.meta.as_ref()) + .scope(async { + let session = self + .sessions + .lock() + .map_err(|_| AcpRuntimeError::ClientClosed)? + .get(¬ification.session_id) + .map(|session| { + ( + session.integration.clone(), + session.background_jobs.clone(), + session.tasks.clone(), + session.structured_completion, + ) + }); + if let Some((handle, background_jobs, tasks, structured_completion)) = session { + // Interrupt admission before waiting for asynchronous task cleanup. + handle.interrupt(); + if structured_completion { + super::cancel_background_jobs(&tasks, &background_jobs).await; + } + } + Ok(()) + }) + .await } async fn close( @@ -1586,7 +1687,7 @@ impl Server { // command arrives. A dropped receiver is already a completed close. if session .commands - .send(Command::Close { reply }) + .send(Command::Close { reply }.scoped()) .await .is_ok() { @@ -1787,20 +1888,22 @@ async fn session_actor driver = active, stopped => { @@ -1811,7 +1914,7 @@ async fn session_actor match command { + command = commands.recv() => { + let (origin, command) = match command { + Some(command) => { let (origin, command) = command.into_scoped(); (origin, Some(command)) }, + None => (DiagnosticScope::capture(), None), + }; + let consumption = if autonomous_external && matches!(&command, Some(Command::Prompt(_))) { + ConsumptionScope::new(Some(DiagnosticScope::default())) + } else { ConsumptionScope::explicit(origin) }; + let stop = consumption.scope(async { + match command { + Some(Command::Scoped { .. }) => unreachable!("command envelopes are not nested"), Some(Command::Snapshot { reply }) => { let result = (|| { let selection = adapter.selection().map_err(AcpRuntimeError::Loop)?; @@ -1868,7 +1981,7 @@ async fn session_actor { hold_branch_reservation(checkout, admission, reply).await; - continue; + return false; } Err(error) => { let _ = reply.send(Err(error)); } } @@ -1893,6 +2006,7 @@ async fn session_actor { @@ -1951,7 +2065,7 @@ async fn session_actor { @@ -1959,24 +2073,30 @@ async fn session_actor { let v1_id = agentkit_acp::SessionId::new(session_id.to_string()); super::clean_up_session(&v1_id, &mut driver, &tasks, &background_jobs).await; - break; + return true; } + } + false + }).await; + if stop { break; } }, _ = admission_released.notified(), if autonomous_pending => {} _ = std::future::ready(()), if autonomous_pending && !busy.load(Ordering::Acquire) => { let generation = handle.cancellation_handle().generation(); - let result = drive_autonomous( + let consumption = ConsumptionScope::new(autonomous_external.then(DiagnosticScope::default)); + let result = consumption.scope(drive_autonomous( &session_id, &integration, &handle, &busy, &mut driver, &sink, &activity, - ) + )) .instrument(crate::telemetry::error_spans::operation("acp_autonomous")) .await; autonomous_pending = matches!(result, Ok(false)); + if !autonomous_pending { autonomous_external = false; } if let Err(error) = result { eprintln!("ACP v2 autonomous turn failed for {session_id}: {error}"); } @@ -1985,7 +2105,7 @@ async fn session_actor { + autonomous_external = true; autonomous_pending = true; } Err(error) => { @@ -2427,40 +2548,42 @@ async fn run_active_turn( // after the client returns to its source without sending session/resume. // Re-establish identity before activity, model, or cleanup diagnostics. establish_diagnostic_route(session_id); - crate::events::scope_diagnostics( - &session_id.to_string(), - activity.execute( - origin, - async { - let result = drive_prompt( - session_id, - driver, - handle, - cancellation_generation, - structured, - &sink.acknowledgements, - ) - .await; - let outcome = super::activity::ExecutionOutcome::new( - result, - handle - .cancellation_handle() - .is_cancelled_since(cancellation_generation), - ); - handle.stop_injection_turn(); - let result = super::activity::finalize( - outcome, - structured, - integration.flush_session_updates(session_id), - |error| sink.update(error_diagnostic_notification(session_id, error)), - ) - .await; - integration.finish_prompt(session_id); - result - }, - |reason| Some(reason.clone()), - ), - ) + crate::events::scope_diagnostics(&session_id.to_string(), async { + ConsumptionScope::bind_local_seed(); + activity + .execute( + origin, + async { + let result = drive_prompt( + session_id, + driver, + handle, + cancellation_generation, + structured, + &sink.acknowledgements, + ) + .await; + let outcome = super::activity::ExecutionOutcome::new( + result, + handle + .cancellation_handle() + .is_cancelled_since(cancellation_generation), + ); + handle.stop_injection_turn(); + let result = super::activity::finalize( + outcome, + structured, + integration.flush_session_updates(session_id), + |error| sink.update(error_diagnostic_notification(session_id, error)), + ) + .await; + integration.finish_prompt(session_id); + result + }, + |reason| Some(reason.clone()), + ) + .await + }) .await } @@ -2827,7 +2950,10 @@ pub(crate) fn component( runtime: Arc, registry: SessionRegistry, ) -> Result, AcpRuntimeError> { - let state = Arc::new(Server::new(runtime, registry)); + component_from_state(Arc::new(Server::new(runtime, registry))) +} + +fn component_from_state(state: Arc) -> Result, AcpRuntimeError> { let agent = agent_client_protocol::Agent .v2() .name("kit") @@ -2844,7 +2970,10 @@ pub(crate) fn component( { let state = Arc::clone(&state); async move |request: wire::LoginAuthRequest, responder, _cx| { - responder.respond_with_result(state.login(request)) + let diagnostic_scope = state.diagnostic_scope(request.meta.as_ref()); + diagnostic_scope + .scope(async { responder.respond_with_result(state.login(request)) }) + .await } }, agent_client_protocol::on_receive_request!(), @@ -2853,10 +2982,11 @@ pub(crate) fn component( { let state = Arc::clone(&state); async move |_request: wire::LogoutAuthRequest, responder, cx| { + let diagnostic_scope = state.diagnostic_scope(_request.meta.as_ref()); let state = Arc::clone(&state); - cx.spawn(async move { + cx.spawn(diagnostic_scope.scope(async move { responder.respond_with_result(state.logout().await.map_err(sdk_error)) - })?; + }))?; Ok(()) } }, @@ -2866,9 +2996,10 @@ pub(crate) fn component( { let state = Arc::clone(&state); async move |request: wire::NewSessionRequest, responder, cx| { + let diagnostic_scope = state.diagnostic_scope(request.meta.as_ref()); let state = Arc::clone(&state); let connection = cx.clone(); - cx.spawn(async move { + cx.spawn(diagnostic_scope.scope(async move { match state.new_session(request, connection.clone()).await { Ok((response, activation)) => { let session_id = response.session_id.clone(); @@ -2878,7 +3009,7 @@ pub(crate) fn component( } Err(error) => responder.respond_with_error(sdk_error(error)), } - })?; + }))?; Ok(()) } }, @@ -2888,15 +3019,16 @@ pub(crate) fn component( { let state = Arc::clone(&state); async move |request: wire::ListSessionsRequest, responder, cx| { + let diagnostic_scope = state.diagnostic_scope(request.meta.as_ref()); let state = Arc::clone(&state); - cx.spawn(async move { + cx.spawn(diagnostic_scope.scope(async move { responder.respond_with_result( state .list_sessions(request) .await .map_err(list_sessions_error), ) - })?; + }))?; Ok(()) } }, @@ -2906,10 +3038,11 @@ pub(crate) fn component( { let state = Arc::clone(&state); async move |request: wire::ResumeSessionRequest, responder, cx| { + let diagnostic_scope = state.diagnostic_scope(request.meta.as_ref()); let state = Arc::clone(&state); let connection = cx.clone(); let session_id = request.session_id.clone(); - cx.spawn(async move { + cx.spawn(diagnostic_scope.scope(async move { match state.resume_session(request, connection.clone()).await { Ok((response, replay, activation)) => { for update in replay { @@ -2921,7 +3054,7 @@ pub(crate) fn component( } Err(error) => responder.respond_with_error(sdk_error(error)), } - })?; + }))?; Ok(()) } }, @@ -2931,12 +3064,13 @@ pub(crate) fn component( { let state = Arc::clone(&state); async move |request: ListPromptBranchesRequest, responder, cx| { + let diagnostic_scope = DiagnosticScope::with_operation(None); let state = Arc::clone(&state); - cx.spawn(async move { + cx.spawn(diagnostic_scope.scope(async move { responder.respond_with_result( state.list_prompt_branches(request).await.map_err(sdk_error), ) - })?; + }))?; Ok(()) } }, @@ -2946,15 +3080,16 @@ pub(crate) fn component( { let state = Arc::clone(&state); async move |request: PreparePromptBranchRequest, responder, cx| { + let diagnostic_scope = DiagnosticScope::with_operation(None); let state = Arc::clone(&state); - cx.spawn(async move { + cx.spawn(diagnostic_scope.scope(async move { responder.respond_with_result( state .prepare_prompt_branch(request) .await .map_err(sdk_error), ) - })?; + }))?; Ok(()) } }, @@ -2966,9 +3101,10 @@ pub(crate) fn component( async move |request: SubmitPromptBranchRequest, responder: Responder, cx| { + let diagnostic_scope = DiagnosticScope::with_operation(None); let state = Arc::clone(&state); let connection = cx.clone(); - cx.spawn(async move { + cx.spawn(diagnostic_scope.scope(async move { match state .submit_prompt_branch(request, connection.clone()) .await @@ -3007,7 +3143,7 @@ pub(crate) fn component( } Err(error) => responder.respond_with_error(sdk_error(error)), } - })?; + }))?; Ok(()) } }, @@ -3017,8 +3153,9 @@ pub(crate) fn component( { let state = Arc::clone(&state); async move |request: wire::PromptRequest, responder, cx| { + let diagnostic_scope = state.diagnostic_scope(request.meta.as_ref()); let state = Arc::clone(&state); - cx.spawn(async move { + cx.spawn(diagnostic_scope.scope(async move { match state.prepare_prompt(request).await { Ok(start) => { responder.respond(wire::PromptResponse::new())?; @@ -3027,7 +3164,7 @@ pub(crate) fn component( } Err(error) => responder.respond_with_error(sdk_error(error)), } - })?; + }))?; Ok(()) } }, @@ -3037,10 +3174,11 @@ pub(crate) fn component( { let state = Arc::clone(&state); async move |request: wire::SetSessionConfigOptionRequest, responder, cx| { + let diagnostic_scope = state.diagnostic_scope(request.meta.as_ref()); let state = Arc::clone(&state); - cx.spawn(async move { + cx.spawn(diagnostic_scope.scope(async move { responder.respond_with_result(state.set_config(request).await) - })?; + }))?; Ok(()) } }, @@ -3052,37 +3190,42 @@ pub(crate) fn component( async move |request: wire::InjectSessionRequest, responder: Responder, cx| { + let diagnostic_scope = state.diagnostic_scope(request.meta.as_ref()); + let session_id = request.session_id.to_string(); let admission = match state.injection_admission(&request.session_id) { Ok(admission) => admission, - Err(error) => return responder.respond_with_error(sdk_error(error)), + Err(error) => { + return responder.respond_with_error(sdk_error(error)); + } }; let integration = Arc::clone(&state.integration); - cx.spawn(async move { - let Some(reserved) = integration - .reserve_inject_request(request, responder) - .await? - else { - return Ok(()); - }; - let id = reserved.response().message_id; - admission - .0 - .lock() - .expect("injection work poisoned") - .pending - .insert(id.clone()); - let mut pending = TrackedInjection { - work: admission.0.clone(), - id, - retained: false, - }; - if let Some(acceptance) = reserved.respond_tracked()? { - acceptance.activate_after_response().await?; - pending.retained = true; - } - drop(admission); - Ok(()) - })?; + cx.spawn(diagnostic_scope.sync_scope(|| { + events::scope_diagnostics(&session_id, async move { + let Some(reserved) = integration + .reserve_inject_request(request, responder) + .await? + else { + return Ok(()); + }; + let id = reserved.response().message_id; + { + let mut work = admission.0.lock().expect("injection work poisoned"); + work.pending.insert(id.clone()); + work.origins.insert(id.clone(), DiagnosticScope::capture()); + } + let mut pending = TrackedInjection { + work: admission.0.clone(), + id, + retained: false, + }; + if let Some(acceptance) = reserved.respond_tracked()? { + acceptance.activate_after_response().await?; + pending.retained = true; + } + drop(admission); + Ok(()) + }) + }))?; Ok(()) } }, @@ -3092,36 +3235,41 @@ pub(crate) fn component( { let state = Arc::clone(&state); async move |request: wire::RevokeInjectSessionRequest, responder, _cx| { - let tracker = state - .sessions - .lock() - .expect("ACP v2 session map poisoned") - .get(&request.session_id) - .map(|session| session.injections.clone()); - let id = request.message_id.clone(); - let result = state.integration.revoke_inject(request).await; - if result.is_ok() - && let Some(tracker) = tracker - { - tracker - .lock() - .expect("injection work poisoned") - .pending - .remove(&id); - } - responder.respond_with_result(result) + let diagnostic_scope = state.diagnostic_scope(request.meta.as_ref()); + diagnostic_scope + .scope(async { + let tracker = state + .sessions + .lock() + .expect("ACP v2 session map poisoned") + .get(&request.session_id) + .map(|session| session.injections.clone()); + let id = request.message_id.clone(); + let result = state.integration.revoke_inject(request).await; + if result.is_ok() + && let Some(tracker) = tracker + { + tracker.lock().expect("injection work poisoned").forget(&id); + } + responder.respond_with_result(result) + }) + .await } }, agent_client_protocol::on_receive_request!(), ) .on_receive_request( { - let integration = Arc::clone(&state.integration); + let state = Arc::clone(&state); async move |request: wire::ReplaceInjectSessionRequest, responder, cx| { - let integration = Arc::clone(&integration); - cx.spawn(async move { - responder.respond_with_result(integration.replace_inject(request).await) - })?; + let diagnostic_scope = state.diagnostic_scope(request.meta.as_ref()); + let state = Arc::clone(&state); + let session_id = request.session_id.to_string(); + cx.spawn(diagnostic_scope.sync_scope(|| { + events::scope_diagnostics(&session_id, async move { + responder.respond_with_result(state.replace_injection(request).await) + }) + }))?; Ok(()) } }, @@ -3131,12 +3279,13 @@ pub(crate) fn component( { let state = Arc::clone(&state); async move |request: DetachComposeRequest, responder, cx| { + let diagnostic_scope = DiagnosticScope::with_operation(None); let state = Arc::clone(&state); - cx.spawn(async move { + cx.spawn(diagnostic_scope.scope(async move { responder.respond_with_result( state.detach_compose(request).await.map_err(sdk_error), ) - })?; + }))?; Ok(()) } }, @@ -3146,8 +3295,14 @@ pub(crate) fn component( { let state = Arc::clone(&state); async move |request: CancelBackgroundRequest, responder, _cx| { - responder - .respond_with_result(state.cancel_background(request).map_err(sdk_error)) + let diagnostic_scope = DiagnosticScope::with_operation(None); + diagnostic_scope + .scope(async { + responder.respond_with_result( + state.cancel_background(request).map_err(sdk_error), + ) + }) + .await } }, agent_client_protocol::on_receive_request!(), @@ -3156,14 +3311,15 @@ pub(crate) fn component( { let state = Arc::clone(&state); async move |request: FileSearchRequest, responder, cx| { + let diagnostic_scope = DiagnosticScope::with_operation(None); let state = Arc::clone(&state); - cx.spawn(async move { + cx.spawn(diagnostic_scope.scope(async move { responder.respond_with_result( state.search_files(request).await.map_err(|error| { agent_client_protocol::util::internal_error(error) }), ) - })?; + }))?; Ok(()) } }, @@ -3183,10 +3339,11 @@ pub(crate) fn component( { let state = Arc::clone(&state); async move |request: wire::CloseSessionRequest, responder, cx| { + let diagnostic_scope = state.diagnostic_scope(request.meta.as_ref()); let state = Arc::clone(&state); - cx.spawn(async move { + cx.spawn(diagnostic_scope.scope(async move { responder.respond_with_result(state.close(request).await.map_err(sdk_error)) - })?; + }))?; Ok(()) } }, @@ -4112,7 +4269,7 @@ mod tests { .await .unwrap(); let child_id = submitted.session_id; - let retired_commands = state.sender(&child_id).unwrap(); + let retired_commands = state.prompt_route(&child_id).unwrap().0; state .cancel(wire::CancelSessionNotification::new(child_id.clone())) .await @@ -4409,6 +4566,10 @@ mod tests { .unwrap() .pending .extend([failed.clone(), retained.clone()]); + work.lock().unwrap().origins.extend([ + (failed.clone(), DiagnosticScope::default()), + (retained.clone(), DiagnosticScope::default()), + ]); drop(TrackedInjection { work: work.clone(), id: failed.clone(), @@ -4419,7 +4580,13 @@ mod tests { id: retained.clone(), retained: true, }); - assert_eq!(work.lock().unwrap().pending, HashSet::from([retained])); + assert_eq!( + work.lock().unwrap().pending, + HashSet::from([retained.clone()]) + ); + assert!(!work.lock().unwrap().origins.contains_key(&failed)); + assert_eq!(work.lock().unwrap().origins.len(), 1); + assert!(work.lock().unwrap().origins.contains_key(&retained)); } #[tokio::test] @@ -4459,99 +4626,812 @@ mod tests { assert!(!work.lock().unwrap().branch_reserved); } - #[derive(Clone, Copy)] - enum RacingBranch { - Prepare, - Reserve, - Abandon, - } + type InjectionDeliveryGate = (oneshot::Sender<()>, oneshot::Receiver<()>); - async fn selected_completion_survives_branch_admission(background: bool, branch: RacingBranch) { - use agentkit_core::{TaskId, ToolCallId, TurnId}; - use agentkit_task_manager::{ContinuePolicy, TaskLaunchRequest, TaskStartContext}; - use agentkit_tools_core::{ - AllowAllPermissions, BasicToolExecutor, OwnedToolContext, ToolRequest, ToolSource, - }; + #[derive(Clone)] + struct InjectionDeliverySink { + work: Arc>, + delivered: Arc>>, + gate: Arc>>, + } - let root = tempfile::tempdir().unwrap(); - let runtime = Runtime::new(root.path(), "gpt-5.4").unwrap(); - let session_id = wire::SessionId::new("selected-completion"); - let loop_id = SessionId::new("selected-completion"); - let integration = Arc::new(AcpIntegration::default()); - let recording = RecordingSink::default(); - let sink = ResponseReplacementSink::new(recording.clone()); - let generations = Arc::new(Mutex::new(Vec::new())); - let activity = SessionActivity::new({ - let generations = generations.clone(); - move |transition| { - generations - .lock() - .unwrap() - .push((transition.id, transition.active)); - Ok(()) + #[async_trait] + impl AcpSessionUpdateSink for InjectionDeliverySink { + fn update( + &self, + notification: wire::UpdateSessionNotification, + ) -> Result<(), AcpRuntimeError> { + // Same production consumption boundary as ConnectionSink; only the + // final transport is an in-memory recording peer. + self.work.lock().unwrap().delivered(¬ification); + self.delivered.lock().unwrap().push(notification); + Ok(()) + } + async fn update_acknowledged( + &self, + notification: wire::UpdateSessionNotification, + ) -> Result<(), AcpRuntimeError> { + let gate = self.gate.lock().unwrap().take(); + if let Some((entered, release)) = gate { + let _ = entered.send(()); + let _ = release.await; } - }); - let handle = integration - .bind_session(AcpSessionBinding::new( + self.update(notification) + } + async fn flush(&self) -> Result<(), AcpRuntimeError> { + Ok(()) + } + } + + #[tokio::test] + async fn injection_ownership_tracks_actual_delivery_replacement_and_revocation() { + for tokens in [false, true] { + let root = tempfile::tempdir().unwrap(); + let runtime = Runtime::new(root.path(), "gpt-5.4").unwrap(); + let server = Arc::new(Server::new(runtime, SessionRegistry::new())); + let session_id = wire::SessionId::new(format!("injection-ownership-{tokens}")); + events::activate_diagnostics(&session_id.to_string()); + let work = Arc::new(Mutex::new(InjectionWork::default())); + let delivered = Arc::new(Mutex::new(Vec::new())); + let gate = Arc::new(Mutex::new(None)); + let handle = server + .integration + .bind_session(AcpSessionBinding::new( + session_id.clone(), + SessionId::new(session_id.to_string()), + InjectionDeliverySink { + work: work.clone(), + delivered: delivered.clone(), + gate: gate.clone(), + }, + )) + .unwrap(); + let tasks = AsyncTaskManager::new().handle(); + let (commands, _receiver) = mpsc::channel(1); + server.sessions.lock().unwrap().insert( session_id.clone(), - loop_id.clone(), - sink.clone(), - )) - .unwrap(); - let cancellation_generation = handle.cancellation_handle().generation(); - let task_manager = - AsyncTaskManager::new().routing(|_: &ToolRequest| RoutingDecision::Background); - let tasks = task_manager.handle(); - let release_task = Arc::new(Notify::new()); - if background { - let tools = ToolRegistry::new().with(BlockingTool { - spec: ToolSpec { - name: ToolName::new("race-completion"), - description: "controlled background completion".into(), - input_schema: json!({"type": "object"}), - output_schema: None, - annotations: ToolAnnotations::default(), - metadata: MetadataMap::new(), + SessionHandle { + injections: work.clone(), + token: 0, + commands, + integration: handle.clone(), + busy: Arc::new(AtomicBool::new(true)), + background_jobs: BackgroundJobs::default(), + structured_completion: false, + tasks, }, - entered: Arc::new(AtomicBool::new(false)), - release: release_task.clone(), - }); - let task_id = TaskId::new("race-task"); - task_manager - .start_task( - TaskLaunchRequest::plain( - Some(task_id.clone()), - ToolRequest { - call_id: ToolCallId::new("race-call"), - tool_name: ToolName::new("race-completion"), - input: json!({}), - session_id: loop_id.clone(), - turn_id: TurnId::new("background-turn"), - metadata: MetadataMap::new(), - }, - ), - TaskStartContext { - executor: Arc::new(BasicToolExecutor::new([ - Arc::new(tools) as Arc - ])), - tool_context: OwnedToolContext { - session_id: loop_id.clone(), - turn_id: TurnId::new("background-turn"), - metadata: MetadataMap::new(), - permissions: Arc::new(AllowAllPermissions), - resources: Arc::new(()), - cancellation: None, - execution_scope: None, - approved_request: None, - }, + ); + let model_scopes = Arc::new(Mutex::new(Vec::new())); + let mut driver = Agent::builder() + .model(OwnershipAdapter { + inner: TestAdapter { + outcome: TestOutcome::Content, + turns: Arc::new(AtomicU64::new(0)), + interrupt: None, }, - ) - .await - .unwrap(); - tasks - .set_continue_policy(task_id, ContinuePolicy::RequestContinue) + model_scopes: model_scopes.clone(), + }) + .build() + .unwrap() + .start(SessionConfig::new(SessionId::new(session_id.to_string())).without_cache()) .await .unwrap(); + let source = DiagnosticScope::with_operation( + tokens.then(|| events::DiagnosticOperation::new("SOURCE".into()).unwrap()), + ) + .scope(async { + events::scope_diagnostics(&session_id.to_string(), async { + DiagnosticScope::capture() + }) + .await + }) + .await; + let (client_transport, agent_transport) = agent_client_protocol::Channel::duplex(); + let router = agent_client_protocol::Agent + .protocol_router() + .with_v2(component_from_state(server.clone()).unwrap()); + let serving = tokio::spawn(async move { router.connect_to(agent_transport).await }); + let client = agent_client_protocol::Client.v2().connect_with( + client_transport, + async move |cx| { + let mut initialize = terminal_auth_initialize_request(); + if tokens { + events::offer_diagnostic_ownership(&mut initialize.meta); + } + let initialized = cx.send_request(initialize).block_task().await?; + assert_eq!( + events::ownership_supported(initialized.meta.as_ref()), + tokens + ); + for scenario in [ + "same", + "mixed", + "replace-same", + "replace-mixed", + "reject", + "revoke", + ] { + handle.prepare_injection_turn(); + handle.start_injection_turn(); + let count = model_scopes.lock().unwrap().len(); + let consumption = ConsumptionScope::new(Some(source.clone())); + consumption + .scope(async { + driver + .submit_input(vec![Item::text( + ItemKind::User, + "source invocation", + )]) + .unwrap(); + driver.next().await.unwrap(); + assert_eq!(model_scopes.lock().unwrap()[count], source); + let original = if scenario == "mixed" || scenario == "replace-same" + { + "INJECT" + } else { + "SOURCE" + }; + let mut request = wire::InjectSessionRequest::new( + session_id.clone(), + wire::SessionInjectMode::Steer, + vec![wire::ContentBlock::Text(wire::TextContent::new( + "original injection", + ))], + ); + if tokens { + events::set_diagnostic_operation( + &mut request.meta, + &events::DiagnosticOperation::new(original.into()).unwrap(), + ); + } + let id = cx + .send_request(request) + .block_task() + .await + .unwrap() + .message_id; + assert_eq!( + work.lock().unwrap().origins[&id] + .operation() + .map(|operation| operation.as_str()), + tokens.then_some(original) + ); + if scenario.starts_with("replace") || scenario == "reject" { + let replacement = if scenario == "replace-mixed" { + "REPLACED" + } else { + "SOURCE" + }; + let content = if scenario == "reject" { + "x".repeat( + agentkit_acp::v2::MAX_PENDING_INJECTION_BYTES + 1, + ) + } else { + "accepted replacement".into() + }; + let mut request = wire::ReplaceInjectSessionRequest::new( + session_id.clone(), + id.clone(), + vec![wire::ContentBlock::Text(wire::TextContent::new( + content, + ))], + ); + if tokens { + events::set_diagnostic_operation( + &mut request.meta, + &events::DiagnosticOperation::new(replacement.into()) + .unwrap(), + ); + } + let result = cx.send_request(request).block_task().await; + assert_eq!(result.is_err(), scenario == "reject"); + let expected = if scenario == "reject" { + original + } else { + replacement + }; + assert_eq!( + work.lock().unwrap().origins[&id] + .operation() + .map(|operation| operation.as_str()), + tokens.then_some(expected) + ); + } + if scenario == "revoke" { + cx.send_request(wire::RevokeInjectSessionRequest::new( + session_id.clone(), + id.clone(), + )) + .block_task() + .await + .unwrap(); + assert!(!work.lock().unwrap().pending.contains(&id)); + assert!(!work.lock().unwrap().origins.contains_key(&id)); + assert!( + cx.send_request(wire::ReplaceInjectSessionRequest::new( + session_id.clone(), + id.clone(), + vec![] + )) + .block_task() + .await + .is_err() + ); + } + let boundary = handle + .handle_injection_boundary(&mut driver, false) + .await + .unwrap(); + if scenario == "revoke" { + assert!(matches!(boundary, AcpInjectionBoundary::Continue)); + } else { + assert!(matches!(boundary, AcpInjectionBoundary::Delivered)); + driver.next().await.unwrap(); + let mixed = tokens + && (scenario == "mixed" || scenario == "replace-mixed"); + assert_eq!( + model_scopes.lock().unwrap()[count + 1], + if mixed { + DiagnosticScope::default() + } else { + source.clone() + }, + "{scenario}" + ); + let last = delivered.lock().unwrap().last().cloned().unwrap(); + let expected = if scenario.starts_with("replace") { + "accepted replacement" + } else { + "original injection" + }; + assert!( + serde_json::to_string(&last).unwrap().contains(expected) + ); + } + assert!(!work.lock().unwrap().origins.contains_key(&id)); + assert!(!work.lock().unwrap().pending.contains(&id)); + }) + .await; + } + // A rejected admission must not leave either a reservation or an origin. + handle.stop_injection_turn(); + assert!( + cx.send_request(wire::InjectSessionRequest::new( + session_id.clone(), + wire::SessionInjectMode::Steer, + vec![] + )) + .block_task() + .await + .is_err() + ); + assert!(work.lock().unwrap().origins.is_empty()); + assert!(work.lock().unwrap().pending.is_empty()); + assert_eq!(work.lock().unwrap().admitting, 0); + + // Poll a replacement while the SDK is delivering that exact ID. + // Its Pending poll must release our lock and preserve old owner. + handle.prepare_injection_turn(); + handle.start_injection_turn(); + let mut request = wire::InjectSessionRequest::new( + session_id.clone(), + wire::SessionInjectMode::Steer, + vec![wire::ContentBlock::Text(wire::TextContent::new( + "racing original", + ))], + ); + if tokens { + events::set_diagnostic_operation( + &mut request.meta, + source.operation().unwrap(), + ); + } + let id = cx.send_request(request).block_task().await?.message_id; + let (entered, at_delivery) = oneshot::channel(); + let (release, released) = oneshot::channel(); + *gate.lock().unwrap() = Some((entered, released)); + let driving_handle = handle.clone(); + let expected = source.clone(); + let driving = tokio::spawn(ConsumptionScope::new(Some(source.clone())).scope( + async move { + driving_handle + .handle_injection_boundary(&mut driver, false) + .await + .unwrap(); + driver.next().await.unwrap(); + driver + }, + )); + at_delivery.await.unwrap(); + let mut replace = wire::ReplaceInjectSessionRequest::new( + session_id.clone(), + id.clone(), + vec![wire::ContentBlock::Text(wire::TextContent::new( + "must not replace", + ))], + ); + if tokens { + events::set_diagnostic_operation( + &mut replace.meta, + &events::DiagnosticOperation::new("RACING".into()).unwrap(), + ); + } + let scope = server.diagnostic_scope(replace.meta.as_ref()); + let replacement = scope.sync_scope(|| { + events::scope_diagnostics( + &session_id.to_string(), + server.replace_injection(replace), + ) + }); + tokio::pin!(replacement); + assert!(futures_util::poll!(replacement.as_mut()).is_pending()); + assert_eq!(work.try_lock().unwrap().origins[&id], expected); + release.send(()).unwrap(); + let _driver = driving.await.unwrap(); + assert!( + replacement.await.is_err(), + "delivered input cannot be replaced" + ); + assert_eq!(model_scopes.lock().unwrap().last(), Some(&expected)); + assert!(work.lock().unwrap().origins.is_empty()); + Ok(()) + }, + ); + let result = timeout(Duration::from_secs(10), client).await; + serving.abort(); + let _ = serving.await; + result + .expect("injection ownership peer timed out") + .expect("injection peer failed"); + } + } + + struct OwnershipAdapter { + inner: A, + model_scopes: Arc>>, + } + + struct OwnershipSession { + inner: S, + model_scopes: Arc>>, + } + + #[async_trait] + impl ModelAdapter for OwnershipAdapter { + type Session = OwnershipSession; + async fn start_session(&self, config: SessionConfig) -> Result { + Ok(OwnershipSession { + inner: self.inner.start_session(config).await?, + model_scopes: self.model_scopes.clone(), + }) + } + } + + #[async_trait] + impl ModelSession for OwnershipSession { + type Turn = S::Turn; + async fn begin_turn( + &mut self, + request: TurnRequest, + cancellation: Option, + ) -> Result { + self.model_scopes + .lock() + .unwrap() + .push(DiagnosticScope::capture()); + self.inner.begin_turn(request, cancellation).await + } + } + + struct OwnershipTool { + spec: ToolSpec, + release: Option>, + entered: Arc, + scopes: Arc>>, + } + + #[async_trait] + impl agentkit_tools_core::Tool for OwnershipTool { + fn spec(&self) -> &ToolSpec { + &self.spec + } + async fn invoke( + &self, + request: agentkit_tools_core::ToolRequest, + _ctx: &mut agentkit_tools_core::ToolContext<'_>, + ) -> Result { + self.scopes.lock().unwrap().push(DiagnosticScope::capture()); + self.entered.store(true, Ordering::SeqCst); + if let Some(release) = &self.release { + release.notified().await; + } + self.scopes.lock().unwrap().push(DiagnosticScope::capture()); + Ok(agentkit_tools_core::ToolResult::new( + agentkit_core::ToolResultPart::success( + request.call_id, + ToolOutput::Text("owned completion".into()), + ), + )) + } + } + + #[tokio::test] + async fn explicit_prompt_keeps_new_owner_while_old_task_remains_live() { + use agentkit_core::{ToolCallId, TurnId}; + use agentkit_task_manager::{ContinuePolicy, TaskLaunchRequest, TaskStartContext}; + use agentkit_tools_core::{ + AllowAllPermissions, BasicToolExecutor, OwnedToolContext, ToolRequest, ToolSource, + }; + for tokens in [false, true] { + timeout(Duration::from_secs(5), async { + let root = tempfile::tempdir().unwrap(); + let runtime = Runtime::new(root.path(), "gpt-5.4").unwrap(); + let server = Arc::new(Server::new(runtime.clone(), SessionRegistry::new())); + let mut initialize = terminal_auth_initialize_request(); + if tokens { + events::offer_diagnostic_ownership(&mut initialize.meta); + } + assert_eq!( + events::ownership_supported( + server.initialize(initialize).unwrap().meta.as_ref() + ), + tokens + ); + let session_id = wire::SessionId::new("independent-prompt"); + let loop_id = SessionId::new(session_id.to_string()); + let background_jobs = BackgroundJobs::default(); + let manager = crate::runtime::diagnostics::DiagnosticTaskManager::new( + AsyncTaskManager::new().routing(|request: &ToolRequest| { + if request.tool_name.0 == "old-task" { + RoutingDecision::Background + } else { + RoutingDecision::Foreground + } + }), + background_jobs.task_origins.clone(), + ); + let tasks = manager.handle(); + let old_scopes = Arc::new(Mutex::new(Vec::new())); + let new_scopes = Arc::new(Mutex::new(Vec::new())); + let model_scopes = Arc::new(Mutex::new(Vec::new())); + let entered = Arc::new(AtomicBool::new(false)); + let release = Arc::new(Notify::new()); + let tools = ToolRegistry::new() + .with(OwnershipTool { + spec: ToolSpec::new( + "old-task", + "blocked old task", + json!({"type": "object"}), + ), + release: Some(release.clone()), + entered: entered.clone(), + scopes: old_scopes.clone(), + }) + .with(OwnershipTool { + // TestAdapter's first turn invokes this real tool. + spec: ToolSpec::new( + "missing-tool", + "new prompt tool", + json!({"type": "object"}), + ), + release: None, + entered: Arc::new(AtomicBool::new(false)), + scopes: new_scopes.clone(), + }); + events::activate_diagnostics(&session_id.to_string()); + let old = DiagnosticScope::with_operation( + tokens.then(|| events::DiagnosticOperation::new("OLD".into()).unwrap()), + ); + let task_id = agentkit_core::TaskId::new("old-task-id"); + let old_scope = old + .scope(async { + events::scope_diagnostics(&session_id.to_string(), async { + let captured = DiagnosticScope::capture(); + manager + .start_task( + TaskLaunchRequest::plain( + Some(task_id.clone()), + ToolRequest::new( + ToolCallId::new("old-call"), + ToolName::new("old-task"), + json!({}), + loop_id.clone(), + TurnId::new("old-turn"), + ), + ), + TaskStartContext { + executor: Arc::new(BasicToolExecutor::new([Arc::new( + tools.clone(), + ) + as Arc])), + tool_context: OwnedToolContext { + session_id: loop_id.clone(), + turn_id: TurnId::new("old-turn"), + metadata: MetadataMap::new(), + permissions: Arc::new(AllowAllPermissions), + resources: Arc::new(()), + cancellation: None, + execution_scope: None, + approved_request: None, + }, + }, + ) + .await + .unwrap(); + captured + }) + .await + }) + .await; + tasks + .set_continue_policy(task_id, ContinuePolicy::RequestContinue) + .await + .unwrap(); + while !entered.load(Ordering::SeqCst) { + tokio::task::yield_now().await; + } + // Reactivation changes local authority too; old continuation must + // recover its full old scope, not merely its upstream token. + events::activate_diagnostics(&session_id.to_string()); + let integration = server.integration.clone(); + let recording = RecordingSink::default(); + let sink = ResponseReplacementSink::new(recording); + let activity = SessionActivity::new(|_| Ok(())); + let handle = integration + .bind_session(AcpSessionBinding::new( + session_id.clone(), + loop_id.clone(), + sink.clone(), + )) + .unwrap(); + let turns = Arc::new(AtomicU64::new(0)); + let driver = Agent::builder() + .model(OwnershipAdapter { + inner: TestAdapter { + outcome: TestOutcome::ToolThenContent, + turns: turns.clone(), + interrupt: None, + }, + model_scopes: model_scopes.clone(), + }) + .add_tool_source(tools) + .task_manager(manager) + .observer(crate::runtime::diagnostics::ConsumptionObserver::new( + background_jobs.task_origins.clone(), + )) + .observer(ResponseReplacementObserver::new( + integration.as_ref().clone(), + sink.clone(), + session_id.clone(), + activity.clone(), + )) + .build() + .unwrap() + .start(SessionConfig::new(loop_id).without_cache()) + .await + .unwrap(); + let busy = Arc::new(AtomicBool::new(false)); + let injections = Arc::new(Mutex::new(InjectionWork::default())); + let (commands, receiver) = mpsc::channel(8); + server.sessions.lock().unwrap().insert( + session_id.clone(), + SessionHandle { + injections: injections.clone(), + token: 0, + commands: commands.clone(), + integration: handle.clone(), + busy: busy.clone(), + background_jobs: background_jobs.clone(), + structured_completion: false, + tasks: tasks.clone(), + }, + ); + let actor = tokio::spawn(session_actor(SessionActor { + initial_generation: None, + admission_released: injections.lock().unwrap().admission_released.clone(), + session_id: session_id.clone(), + runtime: runtime.clone(), + integration: integration.clone(), + handle: handle.clone(), + busy: busy.clone(), + binding: BindingGuard { + integration, + session_id: session_id.clone(), + }, + sink, + activity, + driver, + tasks: tasks.clone(), + background_jobs, + structured_completion: false, + skill_catalog: skill_catalog::SkillCatalogMonitor::new(&[]).unwrap(), + adapter: SelectableAdapter::new_with_credentials( + ProviderKind::OpenAiSubscription, + "gpt-5.4", + crate::credentials::CredentialStorage::Memory, + ) + .unwrap(), + catalog: vec![], + commands: receiver, + mcp_events: runtime.subscribe_mcp(session_id.to_string()), + })); + let mut request = wire::PromptRequest::new( + session_id.clone(), + vec![wire::ContentBlock::Text(wire::TextContent::new( + "new invocation", + ))], + ); + events::set_diagnostic_operation( + &mut request.meta, + &events::DiagnosticOperation::new("NEW".into()).unwrap(), + ); + let admission = server.diagnostic_scope(request.meta.as_ref()); + let start = admission + .scope(server.prepare_prompt(request)) + .await + .unwrap(); + start.send(()).unwrap(); + while busy.load(Ordering::Acquire) { + tokio::task::yield_now().await; + } + assert_eq!(turns.load(Ordering::Relaxed), 2); + assert_eq!( + tasks.list_running().await.len(), + 1, + "old producer must remain live through the NEW prompt" + ); + let expected_new = admission + .scope(async { + events::scope_diagnostics(&session_id.to_string(), async { + DiagnosticScope::capture() + }) + .await + }) + .await; + assert_eq!( + *model_scopes.lock().unwrap(), + vec![expected_new.clone(), expected_new.clone()] + ); + assert_eq!( + *new_scopes.lock().unwrap(), + vec![expected_new.clone(), expected_new] + ); + assert_eq!(*old_scopes.lock().unwrap(), vec![old_scope.clone()]); + release.notify_one(); + while turns.load(Ordering::Relaxed) < 3 || busy.load(Ordering::Acquire) { + tokio::task::yield_now().await; + } + assert_eq!(model_scopes.lock().unwrap()[2], old_scope); + assert_eq!( + *old_scopes.lock().unwrap(), + vec![old_scope.clone(), old_scope] + ); + let (reply, response) = oneshot::channel(); + commands.send(Command::Close { reply }).await.unwrap(); + response.await.unwrap(); + actor.await.unwrap(); + }) + .await + .expect("ownership regression timed out"); + } + } + #[derive(Clone, Copy)] + enum RacingBranch { + Prepare, + Reserve, + Abandon, + } + + async fn selected_completion_survives_branch_admission( + background: bool, + branch: RacingBranch, + mixed: bool, + ) { + use agentkit_core::{TaskId, ToolCallId, TurnId}; + use agentkit_task_manager::{ContinuePolicy, TaskLaunchRequest, TaskStartContext}; + use agentkit_tools_core::{ + AllowAllPermissions, BasicToolExecutor, OwnedToolContext, ToolRequest, ToolSource, + }; + + let root = tempfile::tempdir().unwrap(); + let runtime = Runtime::new(root.path(), "gpt-5.4").unwrap(); + let session_id = wire::SessionId::new("selected-completion"); + let loop_id = SessionId::new("selected-completion"); + let integration = Arc::new(AcpIntegration::default()); + let recording = RecordingSink::default(); + let sink = ResponseReplacementSink::new(recording.clone()); + let generations = Arc::new(Mutex::new(Vec::new())); + let origins = Arc::new(Mutex::new(Vec::new())); + let activity = SessionActivity::new({ + let generations = generations.clone(); + move |transition| { + generations + .lock() + .unwrap() + .push((transition.id, transition.active)); + Ok(()) + } + }); + let handle = integration + .bind_session(AcpSessionBinding::new( + session_id.clone(), + loop_id.clone(), + sink.clone(), + )) + .unwrap(); + let cancellation_generation = handle.cancellation_handle().generation(); + let background_jobs = BackgroundJobs::default(); + let task_manager = crate::runtime::diagnostics::DiagnosticTaskManager::new( + AsyncTaskManager::new().routing(|_: &ToolRequest| RoutingDecision::Background), + background_jobs.task_origins.clone(), + ); + let producer = DiagnosticScope::with_operation(Some( + events::DiagnosticOperation::new("background-origin".into()).unwrap(), + )); + let tasks = task_manager.handle(); + let release_task = Arc::new(Notify::new()); + let mut entered_tasks = Vec::new(); + if background { + for index in 0..if mixed { 2 } else { 1 } { + let producer = if index == 0 { + producer.clone() + } else { + DiagnosticScope::with_operation(Some( + events::DiagnosticOperation::new("another-background-origin".into()) + .unwrap(), + )) + }; + let entered = Arc::new(AtomicBool::new(false)); + entered_tasks.push(entered.clone()); + let tools = ToolRegistry::new().with(BlockingTool { + spec: ToolSpec { + name: ToolName::new("race-completion"), + description: "controlled background completion".into(), + input_schema: json!({"type": "object"}), + output_schema: None, + annotations: ToolAnnotations::default(), + metadata: MetadataMap::new(), + }, + entered, + release: release_task.clone(), + }); + let task_id = TaskId::new(format!("race-task-{index}")); + producer + .scope(task_manager.start_task( + TaskLaunchRequest::plain( + Some(task_id.clone()), + ToolRequest { + call_id: ToolCallId::new(format!("race-call-{index}")), + tool_name: ToolName::new("race-completion"), + input: json!({}), + session_id: loop_id.clone(), + turn_id: TurnId::new("background-turn"), + metadata: MetadataMap::new(), + }, + ), + TaskStartContext { + executor: Arc::new(BasicToolExecutor::new([ + Arc::new(tools) as Arc + ])), + tool_context: OwnedToolContext { + session_id: loop_id.clone(), + turn_id: TurnId::new("background-turn"), + metadata: MetadataMap::new(), + permissions: Arc::new(AllowAllPermissions), + resources: Arc::new(()), + cancellation: None, + execution_scope: None, + approved_request: None, + }, + }, + )) + .await + .unwrap(); + tasks + .set_continue_policy(task_id, ContinuePolicy::RequestContinue) + .await + .unwrap(); + } } let turns = Arc::new(AtomicU64::new(0)); let observer = ResponseReplacementObserver::new( @@ -4563,11 +5443,17 @@ mod tests { let input_settlement = crate::runtime::InputSettlement::default(); let driver = Agent::builder() .mutator(input_settlement.clone()) - .model(TestAdapter { - outcome: TestOutcome::Content, - turns: turns.clone(), - interrupt: None, + .model(OwnershipAdapter { + inner: TestAdapter { + outcome: TestOutcome::Content, + turns: turns.clone(), + interrupt: None, + }, + model_scopes: origins.clone(), }) + .observer(crate::runtime::diagnostics::ConsumptionObserver::new( + background_jobs.task_origins.clone(), + )) .observer(observer) .task_manager(task_manager) .build() @@ -4583,37 +5469,53 @@ mod tests { let (commands, receiver) = mpsc::channel(8); let admission = BranchAdmission::claim(busy.clone(), work.clone()).unwrap(); let completion = tasks.clone(); - let mut actor = Box::pin(session_actor(SessionActor { - initial_generation: None, - admission_released: work.lock().unwrap().admission_released.clone(), - session_id: session_id.clone(), - runtime, - integration: integration.clone(), - handle: handle.clone(), - busy: busy.clone(), - binding: BindingGuard { - integration, + let mut actor = Box::pin( + DiagnosticScope::with_operation(Some( + events::DiagnosticOperation::new("unrelated-admission".into()).unwrap(), + )) + .scope(session_actor(SessionActor { + initial_generation: None, + admission_released: work.lock().unwrap().admission_released.clone(), session_id: session_id.clone(), - }, - sink, - activity, - driver, - tasks, - background_jobs: BackgroundJobs::default(), - structured_completion: false, - skill_catalog: skill_catalog::SkillCatalogMonitor::new(&[]).unwrap(), - adapter: SelectableAdapter::new_with_credentials( - ProviderKind::OpenAiSubscription, - "gpt-5.4", - crate::credentials::CredentialStorage::Memory, - ) - .unwrap(), - catalog: vec![], - commands: receiver, - mcp_events, - })); + runtime, + integration: integration.clone(), + handle: handle.clone(), + busy: busy.clone(), + binding: BindingGuard { + integration, + session_id: session_id.clone(), + }, + sink, + activity, + driver, + tasks, + background_jobs, + structured_completion: false, + skill_catalog: skill_catalog::SkillCatalogMonitor::new(&[]).unwrap(), + adapter: SelectableAdapter::new_with_credentials( + ProviderKind::OpenAiSubscription, + "gpt-5.4", + crate::credentials::CredentialStorage::Memory, + ) + .unwrap(), + catalog: vec![], + commands: receiver, + mcp_events, + })), + ); if background { + while entered_tasks + .iter() + .any(|entered| !entered.load(Ordering::SeqCst)) + { + tokio::task::yield_now().await; + } release_task.notify_one(); + if mixed { + release_task.notify_one(); + } + // Complete every real producer before polling the admission-blocked + // actor: task completion alone must not claim a continuation owner. timeout(Duration::from_secs(5), completion.wait_for_idle()) .await .unwrap(); @@ -4710,11 +5612,21 @@ mod tests { "selected MCP completion" }) .count(), - 1 + if mixed { 2 } else { 1 } ); assert_eq!(transcript.matches("autonomous content").count(), 1); assert_eq!(turns.load(Ordering::Relaxed), 1); assert_eq!(*generations.lock().unwrap(), vec![(1, true), (1, false)]); + for origin in origins.lock().unwrap().iter() { + assert_eq!( + origin.operation().map(|op| op.as_str()), + if background && !mixed { + Some("background-origin") + } else { + None + } + ); + } assert_eq!( handle.cancellation_handle().generation(), cancellation_generation @@ -4729,6 +5641,11 @@ mod tests { assert_eq!(turns.load(Ordering::Relaxed), 1); } + #[tokio::test] + async fn branch_actor_mixed_completion_is_unattributed() { + selected_completion_survives_branch_admission(true, RacingBranch::Abandon, true).await; + } + #[tokio::test] async fn branch_actor_selected_mcp_completion_is_not_stranded() { for branch in [ @@ -4736,7 +5653,7 @@ mod tests { RacingBranch::Reserve, RacingBranch::Abandon, ] { - selected_completion_survives_branch_admission(false, branch).await; + selected_completion_survives_branch_admission(false, branch, false).await; } } @@ -4747,7 +5664,7 @@ mod tests { RacingBranch::Reserve, RacingBranch::Abandon, ] { - selected_completion_survives_branch_admission(true, branch).await; + selected_completion_survives_branch_admission(true, branch, false).await; } } @@ -7517,7 +8434,7 @@ mod tests { .as_u64() .unwrap(); assert_eq!( - crate::events::diagnostic_epoch(&source_id), + crate::events::test_support::diagnostic_epoch(&source_id), Some(first_epoch) ); // Rejected requests must not reactivate a loaded actor. Exercise @@ -7541,7 +8458,7 @@ mod tests { .await .expect_err("invalid loaded resume"); assert_eq!( - crate::events::diagnostic_epoch(&source_id), + crate::events::test_support::diagnostic_epoch(&source_id), Some(first_epoch) ); } @@ -7558,7 +8475,7 @@ mod tests { .unwrap(); assert!(same_id_epoch > first_epoch); assert_eq!( - crate::events::diagnostic_epoch(&source_id), + crate::events::test_support::diagnostic_epoch(&source_id), Some(same_id_epoch) ); let listed = cx @@ -7673,7 +8590,7 @@ mod tests { .unwrap(); assert!(returned_epoch > same_id_epoch); assert_eq!( - crate::events::diagnostic_epoch(&source_id), + crate::events::test_support::diagnostic_epoch(&source_id), Some(returned_epoch) ); assert_eq!(serde_json::to_value(&resumed.config_options).unwrap()[1]["currentValue"], "high"); @@ -7702,7 +8619,7 @@ mod tests { ).block_task().await?; let sibling_epoch = sibling_resume.meta.as_ref().unwrap()[crate::events::ACTIVATION_META_KEY].as_u64().unwrap(); assert!(sibling_epoch > returned_epoch); - assert_eq!(crate::events::diagnostic_epoch(&sibling_id), Some(sibling_epoch)); + assert_eq!(crate::events::test_support::diagnostic_epoch(&sibling_id), Some(sibling_epoch)); assert_eq!(serde_json::to_value(&sibling_resume.config_options).unwrap()[1]["currentValue"], "default"); cx.send_request(wire::CloseSessionRequest::new(sibling_id.clone())).block_task().await?; let source_again = cx.send_request( @@ -7711,7 +8628,7 @@ mod tests { ).block_task().await?; let source_epoch = source_again.meta.as_ref().unwrap()[crate::events::ACTIVATION_META_KEY].as_u64().unwrap(); assert!(source_epoch > sibling_epoch); - assert_eq!(crate::events::diagnostic_epoch(&source_id), Some(source_epoch)); + assert_eq!(crate::events::test_support::diagnostic_epoch(&source_id), Some(source_epoch)); assert_eq!(serde_json::to_value(&source_again.config_options).unwrap()[1]["currentValue"], "high"); assert!(updates.lock().unwrap().iter().any(|update| { update.session_id.to_string() == sibling_id @@ -7763,28 +8680,38 @@ mod tests { #[tokio::test] async fn v2_router_advertises_and_routes_pending_injection_replacement() { - let root = tempfile::tempdir().unwrap(); - let runtime = Runtime::new_with_provider_and_credentials( - root.path(), - "gpt-5.4", - crate::ProviderKind::OpenAiSubscription, - crate::credentials::CredentialStorage::Memory, - ) - .unwrap(); - let (client_transport, agent_transport) = agent_client_protocol::Channel::duplex(); - let router = v2_router(runtime, SessionRegistry::new()).unwrap(); - let server = tokio::spawn(async move { router.connect_to(agent_transport).await }); - let client = - agent_client_protocol::Client - .v2() - .connect_with(client_transport, async move |cx| { - let initialized = cx - .send_request(wire::InitializeRequest::new( - wire::ProtocolVersion::V2, - wire::Implementation::new("replacement-test", "0"), - )) - .block_task() - .await?; + for offered in [ + None, + Some(json!({"kitDiagnosticOwnership": {"version": 99, "transport": "stderr"}})), + Some(json!({"kitDiagnosticOwnership": {"version": 1, "transport": "stderr"}})), + ] { + let root = tempfile::tempdir().unwrap(); + let runtime = Runtime::new_with_provider_and_credentials( + root.path(), + "gpt-5.4", + crate::ProviderKind::OpenAiSubscription, + crate::credentials::CredentialStorage::Memory, + ) + .unwrap(); + let (client_transport, agent_transport) = agent_client_protocol::Channel::duplex(); + let router = v2_router(runtime, SessionRegistry::new()).unwrap(); + let server = tokio::spawn(async move { router.connect_to(agent_transport).await }); + let client = agent_client_protocol::Client.v2().connect_with( + client_transport, + async move |cx| { + let mut initialize = wire::InitializeRequest::new( + wire::ProtocolVersion::V2, + wire::Implementation::new("replacement-test", "0"), + ); + initialize.meta = offered + .as_ref() + .map(|value| value.as_object().unwrap().clone()); + let negotiated = events::ownership_supported(initialize.meta.as_ref()); + let initialized = cx.send_request(initialize).block_task().await?; + assert_eq!( + events::ownership_supported(initialized.meta.as_ref()), + negotiated + ); let pending = initialized .capabilities .session @@ -7814,13 +8741,15 @@ mod tests { ); assert_eq!(error.data, Some(json!({ "sessionId": "missing-session" }))); Ok(()) - }); - let result = timeout(Duration::from_secs(2), client).await; - server.abort(); - let _ = server.await; - result - .expect("replacement client timed out") - .expect("replacement client failed"); + }, + ); + let result = timeout(Duration::from_secs(2), client).await; + server.abort(); + let _ = server.await; + result + .expect("replacement client timed out") + .expect("replacement client failed"); + } } #[test] @@ -7906,6 +8835,58 @@ mod tests { ); } + #[tokio::test] + async fn diagnostic_ownership_negotiates_and_scopes_admission() { + for offer in [ + None, + Some(json!({"kitDiagnosticOwnership": {"version": 99, "transport": "stderr"}})), + Some(json!({"kitDiagnosticOwnership": {"version": 1, "transport": "stderr"}})), + ] { + let root = tempfile::tempdir().unwrap(); + let server = Server::new( + Runtime::new(root.path(), "gpt-5.4").unwrap(), + SessionRegistry::new(), + ); + let mut request = terminal_auth_initialize_request(); + request.meta = offer + .as_ref() + .map(|value| value.as_object().unwrap().clone()); + let expected = events::ownership_supported(request.meta.as_ref()); + let response = server.initialize(request).unwrap(); + assert_eq!( + events::ownership_supported(response.meta.as_ref()), + expected + ); + let meta = json!({"kitDiagnosticOperation": "admission-token", "unrelated": true}); + let scope = server.diagnostic_scope(meta.as_object()); + assert_eq!( + scope.operation().map(|op| op.as_str()), + expected.then_some("admission-token") + ); + // Actual actor envelope crosses a channel and survives request scope exit. + let (commands, mut incoming) = mpsc::channel(1); + let (reply, _response) = oneshot::channel(); + scope + .scope(async { + commands + .send(Command::Close { reply }.scoped()) + .await + .unwrap(); + }) + .await; + let (producer, _) = incoming.recv().await.unwrap().into_scoped(); + assert_eq!(producer.operation(), scope.operation()); + assert!(server.diagnostic_scope(None).operation().is_none()); + let malformed = json!({"kitDiagnosticOperation": {"token": "admission-token"}}); + assert!( + server + .diagnostic_scope(malformed.as_object()) + .operation() + .is_none() + ); + } + } + #[test] fn replay_uses_complete_v2_messages_with_stable_ids() { let session_id = wire::SessionId::new("saved"); diff --git a/src/runtime.rs b/src/runtime.rs index 33b4d1f..ad924a2 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -44,7 +44,7 @@ use crate::{ }, }; -mod diagnostics; +pub(crate) mod diagnostics; mod input_settlement; pub(crate) use input_settlement::{InputSettlement, InputSettlingDriver}; @@ -1540,9 +1540,10 @@ impl Runtime { Some((id, name)) => self.subagents.fresh_with_parent(id, name), None => self.subagents.fresh(), }; - let task_manager = background_task_manager(); - let tasks = task_manager.handle(); let background_jobs = BackgroundJobs::default(); + let mut task_manager = background_task_manager(); + task_manager.origins = background_jobs.task_origins.clone(); + let tasks = task_manager.handle(); let canonical_transcript = opened.transcript.clone(); let mut transcript = opened.transcript; let committed_input = if commit_branch { @@ -1574,6 +1575,9 @@ impl Runtime { // Settlement must fence every mutator, including compaction. .mutator(input_settlement.clone()) .mutator(compactor) + .observer(diagnostics::ConsumptionObserver::new( + background_jobs.task_origins.clone(), + )) .observer(context.integration.as_ref().clone()) .transcript_observer(input_settlement.observer(observer)) .transcript(transcript) @@ -1916,6 +1920,7 @@ pub(crate) struct BackgroundActivity { #[derive(Clone)] pub(crate) struct BackgroundJobs { state: Arc>, + pub(crate) task_origins: diagnostics::TaskOrigins, activity: watch::Sender, } @@ -1924,6 +1929,7 @@ impl Default for BackgroundJobs { let (activity, _) = watch::channel(0); Self { state: Arc::new(Mutex::new(BackgroundJobState::default())), + task_origins: diagnostics::TaskOrigins::default(), activity, } } @@ -2454,7 +2460,10 @@ fn backgroundable_spec(mut spec: ToolSpec) -> ToolSpec { } fn background_task_manager() -> diagnostics::DiagnosticTaskManager { - diagnostics::DiagnosticTaskManager(AsyncTaskManager::new().routing(background_route)) + diagnostics::DiagnosticTaskManager::new( + AsyncTaskManager::new().routing(background_route), + diagnostics::TaskOrigins::default(), + ) } fn background_route(request: &ToolRequest) -> RoutingDecision { diff --git a/src/runtime/diagnostics.rs b/src/runtime/diagnostics.rs index 2921e67..9968637 100644 --- a/src/runtime/diagnostics.rs +++ b/src/runtime/diagnostics.rs @@ -1,8 +1,12 @@ //! Carry producer identity through dependency-owned task and Runlet spawns. -use std::sync::Arc; +use std::{ + collections::HashMap, + sync::{Arc, Mutex}, +}; -use agentkit_core::{TurnCancellation, TurnId}; +use agentkit_core::{ToolCallId, TurnCancellation, TurnId}; +use agentkit_loop::{AgentEvent, LoopObserver, ObservedEvent}; use agentkit_task_manager::{ AsyncTaskManager, PendingLoopUpdates, TaskLaunchRequest, TaskManager, TaskManagerError, TaskManagerHandle, TaskStartContext, TaskStartOutcome, TurnTaskUpdate, @@ -15,7 +19,160 @@ use async_trait::async_trait; use crate::events::DiagnosticScope; -pub(super) struct DiagnosticTaskManager(pub(super) AsyncTaskManager); +/// Ephemeral producer ownership, shared with the actor through BackgroundJobs. +/// Retain entries for the driver lifetime: approvals can restart the same call, +/// and a terminal event can precede consumption by the loop. +#[derive(Clone, Default)] +pub(crate) struct TaskOrigins(Arc>); + +#[derive(Default)] +struct TaskOriginState { + origins: HashMap, +} + +impl TaskOrigins { + pub(crate) fn get(&self, call_id: &ToolCallId) -> Option { + self.0 + .lock() + .expect("task origins poisoned") + .origins + .get(call_id) + .cloned() + } + + fn proposed_origin(&self, call_id: &ToolCallId) -> DiagnosticScope { + self.get(call_id).unwrap_or_else(DiagnosticScope::capture) + } + + fn accepted(&self, call_id: ToolCallId, origin: DiagnosticScope) { + let mut state = self.0.lock().expect("task origins poisoned"); + state.origins.entry(call_id).or_insert(origin); + } +} + +/// A coalesced continuation has an owner only when every cause agrees. +#[derive(Clone, Default)] +pub(crate) struct ContinuationOrigin(Option); + +impl ContinuationOrigin { + pub(crate) fn include(&mut self, origin: Option) { + let origin = origin.unwrap_or_default(); + self.0 = Some(match self.0.take() { + None => origin, + Some(previous) if previous == origin => previous, + Some(_) => DiagnosticScope::default(), + }); + } + + pub(crate) fn scope(&self) -> DiagnosticScope { + self.0.clone().unwrap_or_default() + } +} + +tokio::task_local! { + static CONSUMPTION: ConsumptionScope; +} + +/// Ownership of inputs actually presented to this drive. Task-manager handoff +/// alone is not consumption: LoopDriver buffers updates behind fresh input. +#[derive(Clone)] +pub(crate) struct ConsumptionScope { + causes: Arc>, + initial: DiagnosticScope, + refresh_seed: Arc, +} + +impl ConsumptionScope { + /// None admits background work with no assumed cause; Some seeds a real + /// explicit request, or default() for unsolicited external input. + pub(crate) fn new(seed: Option) -> Self { + let initial = seed.clone().unwrap_or_default(); + Self { + causes: Arc::new(Mutex::new(ContinuationOrigin(seed))), + initial, + refresh_seed: Arc::new(std::sync::atomic::AtomicBool::new(false)), + } + } + + /// v2 establishes local activation inside run_active_turn, after admission. + pub(crate) fn explicit(origin: DiagnosticScope) -> Self { + let scope = Self::new(Some(origin)); + scope + .refresh_seed + .store(true, std::sync::atomic::Ordering::Relaxed); + scope + } + + pub(crate) fn bind_local_seed() { + let _ = CONSUMPTION.try_with(|scope| { + if scope + .refresh_seed + .swap(false, std::sync::atomic::Ordering::Relaxed) + { + *scope.causes.lock().expect("consumption scope poisoned") = + ContinuationOrigin(Some(DiagnosticScope::capture())); + } + scope.current_scope().bind_consumed_origin(); + }); + } + + pub(crate) fn current_scope(&self) -> DiagnosticScope { + self.causes + .lock() + .expect("consumption scope poisoned") + .scope() + } + + pub(crate) fn scope( + &self, + future: F, + ) -> impl std::future::Future + use { + CONSUMPTION.scope(self.clone(), self.initial.scope(future)) + } + + pub(crate) fn consumed(origin: Option) { + let _ = CONSUMPTION.try_with(|scope| { + let mut causes = scope.causes.lock().expect("consumption scope poisoned"); + causes.include(origin); + // LoopDriver can begin the model in the SAME poll after emitting + // ToolResultReceived. Rebind here, not at the next host step. This + // changes no previously captured executor/descendant producer. + causes.scope().bind_consumed_origin(); + }); + } +} + +pub(crate) struct ConsumptionObserver(TaskOrigins); + +impl ConsumptionObserver { + pub(crate) fn new(origins: TaskOrigins) -> Self { + Self(origins) + } +} + +impl LoopObserver for ConsumptionObserver { + fn handle_event(&self, event: ObservedEvent) { + let call_id = match &event.event { + AgentEvent::ToolResultReceived(result) => Some(&result.call_id), + AgentEvent::ApprovalRequired(approval) => approval.call_id.as_ref(), + _ => None, + }; + if let Some(call_id) = call_id { + ConsumptionScope::consumed(self.0.get(call_id)); + } + } +} + +pub(crate) struct DiagnosticTaskManager { + pub(super) inner: AsyncTaskManager, + pub(crate) origins: TaskOrigins, +} + +impl DiagnosticTaskManager { + pub(crate) fn new(inner: AsyncTaskManager, origins: TaskOrigins) -> Self { + Self { inner, origins } + } +} #[async_trait] impl TaskManager for DiagnosticTaskManager { @@ -27,7 +184,8 @@ impl TaskManager for DiagnosticTaskManager { // This runs inside the originating turn, before AsyncTaskManager spawns. // Never recover identity from session IDs: the same ID can be reactivated // while an old producer is still running. - let origin = DiagnosticScope::capture(); + let call_id = request.request.call_id.clone(); + let origin = self.origins.proposed_origin(&call_id); ctx.executor = Arc::new(DiagnosticExecutor { inner: ctx.executor, origin: origin.clone(), @@ -37,10 +195,14 @@ impl TaskManager for DiagnosticTaskManager { // this executor. Scoping only the outer compose future is insufficient. scope.executor = Arc::new(DiagnosticExecutor { inner: scope.executor.clone(), - origin, + origin: origin.clone(), }); } - self.0.start_task(request, ctx).await + let outcome = self.inner.start_task(request, ctx).await?; + // Rejected launches must not register phantom causes. The loop cannot + // consume the resolution until this accepted outcome is returned. + self.origins.accepted(call_id, origin); + Ok(outcome) } async fn wait_for_turn( @@ -48,23 +210,23 @@ impl TaskManager for DiagnosticTaskManager { turn_id: &TurnId, cancellation: Option, ) -> Result, TaskManagerError> { - self.0.wait_for_turn(turn_id, cancellation).await + self.inner.wait_for_turn(turn_id, cancellation).await } async fn take_pending_loop_updates(&self) -> Result { - self.0.take_pending_loop_updates().await + self.inner.take_pending_loop_updates().await } async fn wait_for_loop_update(&self) -> Result<(), TaskManagerError> { - self.0.wait_for_loop_update().await + self.inner.wait_for_loop_update().await } async fn on_turn_interrupted(&self, turn_id: &TurnId) -> Result<(), TaskManagerError> { - self.0.on_turn_interrupted(turn_id).await + self.inner.on_turn_interrupted(turn_id).await } fn handle(&self) -> TaskManagerHandle { - self.0.handle() + self.inner.handle() } } @@ -267,6 +429,78 @@ mod tests { } } + #[tokio::test] + async fn task_origins_survive_delivery_and_approval_readmission() { + let manager = crate::runtime::background_task_manager(); + let old = DiagnosticScope::with_operation(Some( + events::DiagnosticOperation::new("old".into()).unwrap(), + )); + let new = DiagnosticScope::with_operation(Some( + events::DiagnosticOperation::new("new".into()).unwrap(), + )); + let call_id = ToolCallId::new("retained-call"); + let tools = ToolRegistry::new().with(Gate { + spec: ToolSpec::new("gate", "approval gate", json!({"type": "object"})), + entered: Arc::new(Barrier::new(1)), + release: Arc::new(Barrier::new(1)), + }); + let executor: Arc = Arc::new(BasicToolExecutor::new([ + Arc::new(tools) as Arc + ])); + let req = request("retained-call", "gate", json!({"approved": true})); + old.scope(manager.start_task( + TaskLaunchRequest::plain(None, req.clone()), + TaskStartContext { + executor: executor.clone(), + tool_context: context(executor.clone()), + }, + )) + .await + .unwrap(); + let Some(TurnTaskUpdate::Resolution(resolution)) = manager + .wait_for_turn(&TurnId::new("turn"), None) + .await + .unwrap() + else { + panic!("approval resolution"); + }; + let agentkit_task_manager::TaskResolution::Approval(pending) = *resolution else { + panic!("real permission checker must require approval"); + }; + assert_eq!(manager.origins.get(&call_id), Some(old.clone())); + new.scope(manager.start_task( + TaskLaunchRequest { + task_id: Some(pending.task_id), + request: req, + kind: TaskLaunchKind::Approved(pending.approval), + }, + TaskStartContext { + executor: executor.clone(), + tool_context: context(executor), + }, + )) + .await + .unwrap(); + let Some(TurnTaskUpdate::Resolution(resolution)) = manager + .wait_for_turn(&TurnId::new("turn"), None) + .await + .unwrap() + else { + panic!("approved completion"); + }; + assert!(matches!( + *resolution, + agentkit_task_manager::TaskResolution::Item(_) + )); + assert_eq!(manager.origins.get(&call_id), Some(old)); + assert!( + manager + .origins + .get(&ToolCallId::new("never-accepted")) + .is_none() + ); + } + #[test] fn dependency_spawns_keep_origin_on_stderr() { for mode in [ @@ -327,9 +561,21 @@ mod tests { }; if call.starts_with("current") { assert_eq!(diagnostic.activation, latest, "{mode}/{switch}"); + assert_eq!( + diagnostic.operation.as_ref().map(|op| op.as_str()), + Some("current-operation") + ); current += 1; } else { assert!(call.starts_with("old"), "{call}"); + assert_eq!( + diagnostic.operation.as_ref().map(|op| op.as_str()), + if mode == "unscoped" { + None + } else { + Some("old-operation") + } + ); assert_eq!(diagnostic.activation, expected, "{mode}/{switch}: {call}"); if started { old_started += 1; @@ -431,7 +677,10 @@ mod tests { None } }; - let task = if mode == "unscoped" { start.await } else { events::scope_diagnostics("A", start).await }; + let task = if mode == "unscoped" { start.await } else { + DiagnosticScope::with_operation(Some(events::DiagnosticOperation::new("old-operation".into()).unwrap())) + .scope(async { events::scope_diagnostics("A", start).await }).await + }; entered.wait().await; if switch { events::activate_diagnostics("B"); } events::activate_diagnostics("A"); @@ -442,11 +691,11 @@ mod tests { completed(&manager).await; } // A newly-created producer still gets the current activation. - events::scope_diagnostics("A", manager.start_task(TaskLaunchRequest { + DiagnosticScope::with_operation(Some(events::DiagnosticOperation::new("current-operation".into()).unwrap())).scope(async { events::scope_diagnostics("A", manager.start_task(TaskLaunchRequest { task_id: None, request: request("current", "gate", json!({"wait": false})), kind: TaskLaunchKind::Plain, - }, TaskStartContext { executor: executor.clone(), tool_context: context(executor) })).await.unwrap(); + }, TaskStartContext { executor: executor.clone(), tool_context: context(executor) })).await }).await.unwrap(); completed(&manager).await; }).await.expect("barrier-controlled task timed out"); } diff --git a/src/tools/subagent.rs b/src/tools/subagent.rs index 3e44f7b..a80e3e6 100644 --- a/src/tools/subagent.rs +++ b/src/tools/subagent.rs @@ -701,8 +701,11 @@ impl Subagents { let transcript_root = root.clone(); let source_id = source_id.clone(); let branch_id = id.clone(); + let diagnostics = events::DiagnosticScope::capture(); let cloned = tokio::task::spawn_blocking(move || { - session::clone_completed(&transcript_root, &source_id, &branch_id) + diagnostics.sync_scope(|| { + session::clone_completed(&transcript_root, &source_id, &branch_id) + }) }) .await .map_err(|error| ChildError::Failed(format!("transcript clone task failed: {error}"))) @@ -1122,12 +1125,12 @@ impl Subagents { return; }; let mut closed = child.closed_signal(); - tokio::spawn(async move { + tokio::spawn(crate::events::inherit_diagnostics(async move { if !*closed.borrow() { let _ = closed.changed().await; } drop(permit); - }); + })); } fn monitor_child_exit(&self, id: String, state: &Arc>, child: &ChildSession) { diff --git a/src/tools/subagent/tests.rs b/src/tools/subagent/tests.rs index e94cb6f..b6e12b9 100644 --- a/src/tools/subagent/tests.rs +++ b/src/tools/subagent/tests.rs @@ -499,10 +499,16 @@ async fn child_exit_monitor_keeps_spawn_activation_across_reactivation() { manager.monitor_child_exit("source".into(), &state, &child); events::DiagnosticScope::default() } else { - events::scope_diagnostics(&session, async { - let expected = events::DiagnosticScope::capture(); - manager.monitor_child_exit("source".into(), &state, &child); - expected + events::DiagnosticScope::with_operation(events::DiagnosticOperation::new(format!( + "monitor:{route}" + ))) + .scope(async { + events::scope_diagnostics(&session, async { + let expected = events::DiagnosticScope::capture(); + manager.monitor_child_exit("source".into(), &state, &child); + expected + }) + .await }) .await }; diff --git a/src/tui/mod.rs b/src/tui/mod.rs index 0323442..3a99a9a 100644 --- a/src/tui/mod.rs +++ b/src/tui/mod.rs @@ -270,7 +270,9 @@ fn runtime_diagnostic_envelope_update( route: &mut Option, diagnostic: events::DiagnosticEvent, ) -> Option { - let events::DiagnosticEvent { event, activation } = diagnostic; + let events::DiagnosticEvent { + event, activation, .. + } = diagnostic; match event { events::RuntimeEvent::SessionStarted { session_id } => { *route = Some(session_id); @@ -303,6 +305,7 @@ fn runtime_diagnostic_update( runtime_diagnostic_envelope_update( route, events::DiagnosticEvent { + operation: None, event, activation: None, }, @@ -3723,6 +3726,7 @@ mod tests { format!( "{EVENT_MARKER}{}\n", serde_json::to_string(&crate::events::DiagnosticEvent { + operation: None, event, activation: Some(crate::events::DiagnosticActivation { session_id: session.into(), @@ -3827,6 +3831,7 @@ mod tests { ), ] { let diagnostic = DiagnosticEvent { + operation: None, event, activation: Some(DiagnosticActivation { session_id: "source".into(), @@ -3907,6 +3912,7 @@ mod tests { let update = super::runtime_diagnostic_envelope_update( &mut ingress, DiagnosticEvent { + operation: None, event: RuntimeEvent::CompactionStarted { reason: "unverified".into(), at: 3, @@ -4627,6 +4633,7 @@ mod tests { let update = super::runtime_diagnostic_envelope_update( &mut None, DiagnosticEvent { + operation: None, event: RuntimeEvent::ChildStarted { call: format!("compose:{epoch}"), tool: "shell".into(), @@ -4649,6 +4656,7 @@ mod tests { let storage = super::runtime_diagnostic_envelope_update( &mut None, DiagnosticEvent { + operation: None, event: RuntimeEvent::StorageStatus { pending: true, exhausted: false, From ea92454255a2abb19f57b85a5ebce5f4f5340fda Mon Sep 17 00:00:00 2001 From: daniel Date: Sun, 6 Sep 2026 00:51:27 +0100 Subject: [PATCH 3/7] fix(diagnostics): bound child ownership tracking --- docs/user/subagents-and-acp-harnesses.md | 6 + docs/user/tui-and-sessions.md | 2 +- src/acp_child.rs | 839 +++++++++++++++++++++-- 3 files changed, 781 insertions(+), 66 deletions(-) diff --git a/docs/user/subagents-and-acp-harnesses.md b/docs/user/subagents-and-acp-harnesses.md index 8a2bed6..bdbf29a 100644 --- a/docs/user/subagents-and-acp-harnesses.md +++ b/docs/user/subagents-and-acp-harnesses.md @@ -146,6 +146,12 @@ Updated Kit parents and children negotiate `_meta.kitDiagnosticOwnership` with t The parent resolves tokens against immutable caller scopes registered before requests are sent, retained until process exit and stderr drain. Child-local activation IDs remain separate. Each grandchild boundary resolves its own tokens and restores upstream ownership. In negotiated mode, missing, unknown, or malformed tokens cannot acquire launch or current-session attribution. Detached producer ownership survives responses and cancellation. Background continuations recover origins from existing tool-call IDs when their results actually enter the loop. An unrelated live task does not suppress a new prompt. If a continuation actually combines different origins, it has no single attributed owner; previously captured producers retain their individual ownership. +Diagnostic resource limits apply per child process: + +- **Startup stderr: 256 KiB retained while initialization is pending.** Kit continues draining the pipe so a noisy child can still acknowledge initialization. On overflow, Kit warns once, keeps only the complete-line prefix within the budget, and discards further output read before acknowledgment. If discarded data ends mid-line, Kit discards through the next newline even after acknowledgment. It never joins truncated bytes into a runtime event. Later complete diagnostic lines use normal ownership checks. The limit can discard ordinary startup logs as well as runtime diagnostics; it does not reject an otherwise valid handshake or impose a new post-initialization log limit. +- **Ownership: 4,096 issued registrations per negotiated child connection**, shared by its native-fork siblings. New prompts and forks fail before dispatch when this capacity is exhausted, with an instruction to close the subagent and start a new child process. Existing registrations are never evicted, reused, or reassigned, including after completed or cancelled requests. Legacy connections consume no registrations; cancellation or a closed channel detected before dispatch also consumes none. +- **Close and cleanup remain available at exhaustion.** An explicit close can proceed without a new ownership token, with a warning; its child-generated stderr diagnostics remain unattributed rather than borrowing another operation's scope. Parent-owned close/roster/results and cleanup using previously issued provenance continue normally. Only a new child process has a fresh budget: session reactivation or closing one native-fork sibling does not reset the shared connection. Close remaining siblings too when retiring the old process. + The extension is optional for external agents. Missing or unknown negotiation versions keep legacy ACP operations usable but cannot supply absent diagnostic identity. See [live diagnostic compatibility](tui-and-sessions.md) for the older-agent boundary after session reactivation. Configure trusted argv profiles in `~/.kit/config.toml`: diff --git a/docs/user/tui-and-sessions.md b/docs/user/tui-and-sessions.md index 1e9dc02..010a394 100644 --- a/docs/user/tui-and-sessions.md +++ b/docs/user/tui-and-sessions.md @@ -103,7 +103,7 @@ Each durable workspace session appears once, using its existing name, short ID, Enter resumes the selected visible session only while **idle**, outside provisional checkout, submission, pending-message edits, or conflicting queued work. Browsing is still available while working. A different session cannot replace an unsent draft or attachments: send or clear them first. Selecting the current session explicitly closes and reloads committed history while retaining its unsent editor draft and attachments. Reload never reconstructs or submits cancelled input. If loading fails after the current actor closes, Kit makes one restoration attempt. If that also fails, the TUI shows **disconnected**, retains the draft and attachments, and blocks backend-dependent actions. Open F4 and select the current session again to retry; only a successful load reconnects it. -Live diagnostic compatibility: after a session reload or switch back, Kit rejects delayed diagnostics from earlier activations. Updated Kit children explicitly negotiate private diagnostic ownership over ACP metadata. Each invocation receives its own operation token, captured by its producers and retained for the child process lifetime. New nested details from a reused child therefore appear under the new activation, while delayed old details remain tied to the old activation. This also applies to native-fork siblings, detached tools, and nested grandchildren; request completion does not reassign their ownership. +Live diagnostic compatibility: after a session reload or switch back, Kit rejects delayed diagnostics from earlier activations. Updated Kit children explicitly negotiate private diagnostic ownership over ACP metadata. Each invocation receives its own operation token, captured by its producers and retained for the child process lifetime. New nested details from a reused child therefore appear under the new activation, while delayed old details remain tied to the old activation. This also applies to native-fork siblings, detached tools, and nested grandchildren; request completion does not reassign their ownership. Child connections have [bounded startup stderr and ownership registration budgets](subagents-and-acp-harnesses.md#private-nested-diagnostic-ownership). If ownership capacity is exhausted, close the affected subagent and start a new child process; close remains available, and switching or reloading the session does not reset that budget. Older external agents that do not negotiate this extension still support ordinary ACP requests, results, and parent-owned subagent roster events. Their process-wide stderr lacks per-operation identity, so it retains immutable launch provenance: after a parent reactivation, new nested stderr details cannot be recovered from that stream. Kit does not invent missing ownership or retag old output. Ordinary stderr remains log output, and storage status remains process-global. A mixed-origin autonomous continuation or unsolicited external event has no single originating invocation; its diagnostics remain unattributed rather than borrowing the latest prompt. diff --git a/src/acp_child.rs b/src/acp_child.rs index 20d1a0e..2681d64 100644 --- a/src/acp_child.rs +++ b/src/acp_child.rs @@ -37,6 +37,9 @@ use crate::{ tools::mcp::CredentialStorage, }; +/// Bound retained provenance per process, without evicting issued identities. +const MAX_DIAGNOSTIC_OPERATIONS: usize = 4096; + /// Immutable producer ownership, separate from request-lifetime output routes. /// Both the process actor and stderr drain retain this registry. Entries are /// never removed on response, cancellation, or session close. @@ -59,16 +62,28 @@ impl DiagnosticOperationRoutes { .operations .lock() .map_err(|_| ChildError::Failed("diagnostic operation registry unavailable".into()))?; + if routes.1.len() >= MAX_DIAGNOSTIC_OPERATIONS { + return Err(ChildError::Failed(format!( + "nested diagnostic operation capacity exhausted ({MAX_DIAGNOSTIC_OPERATIONS}); close this subagent and start a new child process", + ))); + } routes.0 = routes .0 .checked_add(1) - .ok_or_else(|| ChildError::Failed("diagnostic operation sequence exhausted".into()))?; + .ok_or_else(|| ChildError::Failed("diagnostic operation sequence exhausted; close this subagent and start a new child process".into()))?; let operation = DiagnosticOperation::new(format!("operation-{}", routes.0)) .expect("generated diagnostic operation is valid"); routes.1.insert(operation.clone(), scope); Ok(operation) } + fn register_negotiated(&self) -> Result, ChildError> { + if !self.negotiated.load(Ordering::Acquire) { + return Ok(None); + } + self.register().map(Some) + } + fn resolve(&self, operation: &DiagnosticOperation) -> Option { self.operations.lock().ok()?.1.get(operation).cloned() } @@ -76,11 +91,10 @@ impl DiagnosticOperationRoutes { fn owned_close( session_id: SessionId, - operation: &DiagnosticOperation, - negotiated: bool, + operation: Option<&DiagnosticOperation>, ) -> CloseSessionRequest { let mut request = CloseSessionRequest::new(session_id); - if negotiated { + if let Some(operation) = operation { set_diagnostic_operation(&mut request.meta, operation); } request @@ -89,11 +103,10 @@ fn owned_close( fn owned_model( session_id: SessionId, model: &str, - operation: &DiagnosticOperation, - negotiated: bool, + operation: Option<&DiagnosticOperation>, ) -> SetSessionConfigOptionRequest { let mut request = SetSessionConfigOptionRequest::new(session_id, "model", model); - if negotiated { + if let Some(operation) = operation { set_diagnostic_operation(&mut request.meta, operation); } request @@ -481,7 +494,7 @@ impl std::fmt::Display for ChildError { struct Prompt { // The worker, not the waiting caller, owns serialization until settlement. serial: tokio::sync::OwnedMutexGuard<()>, - operation: DiagnosticOperation, + operation: Option, session_id: SessionId, text: String, cancellation: TurnCancellation, @@ -489,7 +502,7 @@ struct Prompt { } struct Fork { serial: tokio::sync::OwnedMutexGuard<()>, - operation: DiagnosticOperation, + operation: Option, session_id: SessionId, model: Option, parent: Option<(String, String)>, @@ -497,7 +510,7 @@ struct Fork { reply: oneshot::Sender>, } struct Close { - operation: DiagnosticOperation, + operation: Option, session_id: SessionId, reply: oneshot::Sender>, } @@ -621,11 +634,15 @@ impl ChildSession { depth: usize, cancellation: TurnCancellation, ) -> Result { + if cancellation.is_cancelled() { + return Err(ChildError::Cancelled); + } let context = config.harnesses.launch_context(&harness); let actor_context = context.clone(); let diagnostics = Arc::new(DiagnosticOperationRoutes::new()); - // Admission ownership must be captured before the actor is spawned. - let admission = diagnostics.register()?; + // Capture before spawn, but retain a token only after explicit agreement + // and before session/new dispatch. Legacy children need no registrations. + let admission = DiagnosticScope::capture(); let actor_diagnostics = Arc::clone(&diagnostics); let (tx, mut rx) = mpsc::channel(1); let (ready_tx, mut ready_rx) = oneshot::channel(); @@ -699,7 +716,6 @@ impl ChildSession { } pub async fn close(&self) -> Result<(), ChildError> { - let operation = self.diagnostics.register()?; if let Some(ancestor_id) = &self.descendant_parent { // Explicit close belongs to its caller, which may be a new turn // using a reused or native-forked child. Deferred cleanup callers @@ -717,17 +733,24 @@ impl ChildSession { )) }; } + let permit = self.tx.reserve().await.map_err(|_| { + ChildError::TerminalFailed("nested agent process is no longer running".into()) + })?; + // Exhaustion must never strand a live process or native-fork sibling. + // Omit ownership rather than borrowing an issued token or a current route. + let operation = self + .diagnostics + .register_negotiated() + .unwrap_or_else(|error| { + eprintln!("ACP harness closing without diagnostic ownership: {error}"); + None + }); let (reply, response) = oneshot::channel(); - self.tx - .send(Request::Close(Close { - operation, - session_id: self.session_id.clone(), - reply, - })) - .await - .map_err(|_| { - ChildError::TerminalFailed("nested agent process is no longer running".into()) - })?; + permit.send(Request::Close(Close { + operation, + session_id: self.session_id.clone(), + reply, + })); response.await.map_err(|_| { ChildError::TerminalFailed( "nested agent process exited without a close response".into(), @@ -741,10 +764,11 @@ impl ChildSession { parent: Option<(String, String)>, cancellation: &TurnCancellation, ) -> Result { - let operation = self.diagnostics.register()?; + let diagnostics = DiagnosticScope::capture(); let serial = tokio::select! { - serial = self.serial.clone().lock_owned() => serial, + biased; () = cancellation.cancelled() => return Err(ChildError::Cancelled), + serial = self.serial.clone().lock_owned() => serial, }; if !self.supports_native_fork() { return Err(ChildError::Failed( @@ -752,19 +776,18 @@ impl ChildSession { )); } let descendant_parent = parent.as_ref().map(|(id, _)| id.clone()); + let permit = self.request_permit(cancellation).await?; + let operation = diagnostics.sync_scope(|| self.diagnostics.register_negotiated())?; let (reply, response) = oneshot::channel(); - tokio::select! { - sent = self.tx.send(Request::Fork(Fork { - serial, - operation, - session_id: self.session_id.clone(), - model: model.map(str::to_owned), - parent, - cancellation: cancellation.clone(), - reply, - })) => sent.map_err(|_| ChildError::TerminalFailed("nested agent process is no longer running".into()))?, - () = cancellation.cancelled() => return Err(ChildError::Cancelled), - } + permit.send(Request::Fork(Fork { + serial, + operation, + session_id: self.session_id.clone(), + model: model.map(str::to_owned), + parent, + cancellation: cancellation.clone(), + reply, + })); let session_id = response.await.map_err(|_| { ChildError::TerminalFailed("nested agent process exited without a fork response".into()) })??; @@ -779,29 +802,55 @@ impl ChildSession { }) } + async fn request_permit( + &self, + cancellation: &TurnCancellation, + ) -> Result, ChildError> { + if self.is_closed() { + return Err(ChildError::TerminalFailed( + "nested agent process is no longer running".into(), + )); + } + let permit = tokio::select! { + biased; + () = cancellation.cancelled() => return Err(ChildError::Cancelled), + permit = self.tx.reserve() => permit.map_err(|_| { + ChildError::TerminalFailed("nested agent process is no longer running".into()) + })?, + }; + if cancellation.is_cancelled() { + return Err(ChildError::Cancelled); + } + if self.is_closed() { + return Err(ChildError::TerminalFailed( + "nested agent process is no longer running".into(), + )); + } + Ok(permit) + } + pub async fn prompt( &self, text: String, cancellation: TurnCancellation, ) -> Result { - let operation = self.diagnostics.register()?; + let diagnostics = DiagnosticScope::capture(); let serial = tokio::select! { - serial = self.serial.clone().lock_owned() => serial, + biased; () = cancellation.cancelled() => return Err(ChildError::Cancelled), + serial = self.serial.clone().lock_owned() => serial, }; + let permit = self.request_permit(&cancellation).await?; + let operation = diagnostics.sync_scope(|| self.diagnostics.register_negotiated())?; let (reply, response) = oneshot::channel(); - let request = Request::Prompt(Prompt { + permit.send(Request::Prompt(Prompt { serial, operation, session_id: self.session_id.clone(), text, cancellation: cancellation.clone(), reply, - }); - tokio::select! { - sent = self.tx.send(request) => sent.map_err(|_| ChildError::TerminalFailed("nested agent process is no longer running".into()))?, - () = cancellation.cancelled() => return Err(ChildError::Cancelled), - } + })); response.await.map_err(|_| { ChildError::TerminalFailed("nested agent process exited without a response".into()) })? @@ -810,7 +859,7 @@ impl ChildSession { struct RunConfig { diagnostics: Arc, - admission: DiagnosticOperation, + admission: DiagnosticScope, config: ChildConfig, harness: String, persisted: Option<(String, bool)>, @@ -930,11 +979,13 @@ async fn run( let negotiated = ownership_supported(initialized.meta.as_ref()); connection_diagnostics.negotiated.store(negotiated, Ordering::Release); let _ = negotiation_tx.send(Some(negotiated)); + let admission = admission.sync_scope(|| connection_diagnostics.register_negotiated()) + .map_err(|error| agent_client_protocol::util::internal_error(error.to_string()))?; let capabilities = initialized.agent_capabilities; let supports_close = capabilities.session_capabilities.close.is_some(); let mut request = agentkit_acp::NewSessionRequest::new(root.clone()); - if negotiated { - set_diagnostic_operation(&mut request.meta, &admission); + if let Some(operation) = &admission { + set_diagnostic_operation(&mut request.meta, operation); } let session = connection.send_request(request).block_task().await?; if let Some(model) = model { @@ -946,7 +997,7 @@ async fn run( let _ = ready.send(Err(error)); return std::future::pending().await; } - if let Err(error) = connection.send_request(owned_model(session.session_id.clone(), model.as_str(), &admission, negotiated)).block_task().await { + if let Err(error) = connection.send_request(owned_model(session.session_id.clone(), model.as_str(), admission.as_ref())).block_task().await { let error = format!("ACP harness {harness:?} rejected model selection {model:?}: {error}"); let _ = ready.send(Err(error)); return std::future::pending().await; @@ -981,8 +1032,8 @@ async fn run( (FORK_PARENT_NAME_META.into(), Value::String(name)), ])); } - if negotiated { - set_diagnostic_operation(&mut request.meta, &fork.operation); + if let Some(operation) = &fork.operation { + set_diagnostic_operation(&mut request.meta, operation); } let mut request = Box::pin(connection.send_request(request).block_task()); let result = tokio::select! { @@ -998,7 +1049,7 @@ async fn run( connection .send_request(owned_model( session_id.clone(), model.as_str(), - &fork.operation, negotiated, + fork.operation.as_ref(), )) .block_task(), ) @@ -1014,7 +1065,7 @@ async fn run( }; if selected.is_err() && supports_close { let close = connection - .send_request(owned_close(session_id.clone(), &fork.operation, negotiated)) + .send_request(owned_close(session_id.clone(), fork.operation.as_ref())) .block_task(); if tokio::time::timeout(CANCEL_SETTLE, close) .await @@ -1046,7 +1097,7 @@ async fn run( && tokio::time::timeout( CANCEL_SETTLE, cleanup_connection - .send_request(owned_close(session_id.clone(), &fork.operation, negotiated)) + .send_request(owned_close(session_id.clone(), fork.operation.as_ref())) .block_task(), ) .await @@ -1074,7 +1125,7 @@ async fn run( && tokio::time::timeout( CANCEL_SETTLE, cleanup_connection - .send_request(owned_close(session_id.clone(), &fork.operation, negotiated)) + .send_request(owned_close(session_id.clone(), fork.operation.as_ref())) .block_task(), ) .await @@ -1099,7 +1150,7 @@ async fn run( let sessions = Arc::clone(&sessions); tasks.spawn(async move { let request = connection - .send_request(owned_close(close.session_id.clone(), &close.operation, negotiated)) + .send_request(owned_close(close.session_id.clone(), close.operation.as_ref())) .block_task(); let result = tokio::time::timeout(CANCEL_SETTLE, request) .await @@ -1131,8 +1182,8 @@ async fn run( let mut request = agentkit_acp::PromptRequest::new( session_id.clone(), vec![ContentBlock::Text(agentkit_acp::TextContent::new(prompt.text))], ); - if negotiated { - set_diagnostic_operation(&mut request.meta, &prompt.operation); + if let Some(operation) = &prompt.operation { + set_diagnostic_operation(&mut request.meta, operation); } let request = connection.send_request(request).block_task(); tokio::pin!(request); @@ -1141,8 +1192,8 @@ async fn run( result = &mut request => (result.map_err(|error| error.to_string()), false), () = prompt.cancellation.cancelled() => { let mut cancel = CancelNotification::new(session_id.clone()); - if negotiated { - set_diagnostic_operation(&mut cancel.meta, &prompt.operation); + if let Some(operation) = &prompt.operation { + set_diagnostic_operation(&mut cancel.meta, operation); } let _ = connection.send_notification(cancel); match tokio::time::timeout(CANCEL_SETTLE, &mut request).await { @@ -1174,7 +1225,7 @@ async fn run( if supports_close { for (session_id, operation) in session_ids { let close = connection - .send_request(owned_close(session_id, &operation, negotiated)) + .send_request(owned_close(session_id, operation.as_ref())) .block_task(); if let Ok(result) = tokio::time::timeout(CANCEL_SETTLE, close).await { result?; @@ -1221,6 +1272,69 @@ async fn run( }) } +// Retain at most this many bytes while the child negotiates ownership. Draining +// must continue after the limit so a noisy child can still finish initialization. +const PRE_NEGOTIATION_STDERR_LIMIT: usize = 256 * 1024; + +#[derive(Default)] +struct PendingStderr { + bytes: Vec, + overflowed: bool, + discard_partial: bool, +} + +impl PendingStderr { + // Returns true exactly once, when the retention budget is first exceeded. + fn push(&mut self, bytes: &[u8]) -> bool { + if bytes.is_empty() { + return false; + } + if self.overflowed { + self.discard_partial = bytes.last() != Some(&b'\n'); + return false; + } + if self.bytes.capacity() == 0 { + self.bytes.reserve_exact(PRE_NEGOTIATION_STDERR_LIMIT); + } + let remaining = PRE_NEGOTIATION_STDERR_LIMIT - self.bytes.len(); + self.bytes + .extend_from_slice(&bytes[..bytes.len().min(remaining)]); + if bytes.len() <= remaining { + return false; + } + self.overflowed = true; + // Keep only whole lines. Discard everything else until negotiation ends, + // including the rest of this partial line even if it crosses the ack. + let complete = self + .bytes + .iter() + .rposition(|byte| *byte == b'\n') + .map_or(0, |position| position + 1); + self.bytes.truncate(complete); + self.discard_partial = bytes.last() != Some(&b'\n'); + true + } +} + +// Only a newline restores a safe line boundary. EOF or an error must not let +// a later read reinterpret an unconsumed suffix as a fresh runtime event. +async fn discard_stderr_line(stderr: &mut (impl tokio::io::AsyncBufRead + Unpin)) -> bool { + loop { + let Ok(bytes) = stderr.fill_buf().await else { + return false; + }; + if bytes.is_empty() { + return false; + } + let newline = bytes.iter().position(|byte| *byte == b'\n'); + let count = newline.map_or(bytes.len(), |position| position + 1); + stderr.consume(count); + if newline.is_some() { + return true; + } + } +} + /// Translate only explicitly negotiated, known operation tokens. Legacy children /// retain immutable launch provenance; neither reader timing nor child epochs /// authorize retagging their diagnostics after parent reactivation. @@ -1236,7 +1350,7 @@ fn spawn_forward_stderr( let launch = DiagnosticScope::capture(); tokio::spawn(async move { let mut stderr = stderr; - let mut pending = Vec::new(); + let mut pending = PendingStderr::default(); let routes = if let Some((routes, mut negotiation)) = ownership { // Drain while initialization is pending: waiting without reading can // fill the stderr pipe and prevent the child from replying at all. @@ -1253,7 +1367,11 @@ fn spawn_forward_stderr( let _ = negotiation.wait_for(Option::is_some).await; break; } - Ok(count) => pending.extend_from_slice(&buffer[..count]), + Ok(count) => { + if pending.push(&buffer[..count]) { + eprintln!("[{label}] pre-negotiation stderr exceeded {PRE_NEGOTIATION_STDERR_LIMIT}-byte limit; discarding remaining pre-negotiation output and any incomplete line"); + } + } } } } @@ -1264,8 +1382,13 @@ fn spawn_forward_stderr( let negotiated = routes .as_ref() .filter(|routes| routes.negotiated.load(Ordering::Acquire)); + let mut stderr = BufReader::new(stderr); + let resume = !pending.discard_partial || discard_stderr_line(&mut stderr).await; forward_stderr( - pending.as_slice().chain(stderr), + pending + .bytes + .as_slice() + .chain(stderr.take(if resume { u64::MAX } else { 0 })), &label, ancestor_id.as_deref(), negotiated.map(Arc::as_ref), @@ -2513,7 +2636,7 @@ for line in sys.stdin: owned_child_invocation(mode, "close-second", second.close()).await.unwrap(); owned_child_invocation(mode, "close-base", base.close()).await.unwrap(); // Every completed request still has its immutable registration. - assert_eq!(base.diagnostics.operations.lock().unwrap().1.len(), 11); + assert_eq!(base.diagnostics.operations.lock().unwrap().1.len(), if mode == "supported" { 11 } else { 0 }); let mut closed = base.closed_signal(); let registry = Arc::downgrade(&base.diagnostics); drop(first); drop(second); drop(base); @@ -2849,6 +2972,357 @@ for line in sys.stdin: } } + #[test] + fn operation_capacity_bounds_dispatch_without_stranding_close() { + use crate::events::{self, RuntimeEvent}; + const PROBE: &str = "KIT_TEST_OPERATION_CAPACITY"; + if std::env::var_os(PROBE).is_some() { + tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .enable_all() + .build() + .unwrap() + .block_on(async { + let root = tempfile::tempdir().unwrap(); + let script = root.path().join("capacity.py"); + std::fs::write( + &script, + r#" +import json, sys, pathlib +mode, marker, directory = sys.argv[1:] +root = pathlib.Path(directory) +entered = root / 'entered' +entered.unlink(missing_ok=True) +owned = mode != 'legacy' +first = None +cancelled_token = None +pending = None +sessions = 0 +prompts = 0 +forks = 0 +closes = [] +def response(request, result): + print(json.dumps({'jsonrpc':'2.0', 'id':request['id'], 'result':result}), flush=True) +def emit(token, label): + event = {'event':'child_started', 'call':label, 'tool':'shell', 'summary':label, 'at':1} + if token is not None: event['operation'] = token + print(marker + json.dumps(event), file=sys.stderr, flush=True) +for line in sys.stdin: + request = json.loads(line) + method = request['method'] + params = request.get('params', {}) + token = params.get('_meta', {}).get('kitDiagnosticOperation') + if method == 'initialize': + caps = {'close': {}} + if mode != 'no-fork': caps['fork'] = {} + result = {'protocolVersion':1, 'agentCapabilities': {'sessionCapabilities':caps}} + if owned: result['_meta'] = {'kitDiagnosticOwnership':{'version':1,'transport':'stderr'}} + response(request, result) + elif method in ['session/new','session/fork']: + assert (token is not None) == owned + sessions += 1 + if method == 'session/fork': forks += 1 + response(request, {'sessionId':str(sessions)}) + elif method == 'session/prompt': + assert (token is not None) == owned + prompts += 1 + text = params['prompt'][0]['text'] + if first is None: first = token + if text == 'cancel': + pending = request + cancelled_token = token + entered.write_text('ready') + else: + response(request, {'stopReason':'end_turn'}) + elif method == 'session/cancel': + assert pending is not None and token == cancelled_token + response(pending, {'stopReason':'cancelled'}) + pending = None + elif method == 'session/close': + closes.append(token) + if len(closes) == 1: + if first is not None: emit(first, 'old-after-capacity') + if cancelled_token is not None: emit(cancelled_token, 'cancelled-after-capacity') + emit(token, 'close-diagnostic') + response(request, {}) + if len(closes) == sessions: + (root / 'stats').write_text(json.dumps({'prompts':prompts, 'forks':forks, + 'closes':closes, 'first':first, 'cancelled':cancelled_token})) + break +"#, + ) + .unwrap(); + for mode in ["supported", "legacy", "no-fork"] { + let harnesses = AcpHarnesses::new(BTreeMap::from([( + "capacity".into(), + AcpHarnessProfile { + command: "python3".into(), + args: vec![ + script.to_string_lossy().into_owned(), + mode.into(), + events::EVENT_MARKER.into(), + root.path().to_string_lossy().into_owned(), + ], + permissions: AcpPermissionPolicy::Deny, + }, + )])) + .unwrap(); + let config = ChildConfig { + root: root.path().to_path_buf(), + model: "unused".into(), + provider: Default::default(), + reasoning_effort: None, + openrouter_api_key: None, + configured_mcp_config: None, + configured_mcp_config_inherited: false, + legacy_mcp_config: false, + mcp_config: None, + credential_storage: Default::default(), + telemetry: Default::default(), + harnesses, + default_harness: "acp.capacity".into(), + parent_id: None, + parent_name: None, + }; + let original = DiagnosticScope::with_operation(DiagnosticOperation::new( + format!("capacity-original-{mode}"), + )); + let base = original + .scope(ChildSession::start( + config.clone(), + "acp.capacity".into(), + None, + None, + 1, + TurnCancellation::default(), + )) + .await + .unwrap(); + original + .scope(base.prompt("first".into(), TurnCancellation::default())) + .await + .unwrap(); + let registry = base.diagnostics.clone(); + if mode == "no-fork" { + assert!(matches!( + base.fork(None, None, &TurnCancellation::default()).await, + Err(ChildError::Failed(_)) + )); + assert_eq!(registry.operations.lock().unwrap().1.len(), 2); + base.close().await.unwrap(); + let mut closed = base.closed_signal(); + drop(base); + tokio::time::timeout( + Duration::from_secs(5), + closed.wait_for(|closed| *closed), + ) + .await + .unwrap() + .unwrap(); + tokio::time::timeout(Duration::from_secs(5), async { + while Arc::strong_count(®istry) != 1 { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + continue; + } + let newer = DiagnosticScope::with_operation(DiagnosticOperation::new( + format!("capacity-new-{mode}"), + )); + let sibling = newer + .scope(base.fork(None, None, &TurnCancellation::default())) + .await + .unwrap(); + let before_rejection = registry.operations.lock().unwrap().1.len(); + // Poll the actual prompt while ordinary serial admission is + // held, then cancel it. No channel or token is consumed. + let serial = base.serial.lock().await; + let controller = agentkit_core::CancellationController::new(); + let mut waiting = + Box::pin(base.prompt( + "never-dispatched".into(), + controller.handle().checkpoint(), + )); + std::future::poll_fn(|cx| { + assert!(std::future::Future::poll(waiting.as_mut(), cx).is_pending()); + std::task::Poll::Ready(()) + }) + .await; + controller.interrupt(); + assert!(matches!(waiting.await, Err(ChildError::Cancelled))); + drop(serial); + assert_eq!( + registry.operations.lock().unwrap().1.len(), + before_rejection + ); + // A dispatched cancellation DOES retain its original token. + let controller = agentkit_core::CancellationController::new(); + let cancellation = controller.handle().checkpoint(); + let child = base.clone(); + let pending = tokio::spawn(newer.scope(async move { + child.prompt("cancel".into(), cancellation).await + })); + tokio::time::timeout(Duration::from_secs(5), async { + while !root.path().join("entered").exists() { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + controller.interrupt(); + assert!(matches!(pending.await.unwrap(), Err(ChildError::Cancelled))); + let rounds = if mode == "supported" { + MAX_DIAGNOSTIC_OPERATIONS - 4 + } else { + MAX_DIAGNOSTIC_OPERATIONS + 1 + }; + for _ in 0..rounds { + newer + .scope(base.prompt("fill".into(), TurnCancellation::default())) + .await + .unwrap(); + } + let expected = if mode == "supported" { + MAX_DIAGNOSTIC_OPERATIONS + } else { + 0 + }; + assert_eq!(registry.operations.lock().unwrap().1.len(), expected); + if mode == "supported" { + for child in [&base, &sibling] { + let error = child + .prompt("over-capacity".into(), TurnCancellation::default()) + .await + .unwrap_err(); + assert!(error.to_string().contains("capacity exhausted")); + assert!(error.to_string().contains("new child process")); + } + let error = + match base.fork(None, None, &TurnCancellation::default()).await { + Err(error) => error, + Ok(_) => panic!("fork dispatched beyond capacity"), + }; + assert!(error.to_string().contains("capacity exhausted")); + } + // Even with shared native sessions and a full registry, + // both closes cross the real transport and are acknowledged. + newer.scope(base.close()).await.unwrap(); + newer.scope(sibling.close()).await.unwrap(); + let mut closed = base.closed_signal(); + tokio::time::timeout( + Duration::from_secs(5), + closed.wait_for(|closed| *closed), + ) + .await + .unwrap() + .unwrap(); + assert!( + base.prompt("already-exited".into(), TurnCancellation::default()) + .await + .is_err() + ); + assert_eq!(registry.operations.lock().unwrap().1.len(), expected); + drop(sibling); + drop(base); + // The actor and stderr reader must both release their + // references: this observes process exit AND stderr drain. + tokio::time::timeout(Duration::from_secs(5), async { + while Arc::strong_count(®istry) != 1 { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + let stats: Value = serde_json::from_slice( + &std::fs::read(root.path().join("stats")).unwrap(), + ) + .unwrap(); + assert_eq!(stats["prompts"].as_u64().unwrap(), (rounds + 2) as u64); + assert_eq!(stats["forks"], 1); + assert_eq!(stats["closes"], serde_json::json!([null, null])); + if mode == "supported" { + let first: DiagnosticOperation = + serde_json::from_value(stats["first"].clone()).unwrap(); + let cancelled: DiagnosticOperation = + serde_json::from_value(stats["cancelled"].clone()).unwrap(); + assert_eq!(registry.resolve(&first), Some(original)); + assert_eq!(registry.resolve(&cancelled), Some(newer.clone())); + // Only a NEW connection resets capacity. Keep the old + // full registry alive while the fresh process works. + let fresh = ChildSession::start( + config, + "acp.capacity".into(), + None, + None, + 1, + TurnCancellation::default(), + ) + .await + .unwrap(); + assert_eq!(fresh.diagnostics.operations.lock().unwrap().1.len(), 1); + fresh + .prompt("restart".into(), TurnCancellation::default()) + .await + .unwrap(); + fresh.close().await.unwrap(); + let mut closed = fresh.closed_signal(); + drop(fresh); + tokio::time::timeout( + Duration::from_secs(5), + closed.wait_for(|closed| *closed), + ) + .await + .unwrap() + .unwrap(); + assert_eq!( + registry.operations.lock().unwrap().1.len(), + MAX_DIAGNOSTIC_OPERATIONS + ); + } else { + assert_eq!(registry.operations.lock().unwrap().0, 0); + } + } + }); + return; + } + let output = std::process::Command::new(std::env::current_exe().unwrap()) + .args([ + "--exact", + "acp_child::tests::operation_capacity_bounds_dispatch_without_stranding_close", + "--nocapture", + ]) + .env(PROBE, "1") + .env(events::EVENTS_ENV, "1") + .output() + .unwrap(); + assert!( + output.status.success(), + "{}", + String::from_utf8_lossy(&output.stderr) + ); + let stderr = String::from_utf8(output.stderr).unwrap(); + let events = stderr + .lines() + .filter_map(events::parse_diagnostic) + .collect::>(); + for (label, operation) in [ + ("old-after-capacity", "capacity-original-supported"), + ("cancelled-after-capacity", "capacity-new-supported"), + ] { + assert!(events.iter().any(|event| { + event.operation.as_ref().is_some_and(|owner| owner.as_str() == operation) + && matches!(&event.event, RuntimeEvent::ChildStarted { call, .. } if call == label) + }), "missing preserved {label}: {stderr}"); + } + assert!(!events.iter().any(|event| { + event.operation.as_ref().is_some_and(|owner| owner.as_str().ends_with("-supported")) + && matches!(&event.event, RuntimeEvent::ChildStarted { call, .. } if call == "close-diagnostic") + }), "unowned exhausted close borrowed an old/current owner"); + assert!(stderr.contains("closing without diagnostic ownership")); + } + #[test] fn diagnostic_operation_sequence_exhaustion_never_reuses_a_token() { let routes = DiagnosticOperationRoutes::new(); @@ -3072,6 +3546,241 @@ for line in sys.stdin: ); } + mod pre_negotiation_stderr { + use super::*; + + #[test] + fn below_budget_retains_partial_lines_and_split_utf8_exactly() { + let mut pending = PendingStderr::default(); + for bytes in [ + &b"complete\npartial "[..], + &b"\xe2"[..], + &b"\x82\xac"[..], + &b""[..], + ] { + assert!(!pending.push(bytes)); + } + assert_eq!(pending.bytes, "complete\npartial €".as_bytes()); + assert!(!pending.overflowed); + assert!(!pending.discard_partial); + } + + #[test] + fn retention_stays_within_production_budget_and_keeps_whole_prefix_lines() { + let mut pending = PendingStderr::default(); + assert!(!pending.push(b"kept\n")); + let fill = vec![b'x'; PRE_NEGOTIATION_STDERR_LIMIT - pending.bytes.len()]; + assert!(!pending.push(&fill)); + assert_eq!(pending.bytes.len(), PRE_NEGOTIATION_STDERR_LIMIT); + assert_eq!(pending.bytes.capacity(), PRE_NEGOTIATION_STDERR_LIMIT); + assert!(pending.push(b"overflow")); + assert_eq!(pending.bytes, b"kept\n"); + assert!(pending.discard_partial); + for _ in 0..64 { + assert!(!pending.push(&fill)); + assert_eq!(pending.bytes, b"kept\n"); + assert_eq!(pending.bytes.capacity(), PRE_NEGOTIATION_STDERR_LIMIT); + } + assert!(!pending.push(b"\ndiscard this complete line too\n")); + assert!(!pending.discard_partial); + assert_eq!(pending.bytes, b"kept\n"); + } + + #[test] + fn oversized_first_line_is_not_retained_as_a_runtime_prefix() { + let mut pending = PendingStderr::default(); + assert!(!pending.push(crate::events::EVENT_MARKER.as_bytes())); + assert!(pending.push(&vec![b' '; PRE_NEGOTIATION_STDERR_LIMIT])); + assert!(pending.bytes.is_empty()); + assert!(pending.discard_partial); + assert!(!pending.push(b"\n")); + assert!(!pending.discard_partial); + assert!(pending.bytes.is_empty()); + } + + #[tokio::test] + async fn discarded_partial_line_preserves_exact_following_bytes() { + let mut stderr = BufReader::with_capacity(3, &b"partial suffix\nnext exact line\n"[..]); + assert!(discard_stderr_line(&mut stderr).await); + let mut remaining = Vec::new(); + stderr.read_to_end(&mut remaining).await.unwrap(); + assert_eq!(remaining, b"next exact line\n"); + let mut stderr = BufReader::with_capacity(3, &b"partial at EOF"[..]); + assert!(!discard_stderr_line(&mut stderr).await); + assert!(stderr.fill_buf().await.unwrap().is_empty()); + } + + #[test] + fn real_peer_exceeds_production_budget_before_ack_and_keeps_later_ownership() { + use crate::events::{self, RuntimeEvent}; + const PROBE: &str = "KIT_TEST_STDERR_BOUND_PEER"; + if std::env::var_os(PROBE).is_some() { + tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .enable_all() + .build() + .unwrap() + .block_on(async { + let root = tempfile::tempdir().unwrap(); + let script = root.path().join("stderr-bound.py"); + std::fs::write(&script, r#" +import json, sys +mode, marker, budget = sys.argv[1], sys.argv[2], int(sys.argv[3]) +def response(request, result): + print(json.dumps({'jsonrpc': '2.0', 'id': request['id'], 'result': result}), flush=True) +def event(operation, call): + return json.dumps({'event': 'child_started', 'call': call, 'tool': call, + 'summary': call, 'at': 1, 'operation': operation}) +def emit(operation, call): + sys.stderr.write(marker + event(operation, call) + '\n') + sys.stderr.flush() +for line in sys.stdin: + request = json.loads(line) + method, params = request['method'], request.get('params', {}) + if method == 'initialize': + assert params['_meta']['kitDiagnosticOwnership'] == {'version': 1, 'transport': 'stderr'} + if mode == 'complete': + line = marker + event(None, 'discard-before-ack') + '\n' + sys.stderr.write(line * (4 * budget // len(line) + 1)) + else: + prefix = marker if mode == 'split-marker' else 'ordinary partial line:' + sys.stderr.write(prefix + ' ' * (4 * budget)) + sys.stderr.flush() # Cannot acknowledge until the pipe has drained > production cap. + response(request, {'protocolVersion': 1, 'agentCapabilities': {'sessionCapabilities': {'close': {}}}, + '_meta': {'kitDiagnosticOwnership': {'version': 1, 'transport': 'stderr'}}}) + elif method == 'session/new': + operation = params['_meta']['kitDiagnosticOperation'] + assert operation + # This request proves the parent accepted the initialize response. The + # suspicious suffix still belongs to the oversized pre-ack partial line. + if mode != 'complete': + prefix = '' if mode == 'split-marker' else marker + sys.stderr.write(prefix + event(operation, 'suspicious-suffix') + '\n') + sys.stderr.flush() + emit(operation, 'later-new') + response(request, {'sessionId': 'bound-peer'}) + elif method == 'session/prompt': + emit(params['_meta']['kitDiagnosticOperation'], 'later-prompt') + response(request, {'stopReason': 'end_turn'}) + elif method == 'session/close': + emit(params['_meta']['kitDiagnosticOperation'], 'later-close') + response(request, {}) +"#).unwrap(); + for mode in ["complete", "split-marker", "runtime-suffix"] { + let harnesses = AcpHarnesses::new(BTreeMap::from([( + "stderr-bound".into(), + AcpHarnessProfile { + command: "python3".into(), + args: vec![ + script.to_string_lossy().into_owned(), + mode.into(), + events::EVENT_MARKER.into(), + PRE_NEGOTIATION_STDERR_LIMIT.to_string(), + ], + permissions: AcpPermissionPolicy::Deny, + }, + )])).unwrap(); + let config = ChildConfig { + root: root.path().to_path_buf(), + model: "unused".into(), + provider: Default::default(), + reasoning_effort: None, + openrouter_api_key: None, + configured_mcp_config: None, + configured_mcp_config_inherited: false, + legacy_mcp_config: false, + mcp_config: None, + credential_storage: Default::default(), + telemetry: Default::default(), + harnesses, + default_harness: "acp.stderr-bound".into(), + parent_id: None, + parent_name: None, + }; + tokio::time::timeout(Duration::from_secs(10), async { + let base = owned_child_invocation(mode, "launch", ChildSession::start( + config, "acp.stderr-bound".into(), None, None, 1, TurnCancellation::default(), + )).await.unwrap(); + assert!(base.diagnostics.negotiated.load(Ordering::Acquire)); + owned_child_invocation(mode, "new", base.prompt("later".into(), TurnCancellation::default())).await.unwrap(); + owned_child_invocation(mode, "close", base.close()).await.unwrap(); + let mut closed = base.closed_signal(); + let registry = Arc::downgrade(&base.diagnostics); + drop(base); + closed.wait_for(|closed| *closed).await.unwrap(); + while registry.upgrade().is_some() { + tokio::task::yield_now().await; + } + }).await.expect("no stderr backpressure deadlock or negotiation failure"); + } + }); + return; + } + let output = std::process::Command::new(std::env::current_exe().unwrap()) + .args([ + "--exact", + "acp_child::tests::pre_negotiation_stderr::real_peer_exceeds_production_budget_before_ack_and_keeps_later_ownership", + "--nocapture", + ]) + .env(PROBE, "1") + .env(events::EVENTS_ENV, "1") + .output() + .unwrap(); + assert!( + output.status.success(), + "{}", + String::from_utf8_lossy(&output.stderr) + ); + let stderr = String::from_utf8(output.stderr).unwrap(); + let warning = format!( + "pre-negotiation stderr exceeded {PRE_NEGOTIATION_STDERR_LIMIT}-byte limit" + ); + assert_eq!(stderr.matches(&warning).count(), 3, "{stderr}"); + let diagnostics = stderr + .lines() + .filter_map(events::parse_diagnostic) + .collect::>(); + for mode in ["complete", "split-marker", "runtime-suffix"] { + for (call, owner) in [ + ("later-new", "launch"), + ("later-prompt", "new"), + ("later-close", "close"), + ] { + let operation = DiagnosticOperation::new(format!("{mode}:{owner}")).unwrap(); + let matches = diagnostics.iter().filter(|envelope| { + envelope.operation.as_ref() == Some(&operation) + && matches!(&envelope.event, RuntimeEvent::ChildStarted { call: actual, .. } if actual == call) + }).collect::>(); + assert_eq!( + matches.len(), + 1, + "missing or duplicate owned {mode}/{call}: {stderr}" + ); + let admission = diagnostics + .iter() + .find(|envelope| { + envelope.operation.as_ref() == Some(&operation) + && matches!(envelope.event, RuntimeEvent::SessionStarted { .. }) + }) + .unwrap(); + assert!(matches[0].activation.is_some()); + assert_eq!(matches[0].activation, admission.activation); + } + } + assert!( + !diagnostics.iter().any(|envelope| matches!( + &envelope.event, RuntimeEvent::ChildStarted { call, .. } + if call == "suspicious-suffix" || call == "discard-before-ack" + )), + "discarded bytes must never become owned runtime events: {stderr}" + ); + assert!( + !stderr.contains("suspicious-suffix"), + "partial suffix must be discarded, not printed" + ); + } + } + mod removes_descendants_on_exit { use super::*; From f7497289203261d25ea00c283109f27cabd89214 Mon Sep 17 00:00:00 2001 From: daniel Date: Sun, 6 Sep 2026 11:53:15 +0100 Subject: [PATCH 4/7] fix(transcript): preserve input and admission boundaries after restack --- src/protocols/acp/v2.rs | 33 ++++++++++++++++++++++++++++++--- src/session/branch.rs | 8 ++++---- src/tui/app.rs | 9 +++++++-- 3 files changed, 41 insertions(+), 9 deletions(-) diff --git a/src/protocols/acp/v2.rs b/src/protocols/acp/v2.rs index c78425d..498a239 100644 --- a/src/protocols/acp/v2.rs +++ b/src/protocols/acp/v2.rs @@ -3915,6 +3915,8 @@ mod tests { .await .unwrap(); let driver = input_settlement.wrap(driver); + // Mirror successful server admission before exercising route restoration. + events::activate_diagnostics(&session_id.to_string()); let mcp = crate::tools::mcp::empty(); let mcp_events = mcp.subscribe(child_id.clone()); let busy = Arc::new(AtomicBool::new(!autonomous)); @@ -4430,7 +4432,17 @@ mod tests { } } } - assert_eq!(routes, ["failed-child", "source", "failed-child", "source"]); + assert_eq!( + routes, + [ + "source", + "failed-child", + "failed-child", + "source", + "failed-child", + "source" + ] + ); assert_eq!(running, 2); return; } @@ -4457,6 +4469,8 @@ mod tests { self.0.flush().await } } + events::activate_diagnostics("source"); + events::activate_diagnostics("failed-child"); for origin in [ExecutionOrigin::Prompt, ExecutionOrigin::Autonomous] { let integration = AcpIntegration::default(); let recording = RecordingSink::default(); @@ -9104,7 +9118,7 @@ mod tests { #[tokio::test] async fn model_switch_cancel_before_first_step_retires_queued_prompt_and_mcp_wake() { use std::io::{Read as _, Write as _}; - use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _}; + use tokio::io::{AsyncBufReadExt as _, AsyncReadExt as _, AsyncWriteExt as _}; const CHILD: &str = "KIT_TEST_MODEL_SWITCH_CANCEL_CHILD"; const ROUTE: &str = "compact-cancel-io-"; @@ -9123,10 +9137,22 @@ mod tests { .kill_on_drop(true) .spawn() .unwrap(); - let mut stderr = child.stderr.take().unwrap(); + let mut stderr = tokio::io::BufReader::new(child.stderr.take().unwrap()); let mut stdout = child.stdout.take().unwrap(); let mut stdin = child.stdin.take().unwrap(); timeout(Duration::from_secs(15), async { + // Drain successful admission's complete activation marker first. + // Only the later restoration is after submit_input and therefore + // an authority for this test's pre-next cancellation barrier. + let mut admission = Vec::new(); + loop { + admission.clear(); + assert_ne!(stderr.read_until(b'\n', &mut admission).await.unwrap(), 0); + let line = std::str::from_utf8(&admission).unwrap().trim_end(); + if matches!(events::parse(line), Some(events::RuntimeEvent::SessionStarted { session_id }) if session_id.starts_with(ROUTE)) { + break; + } + } let needle = format!("\"session_id\":\"{ROUTE}"); let mut prefix = Vec::new(); let mut chunk = [0; 128]; @@ -9197,6 +9223,7 @@ mod tests { // The transport ID need not be the disk ID. Exceed the diagnostic pipe // capacity so its first bytes are observable while emit still blocks. let session_id = wire::SessionId::new(format!("{ROUTE}{}", "x".repeat(2 * 1024 * 1024))); + events::activate_diagnostics(&session_id.to_string()); let loop_id = SessionId::new(durable_id.clone()); let integration = Arc::new(AcpIntegration::default()); let recording = RecordingSink::default(); diff --git a/src/session/branch.rs b/src/session/branch.rs index a93f5d3..243dbde 100644 --- a/src/session/branch.rs +++ b/src/session/branch.rs @@ -879,7 +879,7 @@ mod tests { let mut transcript = prepared(); transcript[0].metadata.insert(METADATA_KEY.into(), value); let opened = open_branch(root.path(), transcript); - opened.observer.commit_creation(); + opened.observer.commit_creation().unwrap(); drop(opened); let before = disk_snapshot(root.path()); assert!(matches!( @@ -918,7 +918,7 @@ mod tests { .unwrap() .append(&transcript[1]) .unwrap(); - opened.observer.commit_creation(); + opened.observer.commit_creation().unwrap(); drop(opened); let before = disk_snapshot(root.path()); assert!(matches!( @@ -942,7 +942,7 @@ mod tests { } let opened = open_branch(root.path(), transcript.clone()); opened.observer.replace(&transcript).unwrap(); - opened.observer.commit_creation(); + opened.observer.commit_creation().unwrap(); drop(opened); let before = disk_snapshot(root.path()); assert!(matches!( @@ -957,7 +957,7 @@ mod tests { fn catalog_lineage_preserves_legacy_roots_without_rewriting() { let root = tempfile::tempdir().unwrap(); let opened = open_branch(root.path(), vec![Item::text(ItemKind::System, "legacy")]); - opened.observer.commit_creation(); + opened.observer.commit_creation().unwrap(); drop(opened); let before = disk_snapshot(root.path()); assert_eq!(catalog_row(root.path()).lineage, CatalogLineage::Root); diff --git a/src/tui/app.rs b/src/tui/app.rs index 82d164a..cd8e18f 100644 --- a/src/tui/app.rs +++ b/src/tui/app.rs @@ -2509,8 +2509,9 @@ impl App { if self.model_switch.is_some() { return; } - // An explicit bracketed paste is not part of the unbracketed key-burst heuristic. - self.last_key = None; + if self.session_dialog.is_some() { + self.last_key = None; + } if let Some(rename) = self .session_dialog .as_mut() @@ -2546,6 +2547,10 @@ impl App { self.sync_navigation(); return; } + // Navigation also routes pasted key controls through this method; + // keep their receipt timestamp for the next key in that same burst. + // Explicit composer pastes are outside the key-burst heuristic. + self.last_key = None; self.file_picker = None; self.queue_handoff = false; self.editor.insert_str(text); From 1db645556a03ddd136499486c90339fa2bf9f417 Mon Sep 17 00:00:00 2001 From: daniel Date: Sun, 6 Sep 2026 11:56:48 +0100 Subject: [PATCH 5/7] fix(diagnostics): track consumed external compaction provenance --- src/protocols/acp/v2.rs | 294 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 294 insertions(+) diff --git a/src/protocols/acp/v2.rs b/src/protocols/acp/v2.rs index 498a239..2c791cb 100644 --- a/src/protocols/acp/v2.rs +++ b/src/protocols/acp/v2.rs @@ -2040,12 +2040,22 @@ async fn session_actor { inner: A, model_scopes: Arc>>, @@ -9115,6 +9126,289 @@ mod tests { } } + #[tokio::test] + async fn model_switch_consumes_selected_mcp_provenance_before_queued_fresh_prompt() { + timeout(Duration::from_secs(15), async { + for external in [true, false] { + let root = tempfile::tempdir().unwrap(); + let runtime = Runtime::new_with_provider_and_credentials( + root.path(), + "gpt-5.4", + ProviderKind::OpenAiSubscription, + crate::credentials::CredentialStorage::Memory, + ) + .unwrap(); + let session_id = wire::SessionId::new(format!("switch-provenance-{external}")); + events::activate_diagnostics(&session_id.to_string()); + let config_origin = events::scope_diagnostics(&session_id.to_string(), async { + DiagnosticScope::with_operation(events::DiagnosticOperation::new( + "config-owner".into(), + )) + }) + .await; + let fresh_origin = events::scope_diagnostics(&session_id.to_string(), async { + DiagnosticScope::with_operation(events::DiagnosticOperation::new( + "fresh-owner".into(), + )) + }) + .await; + let loop_id = SessionId::new(session_id.to_string()); + let integration = Arc::new(AcpIntegration::default()); + let sink = ResponseReplacementSink::new(RecordingSink::default()); + let activity = native_activity(session_id.clone(), sink.clone()); + let handle = integration + .bind_session(AcpSessionBinding::new( + session_id.clone(), + loop_id.clone(), + sink.clone(), + )) + .unwrap(); + let selection = SelectableAdapter::new_with_credentials( + ProviderKind::OpenAiSubscription, + "gpt-5.4", + crate::credentials::CredentialStorage::Memory, + ) + .unwrap(); + let summary_scopes = Arc::new(Mutex::new(Vec::new())); + let model_scopes = Arc::new(Mutex::new(Vec::new())); + let compactor = crate::compaction::automatic( + OwnershipAdapter { + inner: SwitchSummaryAdapter { + selection: selection.clone(), + seen: Arc::new(Mutex::new(Vec::new())), + outcome: TestOutcome::Content, + interrupt: None, + }, + model_scopes: summary_scopes.clone(), + }, + Default::default(), + None, + loop_id.clone(), + ) + .unwrap(); + let manager = AsyncTaskManager::new(); + let tasks = manager.handle(); + let driver = Agent::builder() + .model(OwnershipAdapter { + inner: TestAdapter { + outcome: TestOutcome::Content, + turns: Arc::new(AtomicU64::new(0)), + interrupt: None, + }, + model_scopes: model_scopes.clone(), + }) + .task_manager(manager) + .mutator(compactor) + .cancellation(handle.cancellation_handle()) + .transcript(vec![ + Item::text(ItemKind::User, "older content ".repeat(20_000)), + Item::text(ItemKind::Assistant, "recent response") + .with_usage(Usage::new(agentkit_core::TokenUsage::new(100, 0))), + ]) + .build() + .unwrap() + .start(SessionConfig::new(loop_id).without_cache()) + .await + .unwrap(); + let mcp = crate::tools::mcp::empty(); + let work = Arc::new(Mutex::new(InjectionWork::default())); + let busy = Arc::new(AtomicBool::new(false)); + let admission = BranchAdmission::claim(busy.clone(), work.clone()).unwrap(); + let (commands, receiver) = mpsc::channel(8); + let mut actor = Box::pin(session_actor(SessionActor { + initial_generation: None, + admission_released: work.lock().unwrap().admission_released.clone(), + session_id: session_id.clone(), + runtime, + integration: integration.clone(), + handle: handle.clone(), + busy: busy.clone(), + binding: BindingGuard { + integration, + session_id: session_id.clone(), + }, + sink, + activity, + driver, + tasks, + background_jobs: BackgroundJobs::default(), + structured_completion: false, + skill_catalog: skill_catalog::SkillCatalogMonitor::new(&[]).unwrap(), + adapter: selection.clone(), + catalog: vec![crate::provider::ModelGroup { + provider: ProviderKind::OpenAiSubscription, + models: vec!["gpt-5.4-mini".into()], + context_windows: [("gpt-5.4-mini".into(), 150)].into_iter().collect(), + }], + commands: receiver, + mcp_events: mcp.subscribe(session_id.to_string()), + })); + if external { + mcp.publish( + &session_id.to_string(), + crate::tools::mcp::McpEvent { + message: "selected external input".into(), + }, + ); + } + // Poll the real actor until the event is selected. The genuine + // branch owner prevents autonomous execution without a pause hook. + assert!(futures_util::poll!(&mut actor).is_pending()); + let generation = handle.cancellation_handle().generation(); + // Neither a rejected config nor a successful noncompacting one + // may consume the selected external input or start model work. + for (id, value, accepted) in [ + ("unknown-config", "unused", false), + ("reasoning_effort", "medium", true), + ] { + let (reply, response) = oneshot::channel(); + commands + .try_send(config_origin.sync_scope(|| { + Command::SetConfig { + request: wire::SetSessionConfigOptionRequest::new( + session_id.clone(), + id, + value, + ), + reply, + cancellation_generation: generation, + } + .scoped() + })) + .unwrap(); + let result = tokio::select! { + result = response => result.unwrap(), + () = &mut actor => panic!("config retired the actor"), + }; + assert_eq!(result.is_ok(), accepted); + assert!(summary_scopes.lock().unwrap().is_empty()); + assert!(model_scopes.lock().unwrap().is_empty()); + assert!(busy.load(Ordering::Acquire)); + } + let request = wire::SetSessionConfigOptionRequest::new( + session_id.clone(), + super::super::MODEL_CONFIG_ID, + "openai-subscription:gpt-5.4-mini", + ); + let (reply, response) = oneshot::channel(); + commands + .try_send(config_origin.sync_scope(|| { + Command::SetConfig { + request: request.clone(), + reply, + cancellation_generation: generation, + } + .scoped() + })) + .unwrap(); + let warning = tokio::select! { + result = response => result.unwrap().unwrap_err(), + () = &mut actor => panic!("warning retired the actor"), + }; + let warning: model_switch::Warning = + serde_json::from_value(warning.data.unwrap()[model_switch::META].clone()) + .unwrap(); + let mut request = request; + request.meta = Some(serde_json::Map::from_iter([( + model_switch::META.into(), + serde_json::to_value(model_switch::Confirmation { + token: warning.token, + action: model_switch::Decision::Compact, + }) + .unwrap(), + )])); + let (reply, compacted) = oneshot::channel(); + commands + .try_send(config_origin.sync_scope(|| { + Command::SetConfig { + request, + reply, + cancellation_generation: generation, + } + .scoped() + })) + .unwrap(); + let (reply, prompted) = oneshot::channel(); + commands + .try_send(fresh_origin.sync_scope(|| { + Command::Prompt(PromptCommand { + request: wire::PromptRequest::new( + session_id.clone(), + vec![wire::ContentBlock::Text(wire::TextContent::new( + "fresh after compaction", + ))], + ), + cancellation_generation: generation, + reply, + }) + .scoped() + })) + .unwrap(); + // Queue both actor commands before releasing admission: no + // intervening autonomous no-work drive can clear stale provenance. + drop(admission); + tokio::select! { + result = compacted => { result.unwrap().unwrap(); }, + () = &mut actor => panic!("compaction retired the actor"), + } + let activate = tokio::select! { + result = prompted => result.unwrap().unwrap(), + () = &mut actor => panic!("fresh prompt retired the actor"), + }; + let expected = if external { + DiagnosticScope::default() + } else { + config_origin.clone() + }; + assert!(!summary_scopes.lock().unwrap().is_empty()); + assert!( + summary_scopes + .lock() + .unwrap() + .iter() + .all(|scope| *scope == expected) + ); + assert!( + model_scopes + .lock() + .unwrap() + .iter() + .all(|scope| *scope == expected) + ); + assert_eq!(selection.selection().unwrap().model, "gpt-5.4-mini"); + let before_fresh = model_scopes.lock().unwrap().len(); + activate.send(()).unwrap(); + let (reply, snapshot) = oneshot::channel(); + commands.try_send(Command::Snapshot { reply }).unwrap(); + let snapshot = tokio::select! { + result = snapshot => result.unwrap().unwrap(), + () = &mut actor => panic!("fresh prompt retired the actor"), + }; + let scopes = model_scopes.lock().unwrap().clone(); + assert!(!scopes[before_fresh..].is_empty()); + assert!( + scopes[before_fresh..] + .iter() + .all(|scope| *scope == fresh_origin) + ); + assert!(snapshot.canonical_transcript.iter().any(|item| { + item.kind == ItemKind::User + && serde_json::to_string(item) + .unwrap() + .contains("fresh after compaction") + })); + assert!(!busy.load(Ordering::Acquire)); + assert_eq!(handle.cancellation_handle().generation(), generation); + let (reply, closed) = oneshot::channel(); + commands.try_send(Command::Close { reply }).unwrap(); + actor.await; + closed.await.unwrap(); + } + }) + .await + .expect("compaction provenance regression timed out"); + } + #[tokio::test] async fn model_switch_cancel_before_first_step_retires_queued_prompt_and_mcp_wake() { use std::io::{Read as _, Write as _}; From 08357e579ac5745904deba80072205159cefae96 Mon Sep 17 00:00:00 2001 From: daniel Date: Sun, 6 Sep 2026 13:28:41 +0100 Subject: [PATCH 6/7] fix(tui): close session search with F4 --- src/tui/app.rs | 123 +++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 123 insertions(+) diff --git a/src/tui/app.rs b/src/tui/app.rs index cd8e18f..10d0ac7 100644 --- a/src/tui/app.rs +++ b/src/tui/app.rs @@ -2881,6 +2881,7 @@ impl App { if dialog.searching { match key.code { + KeyCode::F(4) => self.dismiss_sessions(), KeyCode::Esc | KeyCode::Enter => dialog.searching = false, KeyCode::Backspace => { if let Some((index, _)) = dialog.query.grapheme_indices(true).next_back() { @@ -7880,6 +7881,128 @@ mod tests { } } + #[test] + fn explorer_f4_closes_typed_pasted_and_loading_search_preserving_drafts() { + for (pasted, query_before_load) in [(false, false), (true, false), (false, true)] { + for pending_edit in [false, true] { + let mut app = app(); + app.start_session("source".into()); + app.paste("parked draft"); + app.attach( + PathBuf::from("/tmp/image.png"), + "image/png", + AttachmentKind::Image, + 12, + ); + app.editor.move_left(); + let parked_text = app.editor.text().to_owned(); + let parked_cursor = app.editor.cursor(); + let parked_attachments = app.attachments.clone(); + let received_at = Instant::now(); + if pending_edit { + app.can_steer = true; + app.can_replace_steer = true; + app.apply(Update::SteerAccepted { + editable: true, + id: "a".into(), + text: "pending a".into(), + }); + app.handle_key_at(press(KeyCode::F(2)), received_at); + app.handle_key_at(press(KeyCode::Enter), received_at + super::PASTE_GAP * 2); + assert!(app.editing_steer()); + app.paste(" revised"); + app.editor.move_left(); + } + let draft = app.editor.text().to_owned(); + let cursor = app.editor.cursor(); + let attachments = app.attachments.clone(); + let Action::ListSessions { epoch } = + app.handle_key_at(press(KeyCode::F(4)), received_at + super::PASTE_GAP * 4) + else { + panic!("catalog requested") + }; + if !query_before_load { + app.apply(Update::SessionCatalog { + epoch, + result: Ok(explorer_entries()), + }); + } + if pasted { + app.paste("source"); + } else { + for (index, character) in "source".chars().enumerate() { + app.handle_key_at( + press(KeyCode::Char(character)), + received_at + super::PASTE_GAP * (6 + index as u32 * 2), + ); + } + } + if query_before_load { + app.apply(Update::SessionCatalog { + epoch, + result: Ok(explorer_entries()), + }); + } + let dialog = app.session_dialog.as_ref().unwrap(); + assert!(dialog.searching); + assert_eq!(dialog.query, "source"); + assert_eq!(app.session_matches.len(), 1); + assert!(!app.session_catalog_pending); + assert!(matches!( + app.handle_key_at(press(KeyCode::F(4)), received_at + super::PASTE_GAP * 18), + Action::None + )); + assert!(app.session_dialog.is_none()); + app.apply(Update::SessionCatalog { + epoch, + result: Ok(explorer_entries()), + }); + assert!(app.session_dialog.is_none()); + assert!(!app.session_catalog_pending); + assert_eq!(app.session_id.as_deref(), Some("source")); + assert_eq!(app.editor.text(), draft); + assert_eq!(app.editor.cursor(), cursor); + assert_eq!(app.attachments, attachments); + assert_eq!(app.editing_steer(), pending_edit); + if pending_edit { + app.handle_key_at(press(KeyCode::Esc), received_at + super::PASTE_GAP * 20); + assert!(!app.editing_steer()); + } + assert_eq!(app.editor.text(), parked_text); + assert_eq!(app.editor.cursor(), parked_cursor); + assert_eq!(app.attachments, parked_attachments); + } + } + } + + #[test] + fn loading_explorer_f4_dismisses_search_and_rejects_late_catalog() { + let mut app = app(); + let received_at = Instant::now(); + let Action::ListSessions { epoch } = app.handle_key_at(press(KeyCode::F(4)), received_at) + else { + panic!("catalog requested") + }; + app.handle_key_at( + press(KeyCode::Char('s')), + received_at + super::PASTE_GAP * 2, + ); + assert!(app.session_catalog_pending); + assert!(app.session_dialog.as_ref().unwrap().searching); + assert!(matches!( + app.handle_key_at(press(KeyCode::F(4)), received_at + super::PASTE_GAP * 4), + Action::None + )); + assert!(app.session_dialog.is_none()); + assert!(!app.session_catalog_pending); + app.apply(Update::SessionCatalog { + epoch, + result: Ok(explorer_entries()), + }); + assert!(app.session_dialog.is_none()); + assert!(!app.session_catalog_pending); + } + #[test] fn loading_explorer_escape_reopen_retains_new_query_and_ignores_old_result() { let mut app = app(); From 6a8e6e4d5b44120b9d408da6f3431145161b6e9f Mon Sep 17 00:00:00 2001 From: daniel Date: Sun, 6 Sep 2026 14:26:53 +0100 Subject: [PATCH 7/7] fix(acp): preserve delivery ownership across settlement restack --- src/protocols/acp/v2.rs | 44 ++++++++++++++++++++++++++++++----------- 1 file changed, 32 insertions(+), 12 deletions(-) diff --git a/src/protocols/acp/v2.rs b/src/protocols/acp/v2.rs index 2c791cb..f4f3008 100644 --- a/src/protocols/acp/v2.rs +++ b/src/protocols/acp/v2.rs @@ -31,7 +31,10 @@ use tracing::Instrument as _; use crate::{ events::{self, DiagnosticScope}, provider::{ProviderKind, SelectableAdapter, authentication_method_id}, - runtime::{AcpDriverContext, BackgroundJobs, InputSettlingDriver as LoopDriver, Runtime, diagnostics::ConsumptionScope}, + runtime::{ + AcpDriverContext, BackgroundJobs, InputSettlingDriver as LoopDriver, Runtime, + diagnostics::ConsumptionScope, + }, }; use super::activity::{ExecutionOrigin, SessionActivity}; @@ -192,13 +195,11 @@ impl InjectionWork { self.origins.remove(id); } - fn delivered(&mut self, notification: &wire::UpdateSessionNotification) { - if let wire::SessionUpdate::UserMessage(message) = ¬ification.update - && self.pending.remove(&message.message_id) - { + fn delivered(&mut self, id: &wire::MessageId) { + if self.pending.remove(id) { // The SDK has submitted this exact queued input and cannot begin // its next model step until this acknowledgement returns. - ConsumptionScope::consumed(self.origins.remove(&message.message_id)); + ConsumptionScope::consumed(self.origins.remove(id)); } } } @@ -266,13 +267,20 @@ impl Drop for BranchAdmission { #[async_trait] impl AcpSessionUpdateSink for ConnectionSink { fn update(&self, notification: wire::UpdateSessionNotification) -> Result<(), AcpRuntimeError> { + let delivered = match ¬ification.update { + wire::SessionUpdate::UserMessage(message) => Some(message.message_id.clone()), + _ => None, + }; + // Retain both admission and ownership until notification succeeds. self.0 - .send_notification(notification.clone()) + .send_notification(notification) .map_err(|error| AcpRuntimeError::Sdk(error.to_string()))?; - self.1 - .lock() - .expect("injection work poisoned") - .delivered(¬ification); + if let Some(id) = delivered { + self.1 + .lock() + .expect("injection work poisoned") + .delivered(&id); + } Ok(()) } @@ -4667,8 +4675,14 @@ mod tests { ) -> Result<(), AcpRuntimeError> { // Same production consumption boundary as ConnectionSink; only the // final transport is an in-memory recording peer. - self.work.lock().unwrap().delivered(¬ification); + let delivered = match ¬ification.update { + wire::SessionUpdate::UserMessage(message) => Some(message.message_id.clone()), + _ => None, + }; self.delivered.lock().unwrap().push(notification); + if let Some(id) = delivered { + self.work.lock().unwrap().delivered(&id); + } Ok(()) } async fn update_acknowledged( @@ -5210,7 +5224,9 @@ mod tests { )) .unwrap(); let turns = Arc::new(AtomicU64::new(0)); + let input_settlement = crate::runtime::InputSettlement::default(); let driver = Agent::builder() + .mutator(input_settlement.clone()) .model(OwnershipAdapter { inner: TestAdapter { outcome: TestOutcome::ToolThenContent, @@ -5235,6 +5251,7 @@ mod tests { .start(SessionConfig::new(loop_id).without_cache()) .await .unwrap(); + let driver = input_settlement.wrap(driver); let busy = Arc::new(AtomicBool::new(false)); let injections = Arc::new(Mutex::new(InjectionWork::default())); let (commands, receiver) = mpsc::channel(8); @@ -9188,7 +9205,9 @@ mod tests { .unwrap(); let manager = AsyncTaskManager::new(); let tasks = manager.handle(); + let input_settlement = crate::runtime::InputSettlement::default(); let driver = Agent::builder() + .mutator(input_settlement.clone()) .model(OwnershipAdapter { inner: TestAdapter { outcome: TestOutcome::Content, @@ -9210,6 +9229,7 @@ mod tests { .start(SessionConfig::new(loop_id).without_cache()) .await .unwrap(); + let driver = input_settlement.wrap(driver); let mcp = crate::tools::mcp::empty(); let work = Arc::new(Mutex::new(InjectionWork::default())); let busy = Arc::new(AtomicBool::new(false));