From 915fec2e27f1d748936d315543858b14dcdd370d Mon Sep 17 00:00:00 2001 From: Jared Lunde Date: Wed, 7 Oct 2026 00:06:09 -0700 Subject: [PATCH] =?UTF-8?q?agent:=20MCP=20tasks=20follow-ups=20=E2=80=94?= =?UTF-8?q?=20run=20--continue=20resumes,=20unreachable=20stays=20pending,?= =?UTF-8?q?=20per-session=20journal=20auth?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - run journals the MCP tasks it starts and `run --continue` resumes a session's pending tasks before its turn (serve already did). - A server that cannot be reached is not an answer: the resume redials with backoff; still unreachable (or lost for good mid-poll), the call gets a "[MCP task result pending]" placeholder, no result is journaled, and a later prompt or start retries; the real result replaces the placeholder. A configured server that failed to dial at startup is kept dormant for it. - A placeholder is marked structurally (an `mcp_task_placeholder` journal entry, a reserved host kind), never recognised by its text: a real result that begins like one is an answer. - The task journal is authenticated per session, never per machine (`mcp_resume::JournalAuth`). Service mode trusts its tenant-sealed segments and adds no MAC, so another replica on the same store resumes. Locally each entry carries an HMAC under a key in a 0600 sidecar beside the session (`.jsonl.mcp-task-journal.json`, a suffix so `work.1`/`work.2` never share one; the old with_extension name is carried over), made by its first journal write under an exclusive lock on the sidecar, so concurrent first writers (threads, processes) agree on one key. It moves, trashes and restores with the session, so a session copied to another machine or a fresh $HOME still resumes. A prompt-time resumer restart counts the previous resumer's results as answers, so a finished task is never polled again. Replay ignores unsealed entries, so lines planted in the session .jsonl (write/edit) never cause a poll or reach the model as a result. A session without a key (older, or not yet journaled) is read as before; its next journal write makes the key with the entries already there accepted. Records must also match the call they name and a configured server. Residual local risk documented in ARCHITECTURE. - A session file is tightened to 0600 by the next append if it was looser; task journal entries are confirmed sealed in service storage. - The sessionId ownership rule gets an e2e test that fails without it. - Test fixtures write the flag/pid/state files tests read atomically (temp + rename), and the sticky-task abort test waits for complete content: the cancel flag could be read created-but-empty. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01JimHGjsfk2Ktm5GxyZJKKk --- crates/agent/ARCHITECTURE.md | 64 +- .../agent/src/bin/mcp_apps_fixture_server.rs | 12 +- .../src/bin/mcp_fixture_events_server.rs | 17 +- .../agent/src/bin/mcp_fixture_stdio_server.rs | 16 +- .../agent/src/bin/mcp_fixture_tasks_server.rs | 54 +- .../src/bin/mcp_skills_fixture_server.rs | 12 +- crates/agent/src/file_lock.rs | 12 +- crates/agent/src/main.rs | 155 +++++ crates/agent/src/mcp_resume.rs | 460 ++++++++++++-- crates/agent/src/serve.rs | 61 +- crates/agent/src/session_store.rs | 567 ++++++++++++++++- crates/agent/src/tools/mcp.rs | 175 +++++- crates/agent/tests/mcp_tasks.rs | 8 +- crates/agent/tests/mcp_tasks_env/mod.rs | 60 +- crates/agent/tests/mcp_tasks_resume.rs | 579 ++++++++++++++++++ crates/agent/tests/mcp_tasks_service.rs | 148 +++++ .../agent/tests/session_segments_sealing.rs | 61 ++ 17 files changed, 2342 insertions(+), 119 deletions(-) create mode 100644 crates/agent/tests/mcp_tasks_service.rs diff --git a/crates/agent/ARCHITECTURE.md b/crates/agent/ARCHITECTURE.md index ce8fab49..9af25013 100644 --- a/crates/agent/ARCHITECTURE.md +++ b/crates/agent/ARCHITECTURE.md @@ -3559,13 +3559,17 @@ the task, not the connection or the process, as the unit of work. and the owning `sessionId`. Journal writes come from the run's event sink and from the background resumer. They are drained by the idle loop, the busy loop and at run end, and a write whose `sessionId` is not the session being persisted to is dropped. The journal sits beside the - pre-dispatch checkpoint, which already persisted the assistant's `tool_use`. + pre-dispatch checkpoint, which already persisted the assistant's `tool_use`. Headless `run` + journals the same entries with its persisted session (`RunJournal`, from its event sink). - **Pending.** A journaled task is pending when three things hold: its record belongs to this session, its `tool_use` is anywhere on the active path, and nothing answers it, neither a `tool_result` nor a journaled `mcp_task_result`. Requiring this session's record means only the session that created - a task may resume it. A fork copies messages, not custom entries, so it inherits no journal; even a - copied journal under another session id is ignored. The call then gets the generic "interrupted" - repair. + a task may resume it. A fork copies messages, not custom entries, so it inherits no journal; a + journal that does travel under another session id (a session file copied or restored under a new + id) is ignored. The call then gets the generic "interrupted" repair. A record must also describe the + call it names (`mcp____` equal to the call's tool) and name a configured server, and a + client's `append_custom` cannot write the journal's kinds (`mcp_task`, `mcp_task_result`): a forged + record would otherwise have the resumer poll a task id of the client's choosing. - **Resume, in the background.** When a session is loaded, and the moment a command switches to another one (`switch_session`, `new_session`, `fork`, `clone`), `mcp_resume::Resumer` polls each of that session's pending tasks to a terminal status through its `McpCatalog`. The previous session's @@ -3573,6 +3577,17 @@ the task, not the connection or the process, as the unit of work. reaches the clients. - The `ttlMs` backstop counts from `createdAtMs`. A TTL that ran out while the agent was down resolves as expired without contacting the server, even if the server is gone. + - A server that cannot be reached is not an answer. The resume redials a few times (250 ms + doubling, about 5 s); if it still cannot connect, or the connection is lost for good mid-poll, + the call gets a placeholder `tool_result` reading `[MCP task result pending] …` (the model is + told the result is pending), no result is journaled, and an `mcp_task_placeholder` entry marks + that call's `tool_result` as a placeholder, so the task stays pending. A placeholder is known by + that mark, never by its text: a real result that happens to begin the same way is an answer. A + later prompt (a fresh resumer, when the last one finished with a server unreachable) or start + (`serve`, or `run --continue`) tries again, and the real result replaces the placeholder where it + sits. A configured server that could not be dialed at startup is kept dormant in the catalog for + this. + Only a terminal status, a JSON-RPC error such as `-32602`, or an expired TTL resolves a call. - Answered keys are seeded from the record, so the user is never asked twice. - In-task input reaches the session's host. - Its progress streams as `tool_progress`, starting with a "resuming" notice. That notice is not @@ -3594,9 +3609,46 @@ the task, not the connection or the process, as the unit of work. The session store applies every journaled result on the active path when it materializes messages (`session_store::materialize`), so `get_messages`, the HTML export and `run --continue` show exactly what `serve` sent the model. The file stays append-only and ids stay stable. +- **`run --continue`.** `run` resumes too: before its turn, it polls each pending task of the session it + opened (journaled by `run` or `serve`) to its result, in turn, with a line on stderr (it has no + client to stream progress to; Ctrl-C ends the wait as it ends the run), journals each result, and + splices it into the turn it sends. - **Privacy.** Custom entries never reach the model. A task id can be a bearer token for the - server's stored state. The session file keeps it in the clear because resume needs it, and the file - is as private as the conversation it holds. The HTML export withholds `mcp_task*` entries. + server's stored state. Resume needs it, so the session keeps it, exactly as private as the + conversation it holds. A local session file is created `0600`, and every append tightens an + existing file whose mode is looser (an older version's, a restore's, a `chmod`) before writing. + In service mode the journal goes through the tenant-keyed sealed segments with the rest of the + transcript (`tests/session_segments_sealing.rs`). The HTML export withholds `mcp_task*` entries. +- **Journal authentication** (`mcp_resume::JournalAuth`, held by the session's store, which seals + what it journals and replays only what passes). The key goes where the session goes, never with + the machine, so a session resumes on another replica, another machine, a fresh `$HOME` or after + an upgrade. + - Service mode (segments sealed under the tenant key): no MAC at all. The storage already + authenticates every line, which is strictly stronger, and any replica holding the tenant key + reads the journal (`tests/mcp_tasks_service.rs`). + - Local: every entry carries a `mac`, an HMAC-SHA256 over its kind and content, keyed by 32 random + bytes in a `0600` sidecar beside the session (`.jsonl.mcp-task-journal.json`: a suffix + on the whole file name, so `work.1` and `work.2` never share a key; or inside a segmented + session's directory), made by the session's first journal write, so a session that never + journals a task gets no extra file. Every first writer (threads, other processes) agrees on one + key: it is made under an exclusive lock on the sidecar and re-read once the lock is held, and + every journal write adopts the key on disk. A key under the earlier `with_extension` name + (`.mcp-task-journal.json`) is read, and carried over by the next journal write. It is one of the session's sidecars, so it moves, + trashes and restores with it. On replay an entry without a valid `mac` is ignored: a line appended to the + `.jsonl` by a model with `write`/`edit` neither causes a poll nor is delivered as a result. + - A session without a key (from before per-session keys, or that has not journaled yet) is read + as before: its entries are accepted, and opening it writes nothing. Its next journal write makes + the key and records the MACs of the entries already there as accepted, so a task that resumed + before still resumes. Entries written after that must carry a `mac`. + - The residual risk, honestly: locally the key is a file the agent's user can read, and a session + that never journaled has no key yet, so a model that can also read files, or that targets such a + session, can forge an entry. A model with `write` can also delete the key or corrupt it + (overwrite it with anything that does not parse): the session is then keyless (what is there is + accepted), and its next journal write makes a new key that accepts every line already in the + file, planted ones included; and the transcript itself is unauthenticated (a model + that can write the file can plant an ordinary `tool_result`). Locally the seal turns "append a + line" into a deliberate read-then-forge; it does not close forgery. In service mode the store is + out of the tools' reach (they run in the tenant's sandbox), and forgery is closed. **MCP Apps.** A service session's apps pool (`ServiceSession::mcp_apps_pool`) dials the same grant connectors a second time, with `io.modelcontextprotocol/ui` advertised, the first time the session's diff --git a/crates/agent/src/bin/mcp_apps_fixture_server.rs b/crates/agent/src/bin/mcp_apps_fixture_server.rs index 79dd463d..9ca3096a 100644 --- a/crates/agent/src/bin/mcp_apps_fixture_server.rs +++ b/crates/agent/src/bin/mcp_apps_fixture_server.rs @@ -430,7 +430,7 @@ async fn main() { && let Ok(marker) = std::env::var("MCP_APPS_FIXTURE_FAIL_FIRST_UI") && !std::path::Path::new(&marker).exists() { - let _ = std::fs::write(&marker, "failed once"); + write_atomically(&marker, "failed once"); std::process::exit(3); } let changes_view = method == "tools/call" @@ -459,3 +459,13 @@ async fn main() { }); } } + +/// Write a file a test reads, all at once: to a temporary sibling, then `rename` it into place. A +/// reader polling for the file (or its content) can otherwise see it created but still empty, +/// between `write`'s create and its write. +fn write_atomically(path: &str, contents: &str) { + let tmp = format!("{path}.tmp-{}", std::process::id()); + if std::fs::write(&tmp, contents).is_ok() { + let _ = std::fs::rename(&tmp, path); + } +} diff --git a/crates/agent/src/bin/mcp_fixture_events_server.rs b/crates/agent/src/bin/mcp_fixture_events_server.rs index e0849e72..0b095caa 100644 --- a/crates/agent/src/bin/mcp_fixture_events_server.rs +++ b/crates/agent/src/bin/mcp_fixture_events_server.rs @@ -1758,7 +1758,10 @@ async fn main() { if let Ok(pidfile) = std::env::var("MCP_FIXTURE_ORPHAN_PIDFILE") { let _ = tokio::process::Command::new("sh") .arg("-c") - .arg(format!("sleep 600 & echo $! > {pidfile}")) + // Written to a temporary name and renamed, so a test never reads a half-written pid. + .arg(format!( + "sleep 600 & echo $! > {pidfile}.tmp && mv -f {pidfile}.tmp {pidfile}" + )) .status() .await; } @@ -1767,7 +1770,7 @@ async fn main() { // to "clean up" (as a server closing a browser would), then record that we got to. if let Ok(path) = std::env::var("MCP_FIXTURE_EXIT_MARKER") { tokio::time::sleep(Duration::from_millis(500)).await; - let _ = std::fs::write(path, "clean exit"); + write_atomically(&path, "clean exit"); } } else { println!( @@ -1777,3 +1780,13 @@ async fn main() { let _ = accept.await; } } + +/// Write a file a test reads, all at once: to a temporary sibling, then `rename` it into place. A +/// reader polling for the file (or its content) can otherwise see it created but still empty, +/// between `write`'s create and its write. +fn write_atomically(path: &str, contents: &str) { + let tmp = format!("{path}.tmp-{}", std::process::id()); + if std::fs::write(&tmp, contents).is_ok() { + let _ = std::fs::rename(&tmp, path); + } +} diff --git a/crates/agent/src/bin/mcp_fixture_stdio_server.rs b/crates/agent/src/bin/mcp_fixture_stdio_server.rs index 372b7875..5c8a0090 100644 --- a/crates/agent/src/bin/mcp_fixture_stdio_server.rs +++ b/crates/agent/src/bin/mcp_fixture_stdio_server.rs @@ -69,11 +69,21 @@ fn tasks() -> &'static Mutex> { TASKS.get_or_init(|| Mutex::new(HashMap::new())) } +/// Write a file a test reads, all at once: to a temporary sibling, then `rename` it into place. A +/// reader polling for the file (or its content) can otherwise see it created but still empty, +/// between `write`'s create and its write. +fn write_atomically(path: &str, contents: &str) { + let tmp = format!("{path}.tmp-{}", std::process::id()); + if std::fs::write(&tmp, contents).is_ok() { + let _ = std::fs::rename(&tmp, path); + } +} + fn record_cancel_flag(task_id: &str) { let Ok(path) = std::env::var("MCP_FIXTURE_CANCEL_FLAG") else { return; }; - let _ = std::fs::write(path, task_id); + write_atomically(&path, task_id); } fn capabilities() -> Value { @@ -97,7 +107,9 @@ async fn main() { tokio::time::sleep(std::time::Duration::from_millis(delay_ms)).await; } if let Ok(pidfile) = std::env::var("MCP_FIXTURE_ORPHAN_PIDFILE") { - let script = format!("sleep 600 & echo $! > {pidfile}"); + // Written to a temporary name and renamed, so a test never reads a half-written pid. + let script = + format!("sleep 600 & echo $! > {pidfile}.tmp && mv -f {pidfile}.tmp {pidfile}"); let _ = tokio::process::Command::new("sh") .arg("-c") .arg(script) diff --git a/crates/agent/src/bin/mcp_fixture_tasks_server.rs b/crates/agent/src/bin/mcp_fixture_tasks_server.rs index f80d58d5..99574259 100644 --- a/crates/agent/src/bin/mcp_fixture_tasks_server.rs +++ b/crates/agent/src/bin/mcp_fixture_tasks_server.rs @@ -29,7 +29,8 @@ //! - `ttl_task`: `ttlMs: 400`, never leaves `working`. //! - `blip_task`: over HTTP, polls 2 and 3 have their connection dropped without a response (a //! network blip); completes on poll 5 with `blip-done`. -//! - `gated_task`: `working` until the file at `MCP_TASKS_FIXTURE_GATE` exists, then `gated-done`. +//! - `gated_task`: `working` until the file at `MCP_TASKS_FIXTURE_GATE` exists, then `gated-done` +//! (or `MCP_TASKS_FIXTURE_GATED_TEXT`). //! - `sample_task`: one in-task `sampling/createMessage` (`draft`); completes with `sampled:`. //! - `ttl_shift_task`: created with `ttlMs: null`; every poll then says `ttlMs: 300`, never finishing. //! Only a client honouring the *latest* TTL stops (after 60 polls it completes `ttl-ignored`). @@ -51,7 +52,7 @@ //! Resource `fixture-tasks://doc`: `resources/read` answers a `CreateTaskResult`, which the client //! MUST treat as an invalid response (tasks are defined for `tools/call` only). //! -//! Env: `MCP_TASKS_FIXTURE_LOG` (path), `MCP_TASKS_FIXTURE_LEGACY=1` (answer `server/discover` with +//! Env: `MCP_TASKS_FIXTURE_KEEP` (path: gated tasks survive a restart), `MCP_TASKS_FIXTURE_LOG` (path), `MCP_TASKS_FIXTURE_LEGACY=1` (answer `server/discover` with //! `-32601`, forcing the client onto legacy `initialize`), `MCP_TASKS_FIXTURE_GATE` (path), //! `MCP_TASKS_FIXTURE_HTTP_PORT_FILE` (serve Streamable HTTP on `127.0.0.1:/mcp` instead //! of stdio, writing the port to this file; each HTTP log line also records the `Mcp-Method` / @@ -264,6 +265,9 @@ impl Server { if kind == Kind::CrashOnce { save_state(&task_id); } + if kind == Kind::Gated { + keep_gated(&self.tasks); + } // A long interval for `ttl_task`, so only a client that caps its wait by the TTL notices // the TTL on time. let interval = if kind == Kind::Ttl { 10_000 } else { 60 }; @@ -324,7 +328,8 @@ impl Server { Kind::Blip => completed(task_id, ttl, "blip-done"), Kind::Gated => { if gate_open() { - completed(task_id, ttl, "gated-done") + let text = std::env::var("MCP_TASKS_FIXTURE_GATED_TEXT"); + completed(task_id, ttl, text.as_deref().unwrap_or("gated-done")) } else { task_json(task_id, "working", ttl, 50) } @@ -477,14 +482,54 @@ fn ask(task_id: &str, ttl: Value, key: &str, message: &str, interval: u64) -> Va v } +/// Write a file a test reads, all at once: to a temporary sibling, then `rename` it into place. A +/// reader polling for the file (or its content) can otherwise see it created but still empty, +/// between `write`'s create and its write. +fn write_atomically(path: &str, contents: &str) { + let tmp = format!("{path}.tmp-{}", std::process::id()); + if std::fs::write(&tmp, contents).is_ok() { + let _ = std::fs::rename(&tmp, path); + } +} + /// `crash_once_task`: remember the task across a process restart. fn save_state(task_id: &str) { if let Ok(path) = std::env::var("MCP_TASKS_FIXTURE_STATE") { - let _ = std::fs::write(path, task_id); + write_atomically(&path, task_id); } } /// The `crash_once_task` a previous process saved, already past its crash. +/// `MCP_TASKS_FIXTURE_KEEP=`: the gated tasks this server holds, one id per line, so a +/// restarted HTTP server still knows them (a server whose tasks outlive its process). +fn keep_gated(tasks: &HashMap) { + if let Ok(path) = std::env::var("MCP_TASKS_FIXTURE_KEEP") { + let ids: Vec<&str> = tasks + .iter() + .filter(|(_, t)| t.kind == Kind::Gated) + .map(|(id, _)| id.as_str()) + .collect(); + write_atomically(&path, &ids.join("\n")); + } +} + +fn load_kept(tasks: &mut HashMap) { + let Ok(path) = std::env::var("MCP_TASKS_FIXTURE_KEEP") else { + return; + }; + for id in std::fs::read_to_string(path).unwrap_or_default().lines() { + tasks.insert( + id.to_owned(), + Task { + kind: Kind::Gated, + polls: 0, + answers: HashMap::new(), + updates: 0, + }, + ); + } +} + fn load_state(tasks: &mut HashMap) -> bool { let Ok(path) = std::env::var("MCP_TASKS_FIXTURE_STATE") else { return false; @@ -630,6 +675,7 @@ fn envelope(server: &mut Server, id: Value, method: &str, params: &Value) -> Opt #[tokio::main(flavor = "current_thread")] async fn main() { let mut tasks = HashMap::new(); + load_kept(&mut tasks); if load_state(&mut tasks) && let Some(ms) = std::env::var("MCP_TASKS_FIXTURE_RESTART_DELAY_MS") .ok() diff --git a/crates/agent/src/bin/mcp_skills_fixture_server.rs b/crates/agent/src/bin/mcp_skills_fixture_server.rs index 03c08296..b30821d4 100644 --- a/crates/agent/src/bin/mcp_skills_fixture_server.rs +++ b/crates/agent/src/bin/mcp_skills_fixture_server.rs @@ -469,7 +469,7 @@ fn handle(method: &str, params: &Value) -> Result { }] })), "tools/call" if params.get("name").and_then(Value::as_str) == Some("publish_late") => { if let Ok(path) = std::env::var("MCP_SKILLS_FIXTURE_LATE_FLAG") { - let _ = std::fs::write(path, "published"); + write_atomically(&path, "published"); } Ok( json!({ "content": [{ "type": "text", "text": "LATE-PUBLISHED" }], "isError": false }), @@ -821,3 +821,13 @@ async fn serve_http() { () = eof => {} } } + +/// Write a file a test reads, all at once: to a temporary sibling, then `rename` it into place. A +/// reader polling for the file (or its content) can otherwise see it created but still empty, +/// between `write`'s create and its write. +fn write_atomically(path: &str, contents: &str) { + let tmp = format!("{path}.tmp-{}", std::process::id()); + if std::fs::write(&tmp, contents).is_ok() { + let _ = std::fs::rename(&tmp, path); + } +} diff --git a/crates/agent/src/file_lock.rs b/crates/agent/src/file_lock.rs index d255d7ca..c3bf9d72 100644 --- a/crates/agent/src/file_lock.rs +++ b/crates/agent/src/file_lock.rs @@ -23,10 +23,18 @@ const RETRIES: usize = 5; /// A held lock. Dropping it closes the descriptor (releasing the kernel lock) and frees the /// registration. pub struct FileLock { - _file: File, + file: File, _registration: Registration, } +impl FileLock { + /// The locked file, for a holder that writes the lock file's own content. Write through this + /// descriptor, never a second open: over NFS, closing *any* descriptor to the file drops the lock. + pub fn file(&self) -> &File { + &self.file + } +} + /// "This process holds the lock at this path." struct Registration(PathBuf); @@ -80,7 +88,7 @@ pub fn try_lock(lock_path: &Path) -> std::io::Result> { } if same_file(&file, lock_path)? { return Ok(Some(FileLock { - _file: file, + file, _registration: registration, })); } diff --git a/crates/agent/src/main.rs b/crates/agent/src/main.rs index f1046932..eb5f11e0 100644 --- a/crates/agent/src/main.rs +++ b/crates/agent/src/main.rs @@ -3666,6 +3666,9 @@ async fn run_turn_once( // on `Compacted`. Both are no-ops without a session mount. let pressure_armed = std::sync::atomic::AtomicBool::new(true); let session_steers = |ev: &agent_core::AgentEvent| { + if let Some(journal) = RUN_JOURNAL.get() { + journal.observe(ev); + } if !session_memory_active { return; } @@ -3792,6 +3795,141 @@ async fn run_turn_once( /// `persist_run_tail`, after `run_turn` returns), so a crash mid-turn — after several tool /// round-trips already ran real commands or edited real files — lost all record of them, with the /// session file (if any) unable to distinguish that from "nothing happened yet". +/// `run`'s MCP task journal: the same `mcp_task` custom entries `serve` writes (see +/// `beyond_ai_agent::mcp_resume`), so a `run` that dies mid-task can be resumed by +/// `run --continue` (or by `serve`). Process-wide because `run` drives exactly one session per +/// process; set once, when that session is persisted. +static RUN_JOURNAL: std::sync::OnceLock = std::sync::OnceLock::new(); + +struct RunJournal { + store: Arc>>, + session_id: String, +} + +impl RunJournal { + /// Journal a task record a tool emitted (`tool_progress` `details.mcpTask`). + fn observe(&self, ev: &agent_core::AgentEvent) { + let agent_core::AgentEvent::ToolProgress { + id, + details: Some(details), + .. + } = ev + else { + return; + }; + let Some(record) = details.get(tools::mcp::MCP_TASK_DETAILS_KEY) else { + return; + }; + self.append( + beyond_ai_agent::mcp_resume::TASK_ENTRY_KIND, + beyond_ai_agent::mcp_resume::task_entry(record, id, &self.session_id), + ); + } + + fn append(&self, kind: &str, data: serde_json::Value) { + let mut guard = self.store.lock().unwrap_or_else(|e| e.into_inner()); + if let Some(store) = guard.as_mut() + && let Err(e) = store.append_journal(kind, data) + { + eprintln!("run: failed to journal an MCP task: {e}"); + } + } +} + +/// Resolve the MCP tasks a previous process left in flight on this session (journaled by `run` or +/// `serve`), before the next turn, so the model sees what actually happened. `run` has no client to +/// keep company while it waits, so each is polled to its result in turn, with a line on stderr; +/// Ctrl-C ends the wait as it ends the run. Every result is journaled as it lands, and returned for +/// the caller to splice into the turn it is about to send, with the journal's placeholder marks +/// (the calls whose `tool_result` on the path is only a "result pending" placeholder). +async fn resume_mcp_tasks_for_run( + journal: Option<&RunJournal>, + session: &Session, + catalog: &tools::mcp::McpCatalog, +) -> ( + Vec, + std::collections::HashSet, +) { + use beyond_ai_agent::mcp_resume::{ + PLACEHOLDER_ENTRY_KIND, RESULT_ENTRY_KIND, Resolved, TASK_ENTRY_KIND, + }; + let Some(journal) = journal else { + return Default::default(); + }; + let (tasks, results, mut placeholders) = { + let guard = journal.store.lock().unwrap_or_else(|e| e.into_inner()); + match guard.as_ref() { + Some(store) => ( + store.active_journal(TASK_ENTRY_KIND), + beyond_ai_agent::mcp_resume::results(&store.active_journal(RESULT_ENTRY_KIND)), + beyond_ai_agent::mcp_resume::placeholder_ids( + &store.active_journal(PLACEHOLDER_ENTRY_KIND), + ), + ), + None => return Default::default(), + } + }; + let mut pending = beyond_ai_agent::mcp_resume::pending( + &session.messages, + &tasks, + &results, + &placeholders, + &journal.session_id, + ); + // Only a configured server is ever polled. + pending.retain(|p| catalog.has_server(&p.record.server)); + let mut resolved = Vec::with_capacity(pending.len()); + for task in pending { + eprintln!( + "[resuming MCP task {} ({}) left in flight by a previous run]", + task.record.task_id, task.name + ); + let result = match catalog + .resume_task(task.record.clone(), &tools::mcp::host(), None) + .await + { + Ok(output) => Resolved { + tool_use_id: task.tool_use_id, + name: task.name, + content: output.text, + is_error: false, + images: output.images, + pending: false, + }, + // The server is unreachable: the turn is told the result is pending, and the task + // stays journaled for the next `run --continue`. + Err(tools::mcp::ResumeError::Unreachable(why)) => { + eprintln!("[MCP task {} is still pending: {why}]", task.record.task_id); + journal.append( + PLACEHOLDER_ENTRY_KIND, + beyond_ai_agent::mcp_resume::placeholder_entry( + &task.tool_use_id, + &journal.session_id, + ), + ); + placeholders.insert(task.tool_use_id.clone()); + resolved.push(beyond_ai_agent::mcp_resume::pending_placeholder( + &task, &why, + )); + continue; + } + Err(tools::mcp::ResumeError::Final(e)) => Resolved { + tool_use_id: task.tool_use_id, + name: task.name, + content: e.to_string(), + is_error: true, + images: Vec::new(), + pending: false, + }, + }; + let mut data = serde_json::to_value(&result).unwrap_or_default(); + data["sessionId"] = serde_json::json!(journal.session_id); + journal.append(RESULT_ENTRY_KIND, data); + resolved.push(result); + } + (resolved, placeholders) +} + struct DirectCheckpoint(Arc>>); #[async_trait::async_trait] @@ -4828,6 +4966,16 @@ async fn run_task( // crash mid-turn — after several tool round-trips already ran real commands/edited real files — // lost all record of them with no session trace at all. let store = Arc::new(std::sync::Mutex::new(store)); + let run_journal = store + .lock() + .unwrap_or_else(|e| e.into_inner()) + .is_some() + .then(|| { + RUN_JOURNAL.get_or_init(|| RunJournal { + store: store.clone(), + session_id: meta.id.clone(), + }) + }); // Matches `serve`'s own `build_agent`: defaults to the model's own capability-table context // window when `--context-window` isn't given, then applies the reserve/keep-recent overrides. let mut compaction = agent_core::CompactionConfig { @@ -4980,6 +5128,8 @@ async fn run_task( .restore_from_transcript(&session.messages); let initial_message = expand_message(&initial_message, &skills, &prompt_templates, &mcp_skills).await; + let (resumed, placeholders) = + resume_mcp_tasks_for_run(run_journal, &session, &mcp_catalog).await; if initial_images.is_empty() { session.user(initial_message); } else { @@ -4988,6 +5138,11 @@ async fn run_task( initial_images, )); } + if !resumed.is_empty() { + // Into the turn just pushed: the call it answers is at the tip. + let messages: &mut Vec = Arc::make_mut(&mut session.messages); + beyond_ai_agent::mcp_resume::splice(messages, &resumed, &placeholders); + } let life = lifecycle.enabled().then(|| { beyond_ai_agent::lifecycle::Run::begin( std::sync::Arc::clone(&lifecycle), diff --git a/crates/agent/src/mcp_resume.rs b/crates/agent/src/mcp_resume.rs index a330ffcb..5a2b31b6 100644 --- a/crates/agent/src/mcp_resume.rs +++ b/crates/agent/src/mcp_resume.rs @@ -40,6 +40,160 @@ pub const TASK_ENTRY_KIND: &str = crate::tools::mcp::MCP_TASK_ENTRY_KIND; /// Custom-entry kind of a resolved task's result. pub const RESULT_ENTRY_KIND: &str = "mcp_task_result"; +/// Custom-entry kind marking a `tool_use` whose `tool_result` on the path is a "result pending" +/// placeholder (see [`pending_placeholder`]), not an answer. A placeholder is known by this entry, +/// never by its content, so a real result that happens to read like one is still an answer. +pub const PLACEHOLDER_ENTRY_KIND: &str = "mcp_task_placeholder"; + +/// The journal kinds [`JournalAuth`] seals and checks. +fn is_journal_kind(kind: &str) -> bool { + kind == TASK_ENTRY_KIND || kind == RESULT_ENTRY_KIND || kind == PLACEHOLDER_ENTRY_KIND +} + +/// How one session's journal entries are authenticated on replay. Held by the session's store, +/// which seals what it journals ([`seal`](Self::seal)) and replays only what passes +/// ([`accepts`](Self::accepts)). +/// +/// - **Storage**: the store already authenticates every line (service mode's segments, sealed under +/// the tenant's key, which is strictly stronger than a MAC). No MAC, so any replica holding the +/// tenant key reads the journal. +/// - **Keyed**: a key that lives *with the session*, in a `0600` sidecar ([`SIDECAR`](Self::SIDECAR)) +/// that moves, trashes and restores with it, so the journal resumes wherever the session goes: +/// another machine, a fresh `$HOME`, an upgrade. It is made by the session's first journal write +/// (a session that never journals gets no file). Each entry carries a `mac` (HMAC-SHA256 over its +/// kind and content). `legacy` holds the MACs of the entries the session already had when its key +/// was made (written before per-session keys): accepted, so a session that resumed before still +/// does. +/// - **Unkeyed**: no key yet (the session has not journaled since per-session keys, or its key +/// could not be written): entries are accepted as they were before journals were sealed, and the +/// next journal write makes the key. +/// +/// What a key does and does not stop, honestly: a session's `.jsonl` is ordinary file content, and +/// a model with the ungated `write`/`edit` tools can append lines to it. A line it plants without +/// the key carries no valid `mac` and is ignored on replay, so it neither causes a poll nor is +/// delivered as a tool result. But the key is a file beside the session that the agent's own user +/// can read, so a model that can also read files and goes looking can forge an entry. A model with +/// `write` can also delete the key or overwrite it with anything that does not parse: the session +/// is then unkeyed (what is there is accepted) until its next journal write makes a new key, which +/// accepts every line already in the file, planted ones included; and the transcript itself +/// is unauthenticated, so a model that can write the file can always plant an ordinary +/// `tool_result`. Locally this raises the bar from "append a line" to a deliberate read-then-forge, +/// not more. In service mode the store is out of the tools' reach (they run in the tenant's +/// sandbox), and forgery is closed. +#[derive(Clone, Debug)] +pub enum JournalAuth { + Storage, + Keyed { + key: [u8; 32], + legacy: Arc>, + }, + Unkeyed, +} + +fn hex(bytes: &[u8]) -> String { + bytes.iter().map(|b| format!("{b:02x}")).collect() +} + +impl JournalAuth { + /// The sidecar's name: a suffix on a single-file session's whole file name + /// (`.jsonl.mcp-task-journal.json`), or this name inside a segmented session's directory + /// name inside a segmented session's directory (one without a tenant codec). + pub const SIDECAR: &'static str = "mcp-task-journal.json"; + + /// The auth a sidecar's bytes describe; `None` when they describe none. + pub fn from_sidecar(bytes: &[u8]) -> Option { + let v: Value = serde_json::from_slice(bytes).ok()?; + let text = v.get("key")?.as_str()?; + if text.len() != 64 { + return None; + } + let mut key = [0u8; 32]; + for (i, byte) in key.iter_mut().enumerate() { + *byte = u8::from_str_radix(text.get(i * 2..i * 2 + 2)?, 16).ok()?; + } + let legacy = v + .get("legacy") + .and_then(Value::as_array) + .map(|macs| { + macs.iter() + .filter_map(|m| m.as_str().map(str::to_owned)) + .collect() + }) + .unwrap_or_default(); + Some(Self::Keyed { + key, + legacy: Arc::new(legacy), + }) + } + + /// A new key for a session whose entries so far are `existing` (`(kind, data)` of every custom + /// entry it holds, on any branch): its journal entries among them stay accepted. Returns the + /// auth and the sidecar bytes to write; `None` if no randomness is available. + pub fn fresh<'a>( + existing: impl IntoIterator, + ) -> Option<(Self, Vec)> { + let mut key = [0u8; 32]; + getrandom::fill(&mut key).ok()?; + let legacy: HashSet = existing + .into_iter() + .filter(|(kind, _)| is_journal_kind(kind)) + .filter_map(|(kind, data)| mac_of(&key, kind, data)) + .collect(); + let mut sorted: Vec<&String> = legacy.iter().collect(); + sorted.sort(); + let bytes = serde_json::to_vec(&json!({ "key": hex(&key), "legacy": sorted })).ok()?; + Some(( + Self::Keyed { + key, + legacy: Arc::new(legacy), + }, + bytes, + )) + } + + /// Bind a journal entry to its session before it is written (a `mac`, when keyed). + pub fn seal(&self, kind: &str, data: &mut Value) { + if let Self::Keyed { key, .. } = self + && let Some(mac) = mac_of(key, kind, data) + && let Value::Object(map) = data + { + map.insert("mac".into(), json!(mac)); + } + } + + /// Whether a journal entry read back may be replayed: planted or tampered entries are not. + pub fn accepts(&self, kind: &str, data: &Value) -> bool { + let Self::Keyed { key, legacy } = self else { + return true; + }; + let Some(want) = mac_of(key, kind, data) else { + return false; + }; + let sealed = data.get("mac").and_then(Value::as_str).is_some_and(|got| { + got.len() == want.len() + && got + .bytes() + .zip(want.bytes()) + .fold(0u8, |acc, (a, b)| acc | (a ^ b)) + == 0 + }); + sealed || legacy.contains(&want) + } +} + +fn mac_of(key: &[u8; 32], kind: &str, data: &Value) -> Option { + use hmac::{Hmac, Mac}; + let mut body = data.clone(); + if let Value::Object(map) = &mut body { + map.remove("mac"); + } + let mut mac = Hmac::::new_from_slice(key).ok()?; + mac.update(kind.as_bytes()); + mac.update(b"\n"); + mac.update(serde_json::to_string(&body).ok()?.as_bytes()); + Some(hex(&mac.finalize().into_bytes())) +} + /// Where journal writes go: `(kind, data)`, `data` carrying `sessionId` so a write that outlives a /// session switch is dropped instead of landing in the wrong session's file. pub type JournalTx = mpsc::UnboundedSender<(String, Value)>; @@ -66,6 +220,47 @@ pub struct Resolved { pub is_error: bool, #[serde(default, skip_serializing_if = "Vec::is_empty")] pub images: Vec, + /// Not an answer: the server could not be reached, so the call is told its result is pending + /// (see [`pending_placeholder`]). Never journaled as a result; the task stays pending and is + /// tried again. + #[serde(default, skip_serializing_if = "std::ops::Not::not")] + pub pending: bool, +} + +/// How a "result pending" placeholder `tool_result` reads, for the model. Only the text: what makes +/// a `tool_result` a placeholder is a [`PLACEHOLDER_ENTRY_KIND`] entry for its call, never this. +pub const PENDING_PREFIX: &str = "[MCP task result pending] "; + +/// The placeholder answer for a task whose server could not be reached. Whoever puts it on the path +/// journals [`placeholder_entry`] for it, so the call stays unanswered and the real result replaces +/// it when it lands. +pub fn pending_placeholder(task: &PendingTask, why: &str) -> Resolved { + Resolved { + tool_use_id: task.tool_use_id.clone(), + name: task.name.clone(), + content: format!( + "{PENDING_PREFIX}{why}. Task `{}` may still be running on the server; its result will \ + be delivered on a later turn once the server is reachable.", + task.record.task_id + ), + is_error: true, + images: Vec::new(), + pending: true, + } +} + +/// The journal data marking `tool_use_id`'s `tool_result` as a placeholder. +pub fn placeholder_entry(tool_use_id: &str, session_id: &str) -> Value { + json!({ "toolUseId": tool_use_id, "sessionId": session_id }) +} + +/// The `tool_use` ids whose `tool_result` on the path is a placeholder, from the journal's +/// [`PLACEHOLDER_ENTRY_KIND`] entries. +pub fn placeholder_ids(journal: &[Value]) -> HashSet { + journal + .iter() + .filter_map(|entry| entry["toolUseId"].as_str().map(str::to_owned)) + .collect() } /// The journal data for a task record emitted by tool `tool_use_id` in session `session_id`. @@ -76,12 +271,19 @@ pub fn task_entry(record: &Value, tool_use_id: &str, session_id: &str) -> Value data } -/// Every `tool_use` id answered on the path: by a `tool_result`, or by a journaled result. -fn answered_ids(messages: &[Message], results: &[Resolved]) -> HashSet { +/// Every `tool_use` id answered on the path: by a `tool_result` that is not a placeholder, or by a +/// journaled result. +fn answered_ids( + messages: &[Message], + results: &[Resolved], + placeholders: &HashSet, +) -> HashSet { let mut ids: HashSet = results.iter().map(|r| r.tool_use_id.clone()).collect(); for message in messages { for block in &message.content { - if let ContentBlock::ToolResult { tool_use_id, .. } = block { + if let ContentBlock::ToolResult { tool_use_id, .. } = block + && !placeholders.contains(tool_use_id) + { ids.insert(tool_use_id.clone()); } } @@ -89,7 +291,8 @@ fn answered_ids(messages: &[Message], results: &[Resolved]) -> HashSet { ids } -/// Journaled results on the active path, oldest first. +/// Journaled results on the active path, oldest first. `journal` holds only entries the session's +/// store accepted (see [`JournalAuth`]). pub fn results(journal: &[Value]) -> Vec { journal .iter() @@ -97,14 +300,16 @@ pub fn results(journal: &[Value]) -> Vec { .collect() } -/// The journaled tasks of `session_id` that nothing on the path answers yet. +/// The journaled tasks of `session_id` that nothing on the path answers yet. `tasks` holds only +/// entries the session's store accepted (see [`JournalAuth`]). pub fn pending( messages: &[Message], tasks: &[Value], results: &[Resolved], + placeholders: &HashSet, session_id: &str, ) -> Vec { - let answered = answered_ids(messages, results); + let answered = answered_ids(messages, results, placeholders); let calls: HashMap<&str, &str> = messages .iter() .filter(|m| m.role == Role::Assistant) @@ -133,9 +338,14 @@ pub fn pending( .filter(|(id, _)| !answered.contains(id)) .filter_map(|(id, entry)| { let name = calls.get(id.as_str())?; + let record: McpTaskRecord = serde_json::from_value(entry).ok()?; + // The record must describe the call it claims to answer: the same server and tool. + if *name != crate::tools::mcp::registered_name(&record.server, &record.tool) { + return None; + } Some(PendingTask { name: (*name).to_owned(), - record: serde_json::from_value(entry).ok()?, + record, tool_use_id: id, }) }) @@ -146,14 +356,42 @@ pub fn pending( /// holding its `tool_use`: exactly where the model would have seen it, with no turn added, so role /// alternation holds whether the call is at the tip (the new prompt's turn) or further back (an /// aborted prompt's turn). A result already answered on the path, or with no user turn after its -/// call yet, is skipped. Returns the results spliced. -pub fn splice(messages: &mut [Message], results: &[Resolved]) -> Vec { - let answered = answered_ids(messages, &[]); +/// call yet, is skipped. `placeholders` names the calls whose `tool_result` is a placeholder (see +/// [`placeholder_ids`]): the real result replaces it where it is. Returns the results spliced. +pub fn splice( + messages: &mut [Message], + results: &[Resolved], + placeholders: &HashSet, +) -> Vec { + let answered = answered_ids(messages, &[], placeholders); let mut spliced = Vec::new(); for result in results { if answered.contains(&result.tool_use_id) { continue; } + // A placeholder already in place: the real result replaces it where it is (once: a block + // already holding this result is left alone); a second placeholder adds nothing. + let placeholder = messages + .iter_mut() + .flat_map(|m| m.content.iter_mut()) + .find(|b| { + matches!(b, ContentBlock::ToolResult { tool_use_id, .. } + if *tool_use_id == result.tool_use_id) + }); + if let Some(block) = placeholder { + let applied = matches!(block, ContentBlock::ToolResult { content, is_error, .. } + if *is_error == result.is_error && &**content == result.content.as_str()); + if !result.pending && !applied { + *block = ContentBlock::ToolResult { + tool_use_id: result.tool_use_id.clone(), + content: result.content.clone().into(), + is_error: result.is_error, + images: result.images.clone(), + }; + spliced.push(result.clone()); + } + continue; + } let Some(call) = messages.iter().position(|m| { m.role == Role::Assistant && m.content.iter().any( @@ -243,9 +481,18 @@ impl Resumer { let work = async move { let resolved = resolve(&task, &catalog, &task_host, &emit, &journal, &session_id).await; - let mut data = serde_json::to_value(&resolved).unwrap_or_default(); - data["sessionId"] = json!(session_id); - let _ = journal.send((RESULT_ENTRY_KIND.to_owned(), data)); + // A pending placeholder is not an answer: the task stays journaled as pending, and + // the placeholder is marked as one wherever it lands on the path. + if resolved.pending { + let _ = journal.send(( + PLACEHOLDER_ENTRY_KIND.to_owned(), + placeholder_entry(&resolved.tool_use_id, &session_id), + )); + } else { + let mut data = serde_json::to_value(&resolved).unwrap_or_default(); + data["sessionId"] = json!(session_id); + let _ = journal.send((RESULT_ENTRY_KIND.to_owned(), data)); + } if let Ok(mut state) = state.lock() { state .outstanding @@ -266,6 +513,22 @@ impl Resumer { &self.session_id } + /// Whether this resumer has finished and some task's server was unreachable, so a fresh + /// resumer should try again (at the next prompt). + pub fn needs_retry(&self) -> bool { + self.state + .lock() + .is_ok_and(|s| s.outstanding.is_empty() && s.done.iter().any(|r| r.pending)) + } + + /// The definitive results this resumer resolved (journaled, or about to be). + pub fn finals(&self) -> Vec { + self.state + .lock() + .map(|s| s.done.iter().filter(|r| !r.pending).cloned().collect()) + .unwrap_or_default() + } + /// The tasks still being polled. pub fn outstanding(&self) -> Vec { self.state @@ -357,16 +620,24 @@ async fn resolve( result }; let (result, ()) = tokio::join!(work, forward); - let (content, is_error, images) = match result { - Ok(output) => (output.text, false, output.images), - Err(e) => (e.to_string(), true, Vec::new()), - }; - Resolved { - tool_use_id: task.tool_use_id.clone(), - name: task.name.clone(), - content, - is_error, - images, + match result { + Ok(output) => Resolved { + tool_use_id: task.tool_use_id.clone(), + name: task.name.clone(), + content: output.text, + is_error: false, + images: output.images, + pending: false, + }, + Err(crate::tools::mcp::ResumeError::Unreachable(why)) => pending_placeholder(task, &why), + Err(crate::tools::mcp::ResumeError::Final(e)) => Resolved { + tool_use_id: task.tool_use_id.clone(), + name: task.name.clone(), + content: e.to_string(), + is_error: true, + images: Vec::new(), + pending: false, + }, } } @@ -374,11 +645,13 @@ async fn resolve( /// the resumer is still polling, then splice every resolved result into the transcript and report /// each through the run's own event sink (a `tool_start` and a `tool_end`, so the session's live /// stats and lifecycle observer see the call finish). An abort returns early with nothing spliced; -/// the tasks keep polling and the next prompt waits again. +/// the tasks keep polling and the next prompt waits again. `placeholders` are the journal's +/// placeholder marks (see [`placeholder_ids`]); this resumer's own are added. pub async fn resume_into_turn( session: &mut agent_core::Session, resumer: &Resumer, journaled: &[Resolved], + placeholders: &HashSet, sink: &mut F, cancel: &CancellationToken, ) { @@ -416,6 +689,12 @@ pub async fn resume_into_turn( () = cancel.cancelled() => return, done = resumer.wait() => done, }; + let mut placeholders = placeholders.clone(); + placeholders.extend( + done.iter() + .filter(|r| r.pending) + .map(|r| r.tool_use_id.clone()), + ); let mut results: Vec = journaled.to_vec(); for result in done { if !results.iter().any(|r| r.tool_use_id == result.tool_use_id) { @@ -424,7 +703,7 @@ pub async fn resume_into_turn( } // Only touch the transcript when a result has somewhere to go: `make_mut` copies a shared // history, and an untouched turn must leave `session.messages` exactly as it was. - let answered = answered_ids(&session.messages, &[]); + let answered = answered_ids(&session.messages, &[], &placeholders); let calls: HashSet<&str> = session .messages .iter() @@ -440,7 +719,7 @@ pub async fn resume_into_turn( return; } let messages: &mut Vec = Arc::make_mut(&mut session.messages); - let spliced = splice(messages, &results); + let spliced = splice(messages, &results, &placeholders); for result in spliced { if !outstanding .iter() @@ -465,6 +744,94 @@ pub async fn resume_into_turn( mod tests { use super::*; + fn none() -> HashSet { + HashSet::new() + } + + fn keyed() -> JournalAuth { + JournalAuth::fresh([]).unwrap().0 + } + + /// Under a per-session key, a journal line planted without it (or altered after sealing) is + /// refused; an entry sealed under another session's key is too. + #[test] + fn a_keyed_journal_refuses_unsealed_tampered_and_foreign_entries() { + let auth = keyed(); + let planted = task_entry(&record("t1"), "tu1", "s"); + assert!(!auth.accepts(TASK_ENTRY_KIND, &planted)); + let mut good = planted.clone(); + auth.seal(TASK_ENTRY_KIND, &mut good); + assert!(auth.accepts(TASK_ENTRY_KIND, &good)); + // The MAC covers the kind too: a sealed record is not a sealed result. + assert!(!auth.accepts(RESULT_ENTRY_KIND, &good)); + let mut tampered = good.clone(); + tampered["taskId"] = json!("someone-elses"); + assert!(!auth.accepts(TASK_ENTRY_KIND, &tampered)); + assert!(!keyed().accepts(TASK_ENTRY_KIND, &good)); + } + + /// The key travels as the sidecar's bytes; the entries a pre-key session already held when its + /// key was made stay accepted (once made, the legacy set never grows). + #[test] + fn a_fresh_key_round_trips_through_its_sidecar_and_keeps_existing_entries() { + let old = task_entry(&record("t1"), "tu1", "s"); + let custom = json!({ "anything": 1 }); + let (auth, bytes) = + JournalAuth::fresh([(TASK_ENTRY_KIND, &old), ("someone_elses_kind", &custom)]).unwrap(); + assert!(auth.accepts(TASK_ENTRY_KIND, &old)); + let JournalAuth::Keyed { legacy, .. } = &auth else { + panic!("keyed") + }; + assert_eq!(legacy.len(), 1, "only journal kinds are carried over"); + let reread = JournalAuth::from_sidecar(&bytes).unwrap(); + assert!(reread.accepts(TASK_ENTRY_KIND, &old)); + let mut sealed = task_entry(&record("t2"), "tu2", "s"); + auth.seal(TASK_ENTRY_KIND, &mut sealed); + assert!(reread.accepts(TASK_ENTRY_KIND, &sealed)); + assert!(!reread.accepts(TASK_ENTRY_KIND, &task_entry(&record("t3"), "tu3", "s"))); + assert!(JournalAuth::from_sidecar(b"{\"key\":\"short\"}").is_none()); + } + + /// Sealed storage authenticates every line itself: nothing is sealed or refused on top. + #[test] + fn storage_and_unkeyed_journals_accept_entries_as_they_are() { + for auth in [JournalAuth::Storage, JournalAuth::Unkeyed] { + let mut entry = task_entry(&record("t1"), "tu1", "s"); + auth.seal(TASK_ENTRY_KIND, &mut entry); + assert!(entry.get("mac").is_none()); + assert!(auth.accepts(TASK_ENTRY_KIND, &entry)); + } + } + + /// A placeholder is known by its journal mark, never by its text: a real result that happens to + /// begin like one answers its call, and a marked one does not. + #[test] + fn a_placeholder_is_known_by_its_mark_not_its_content() { + let lookalike = format!("{PENDING_PREFIX}but this is the real answer"); + let messages = vec![ + Message::user("go"), + call("tu1"), + Message::tool_result("tu1", lookalike.as_str(), false), + ]; + let journal = vec![task_entry(&record("t1"), "tu1", "s")]; + assert!(pending(&messages, &journal, &[], &none(), "s").is_empty()); + let marked: HashSet = placeholder_ids(&[placeholder_entry("tu1", "s")]); + assert_eq!(pending(&messages, &journal, &[], &marked, "s").len(), 1); + + // The real result replaces a marked placeholder where it is, once; an unmarked look-alike + // is an answer and stays. + let mut spliced = messages.clone(); + assert!(splice(&mut spliced, &[resolved("tu1")], &none()).is_empty()); + assert_eq!(spliced, messages); + assert_eq!(splice(&mut spliced, &[resolved("tu1")], &marked).len(), 1); + assert!(matches!( + &spliced[2].content[0], + ContentBlock::ToolResult { content, .. } if &**content == "done" + )); + assert_eq!(spliced[2].content.len(), 1); + assert!(splice(&mut spliced, &[resolved("tu1")], &marked).is_empty()); + } + fn record(task_id: &str) -> Value { json!({ "server": "s", "tool": "job", "taskId": task_id, "createdAtMs": 0, "ttlMs": null, @@ -482,17 +849,34 @@ mod tests { content: "done".into(), is_error: false, images: Vec::new(), + pending: false, } } + /// A record that does not describe the call it names (another server or tool) is not + /// resumed: whoever wrote it, it cannot make the resumer poll something the call never was. + #[test] + fn pending_ignores_a_record_that_does_not_match_its_call() { + let messages = vec![Message::user("go"), call("tu1")]; + let mut other = record("t1"); + other["server"] = json!("elsewhere"); + let journal = vec![(task_entry(&other, "tu1", "s"))]; + assert!(pending(&messages, &journal, &[], &none(), "s").is_empty()); + let journal = vec![(task_entry(&record("t1"), "tu1", "s"))]; + assert_eq!(pending(&messages, &journal, &[], &none(), "s").len(), 1); + } + /// F9: only the session that journaled a task may resume it; a copy of the journal under /// another session id (a fork, a copied file) leaves the call to the generic repair. #[test] fn pending_ignores_tasks_journaled_by_another_session() { let messages = vec![Message::user("go"), call("tu1")]; - let journal = vec![task_entry(&record("t1"), "tu1", "parent")]; - assert!(pending(&messages, &journal, &[], "fork").is_empty()); - assert_eq!(pending(&messages, &journal, &[], "parent").len(), 1); + let journal = vec![(task_entry(&record("t1"), "tu1", "parent"))]; + assert!(pending(&messages, &journal, &[], &none(), "fork").is_empty()); + assert_eq!( + pending(&messages, &journal, &[], &none(), "parent").len(), + 1 + ); } /// F1: a journaled call is pending wherever it sits on the path, not only at the tip (an @@ -505,11 +889,11 @@ mod tests { Message::user("p1"), Message::assistant(Vec::new()), ]; - let journal = vec![task_entry(&record("t1"), "tu1", "s")]; - let found = pending(&messages, &journal, &[], "s"); + let journal = vec![(task_entry(&record("t1"), "tu1", "s"))]; + let found = pending(&messages, &journal, &[], &none(), "s"); assert_eq!(found.len(), 1); assert_eq!(found[0].record.task_id, "t1"); - assert!(pending(&messages, &journal, &[resolved("tu1")], "s").is_empty()); + assert!(pending(&messages, &journal, &[resolved("tu1")], &none(), "s").is_empty()); messages[2].content.insert( 0, ContentBlock::ToolResult { @@ -519,7 +903,7 @@ mod tests { images: Vec::new(), }, ); - assert!(pending(&messages, &journal, &[], "s").is_empty()); + assert!(pending(&messages, &journal, &[], &none(), "s").is_empty()); } /// F10: the latest record for a call wins (it carries the keys answered since creation). @@ -529,10 +913,10 @@ mod tests { let mut later = record("t1"); later["answered"] = json!(["name"]); let journal = vec![ - task_entry(&record("t1"), "tu1", "s"), - task_entry(&later, "tu1", "s"), + (task_entry(&record("t1"), "tu1", "s")), + (task_entry(&later, "tu1", "s")), ]; - let found = pending(&messages, &journal, &[], "s"); + let found = pending(&messages, &journal, &[], &none(), "s"); assert_eq!(found[0].record.answered, vec!["name".to_owned()]); } @@ -547,7 +931,7 @@ mod tests { Message::assistant(Vec::new()), Message::user("p2"), ]; - let spliced = splice(&mut messages, &[resolved("tu1")]); + let spliced = splice(&mut messages, &[resolved("tu1")], &none()); assert_eq!(spliced.len(), 1); assert_eq!(messages.len(), 5, "no turn added"); assert!(matches!( @@ -558,6 +942,6 @@ mod tests { assert_ne!(pair[0].role, pair[1].role); } // Idempotent: already answered, so a second splice adds nothing. - assert!(splice(&mut messages, &[resolved("tu1")]).is_empty()); + assert!(splice(&mut messages, &[resolved("tu1")], &none()).is_empty()); } } diff --git a/crates/agent/src/serve.rs b/crates/agent/src/serve.rs index e24d3fcf..d300b639 100644 --- a/crates/agent/src/serve.rs +++ b/crates/agent/src/serve.rs @@ -2031,14 +2031,27 @@ impl Persistence { .unwrap_or(&[]) } - /// Custom entries of `kind` on the active path; empty without persistence. - fn active_custom(&self, kind: &str) -> Vec { + /// MCP task journal entries of `kind` on the active path that the session accepts (see + /// `SessionStore::active_journal`); empty without persistence. + fn active_journal(&self, kind: &str) -> Vec { self.store .as_ref() - .map(|store| store.active_custom(kind)) + .map(|store| store.active_journal(kind)) .unwrap_or_default() } + /// Append an MCP task journal entry, sealed for the session (see + /// `SessionStore::append_journal`). Same persistence-required contract as `append_custom`. + fn append_journal(&mut self, kind: &str, data: Value) -> std::io::Result { + let store = self.store.as_mut().ok_or_else(|| { + std::io::Error::new( + std::io::ErrorKind::Unsupported, + "no session persistence configured (start serve with --session-file or --session-dir)", + ) + })?; + store.append_journal(kind, data) + } + /// Every user-turn message anywhere in the session's tree — every branch, not just the active path /// — as `(id, Message)` pairs (empty unless persistence is configured). What `get_fork_messages` /// surfaces, matching pi's own whole-tree `getUserMessagesForForking` rather than only the active @@ -3401,7 +3414,7 @@ pub(crate) async fn serve_session( ($kind:expr, $data:expr) => {{ let data: Value = $data; if data["sessionId"].as_str() == Some(persistence.session_id()) { - match persistence.append_custom($kind, data) { + match persistence.append_journal(&$kind, data) { Err(e) if e.kind() != std::io::ErrorKind::Unsupported => { tracing::warn!(error = %e, "MCP task journal entry failed to persist"); } @@ -3522,18 +3535,32 @@ pub(crate) async fn serve_session( // MCP tasks a previous process left in flight resume now, in the background, not on the next // prompt (see `crate::mcp_resume`). Restarted for a session switched to later. + // + // `$carried`: results the previous resumer already resolved, which count as answers whether or + // not their journal writes have landed yet (they travel through `mcp_task_rx`, which a prompt + // can overtake), so a restart never polls a finished task again. macro_rules! start_mcp_resumer { - () => {{ - let journal = persistence.active_custom(crate::mcp_resume::TASK_ENTRY_KIND); - let results = crate::mcp_resume::results( - &persistence.active_custom(crate::mcp_resume::RESULT_ENTRY_KIND), + () => { + start_mcp_resumer!(&[]) + }; + ($carried:expr) => {{ + let journal = persistence.active_journal(crate::mcp_resume::TASK_ENTRY_KIND); + let mut results = crate::mcp_resume::results( + &persistence.active_journal(crate::mcp_resume::RESULT_ENTRY_KIND), + ); + results.extend_from_slice($carried); + let placeholders = crate::mcp_resume::placeholder_ids( + &persistence.active_journal(crate::mcp_resume::PLACEHOLDER_ENTRY_KIND), ); - let pending = crate::mcp_resume::pending( + let mut pending = crate::mcp_resume::pending( &session.messages, &journal, &results, + &placeholders, persistence.session_id(), ); + // Only a configured server is ever polled. + pending.retain(|p| cfg.mcp_catalog.has_server(&p.record.server)); if pending.is_empty() { crate::mcp_resume::Resumer::idle(persistence.session_id()) } else { @@ -4474,12 +4501,22 @@ pub(crate) async fn serve_session( let images = parse_images(cmd.get("images")); // MCP tasks a previous process left in flight are waited on and answered before this // turn's model call (see `crate::mcp_resume`), by the first attempt only. + // A resumer that finished with a task's server unreachable left that task pending: + // a fresh one tries it again now (its finished results carry over). + let mut mcp_carried = Vec::new(); if mcp_resumer.session_id() != persistence.session_id() { mcp_resumer = start_mcp_resumer!(); + } else if mcp_resumer.needs_retry() { + mcp_carried = mcp_resumer.finals(); + mcp_resumer = start_mcp_resumer!(&mcp_carried); } let mut mcp_resume_first = true; - let mcp_journaled = crate::mcp_resume::results( - &persistence.active_custom(crate::mcp_resume::RESULT_ENTRY_KIND), + let mut mcp_journaled = crate::mcp_resume::results( + &persistence.active_journal(crate::mcp_resume::RESULT_ENTRY_KIND), + ); + mcp_journaled.extend(mcp_carried); + let mcp_placeholders = crate::mcp_resume::placeholder_ids( + &persistence.active_journal(crate::mcp_resume::PLACEHOLDER_ENTRY_KIND), ); if images.is_empty() { session.user(message); @@ -4778,6 +4815,7 @@ pub(crate) async fn serve_session( let resume_first = std::mem::take(&mut mcp_resume_first); let resumer = &mcp_resumer; let journaled = &mcp_journaled; + let placeholders = &mcp_placeholders; let run_cancel = cancel.clone(); let run = crate::tools::mcp::with_session_host(mcp_host.clone(), async { if resume_first { @@ -4785,6 +4823,7 @@ pub(crate) async fn serve_session( &mut session, resumer, journaled, + placeholders, &mut sink, &run_cancel, ) diff --git a/crates/agent/src/session_store.rs b/crates/agent/src/session_store.rs index 568d6b5d..6a8bec69 100644 --- a/crates/agent/src/session_store.rs +++ b/crates/agent/src/session_store.rs @@ -345,7 +345,8 @@ pub fn search_sessions(sessions: Vec, query: Option<&str>) -> Vec, query: Option<&str>) -> Vec Message { /// result journaled on the path (a `mcp_task_result` custom entry, see `crate::mcp_resume`) spliced /// in front of the first user turn after the call it answers. The splice is part of the transcript, /// not just of the next request, so `get_messages`, the HTML export and `run --continue` agree with -/// what `serve` sent the model. Custom entries are otherwise invisible here. -fn materialize(nodes: &HashMap, path: &[String]) -> Vec { +/// what `serve` sent the model. Only journal entries `journal` accepts count; custom entries are +/// otherwise invisible here. +fn materialize( + nodes: &HashMap, + path: &[String], + journal: &crate::mcp_resume::JournalAuth, +) -> Vec { + use crate::mcp_resume::{PLACEHOLDER_ENTRY_KIND, RESULT_ENTRY_KIND}; let mut messages = Vec::new(); let mut results = Vec::new(); + let mut placeholders = Vec::new(); for id in path { match nodes.get(id).map(|n| n.content.as_ref()) { Some(NodeContent::Message(m)) => messages.push(m.clone()), Some(NodeContent::Custom { kind, data }) - if kind == crate::mcp_resume::RESULT_ENTRY_KIND => + if kind == RESULT_ENTRY_KIND && journal.accepts(kind, data) => { if let Ok(result) = serde_json::from_value(data.clone()) { results.push(result); } } + Some(NodeContent::Custom { kind, data }) + if kind == PLACEHOLDER_ENTRY_KIND && journal.accepts(kind, data) => + { + placeholders.push(data.clone()); + } _ => {} } } if !results.is_empty() { - crate::mcp_resume::splice(&mut messages, &results); + let placeholders = crate::mcp_resume::placeholder_ids(&placeholders); + crate::mcp_resume::splice(&mut messages, &results, &placeholders); } messages } +/// The journal auth for a session whose log is `log` (see [`crate::mcp_resume::JournalAuth`]): +/// sealed storage needs none; otherwise the key in the session's own sidecar, if it has one yet. A +/// session without one (it never journaled a task, or predates per-session keys) is unkeyed until +/// its first journal write makes one ([`make_journal_key`]). Reading never writes. +fn journal_auth(log: &Log) -> crate::mcp_resume::JournalAuth { + use crate::mcp_resume::JournalAuth; + let Some(path) = log.journal_key_path() else { + return JournalAuth::Storage; + }; + read_journal_key(&path) + .or_else(|| { + log.legacy_journal_key_path() + .and_then(|p| read_journal_key(&p)) + }) + .unwrap_or(JournalAuth::Unkeyed) +} + +fn read_journal_key(path: &Path) -> Option { + crate::mcp_resume::JournalAuth::from_sidecar(&fs::read(path).ok()?) +} + +/// How long a first journal write waits for another writer making the same session's key. +const JOURNAL_KEY_LOCK_WAIT: std::time::Duration = std::time::Duration::from_secs(5); + +/// Make the journal key of a session that has none yet, at its first journal write: the entries it +/// already holds (`nodes`, any branch; from before per-session keys) stay accepted. `None` when no +/// key belongs in a sidecar or none could be written (the write then goes out unsealed, as before). +/// +/// Every first writer (threads of this process, other processes) must end up with the *same* key, +/// or the loser's entries would be refused on replay. So the key is made under an exclusive lock on +/// the sidecar itself (`file_lock`), and re-read once the lock is held: whoever got there first has +/// written it, and everyone after uses theirs. The key is written through the locked descriptor, in +/// place: no temp file for two writers to share, and none left beside the session. A sidecar that +/// does not parse (torn, or overwritten) is replaced the same way. A key file under the +/// pre-suffix name (`legacy_journal_key_path`) is carried over rather than replaced. +fn make_journal_key( + log: &Log, + nodes: &HashMap, +) -> Option { + use crate::mcp_resume::JournalAuth; + use std::io::{Seek, SeekFrom}; + let path = log.journal_key_path()?; + let failed = |e: &dyn std::fmt::Display| { + tracing::warn!(path = %path.display(), error = %e, "MCP task journal key not written"); + }; + // Never write a key through a link someone planted at the sidecar's name. + if fs::symlink_metadata(&path).is_ok_and(|m| m.file_type().is_symlink()) { + failed(&"the sidecar is a symlink"); + return None; + } + let deadline = std::time::Instant::now() + JOURNAL_KEY_LOCK_WAIT; + let lock = loop { + match crate::file_lock::try_lock(&path) { + Ok(Some(lock)) => break lock, + Ok(None) if std::time::Instant::now() < deadline => { + std::thread::sleep(std::time::Duration::from_millis(2)); + } + Ok(None) => { + failed(&"another writer held its lock too long"); + return None; + } + Err(e) => { + failed(&e); + return None; + } + } + }; + if let Some(theirs) = read_journal_key(&path) { + return Some(theirs); + } + let (auth, bytes) = match log.legacy_journal_key_path().and_then(|p| { + let bytes = fs::read(&p).ok()?; + Some((JournalAuth::from_sidecar(&bytes)?, bytes)) + }) { + Some(carried) => carried, + None => { + let entries = nodes.values().filter_map(|node| match &*node.content { + NodeContent::Custom { kind, data } => Some((kind.as_str(), data)), + NodeContent::Message(_) => None, + }); + JournalAuth::fresh(entries)? + } + }; + let written = (|| -> std::io::Result<()> { + let mut file = lock.file(); + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + file.set_permissions(fs::Permissions::from_mode(0o600))?; + } + file.set_len(0)?; + file.seek(SeekFrom::Start(0))?; + file.write_all(&bytes)?; + file.sync_all()?; + fsync_dir(&path) + })(); + match written { + Ok(()) => Some(auth), + Err(e) => { + failed(&e); + None + } + } +} + fn path_from_root(nodes: &HashMap, tip: Option<&str>) -> Vec { let mut rev = Vec::new(); let mut visited = HashSet::new(); @@ -829,6 +949,9 @@ pub struct SessionStore { /// tree logic below never looks inside it: every read goes through `log.lines()` and every write /// through `log.append`/`log.replace_all`. log: Log, + /// How this session's MCP task journal entries are sealed and checked (see + /// `crate::mcp_resume::JournalAuth`). + journal: crate::mcp_resume::JournalAuth, meta: SessionMeta, /// How many messages are already on disk (on the active path) — the append cursor. persisted: usize, @@ -898,7 +1021,9 @@ impl SessionStore { let mut header = Vec::new(); write_line(&mut header, &Entry::Session(meta.clone()))?; let log = Log::create(path, &meta.id, layout, &header)?; + let journal = journal_auth(&log); Ok(Self { + journal, log, meta, persisted: 0, @@ -1216,7 +1341,8 @@ impl SessionStore { // A custom entry (`NodeContent::Custom`) contributes nothing here — it's a real, positioned // node in `active`'s chain (see `Entry::Custom`'s doc comment), but not a message, so // `as_message` filters it out of the materialized `Session.messages`/LLM context. - let messages: Vec = materialize(&nodes, &active); + let journal = journal_auth(&log); + let messages: Vec = materialize(&nodes, &active, &journal); let persisted = messages.len(); let mut session = Session::new(); session.messages = Arc::new(messages); @@ -1281,6 +1407,7 @@ impl SessionStore { } Ok(( Self { + journal, log, meta, persisted, @@ -1342,10 +1469,48 @@ impl SessionStore { /// The `data` of every custom entry of `kind` on the active path, oldest first — see /// [`Self::append_custom`]. pub fn active_custom(&self, kind: &str) -> Vec { + self.active_custom_where(kind, |_| true) + } + + /// The MCP task journal entries of `kind` on the active path that this session's journal auth + /// accepts (see [`crate::mcp_resume::JournalAuth`]): planted or tampered lines are left out. + pub fn active_journal(&self, kind: &str) -> Vec { + self.active_custom_where(kind, |data| self.journal.accepts(kind, data)) + } + + /// Append an MCP task journal entry, sealed for this session (see + /// [`crate::mcp_resume::JournalAuth`]); otherwise [`append_custom`](Self::append_custom). + pub fn append_journal( + &mut self, + kind: &str, + mut data: serde_json::Value, + ) -> std::io::Result { + // The key on disk is the session's key: another writer (a thread, another process) may + // have made it since this store opened. With none there (none yet, one under the old name + // to carry over, or one removed or torn since), make it now. One small read per journal + // write. + if !self.log.read_only() + && let Some(key) = self.log.journal_key_path() + && let Some(auth) = + read_journal_key(&key).or_else(|| make_journal_key(&self.log, &self.nodes)) + { + self.journal = auth; + } + self.journal.seal(kind, &mut data); + self.append_custom(kind, data) + } + + fn active_custom_where( + &self, + kind: &str, + keep: impl Fn(&serde_json::Value) -> bool, + ) -> Vec { self.active .iter() .filter_map(|id| match self.nodes.get(id).map(|n| n.content.as_ref()) { - Some(NodeContent::Custom { kind: k, data }) if k == kind => Some(data.clone()), + Some(NodeContent::Custom { kind: k, data }) if k == kind && keep(data) => { + Some(data.clone()) + } _ => None, }) .collect() @@ -2199,7 +2364,7 @@ impl SessionStore { /// a redundant marker every time a client re-confirms the current position. pub fn switch_active(&mut self, target_id: &str) -> std::io::Result> { if self.active.last().is_some_and(|id| id == target_id) { - return Ok(materialize(&self.nodes, &self.active)); + return Ok(materialize(&self.nodes, &self.active, &self.journal)); } if !self.nodes.contains_key(target_id) { return Err(std::io::Error::new( @@ -2217,7 +2382,7 @@ impl SessionStore { self.log.append(&buf)?; let active = path_from_root(&self.nodes, Some(target_id)); - let messages: Vec = materialize(&self.nodes, &active); + let messages: Vec = materialize(&self.nodes, &active, &self.journal); self.persisted = messages.len(); self.active = active; Ok(messages) @@ -2297,7 +2462,7 @@ impl SessionStore { self.branch_summary_details .insert(entry_id.clone(), details_for_index); self.active = path_from_root(&self.nodes, Some(&entry_id)); - let messages: Vec = materialize(&self.nodes, &self.active); + let messages: Vec = materialize(&self.nodes, &self.active, &self.journal); self.persisted = messages.len(); Ok(messages) } @@ -2389,7 +2554,7 @@ impl SessionStore { self.branch_summary_details .insert(entry_id.clone(), details_for_index); self.active = vec![entry_id]; - let messages: Vec = materialize(&self.nodes, &self.active); + let messages: Vec = materialize(&self.nodes, &self.active, &self.journal); self.persisted = messages.len(); Ok(messages) } @@ -2719,13 +2884,17 @@ fn move_sibling_memory(session_jsonl: &Path, dst_dir: &Path) { { let _ = fs::rename(&src_mem, dst_dir.join(name)); } - // The session's sidecar files (MCP Events state, MCP App view state) travel with it too. - for ext in SESSION_SIDECARS { - let events = session_jsonl.with_extension(ext); - if let Some(name) = events.file_name() - && events.is_file() + // The session's sidecar files (MCP Events state, MCP App view state, the journal key) travel + // with it too. + let sidecars = SESSION_SIDECARS + .iter() + .map(|ext| session_jsonl.with_extension(ext)) + .chain([journal_key_beside(session_jsonl)]); + for sidecar in sidecars { + if let Some(name) = sidecar.file_name() + && sidecar.is_file() { - let _ = fs::rename(&events, dst_dir.join(name)); + let _ = fs::rename(&sidecar, dst_dir.join(name)); } } } @@ -2735,8 +2904,9 @@ fn move_sibling_memory(session_jsonl: &Path, dst_dir: &Path) { /// the append-only log of undelivered events) and the MCP App view state (`tools::mcp_apps::Sidecar` /// — the attached view context, and the view replay store holding each kept view's HTML and full /// tool result). Each trashes, restores and is removed with its session: left behind, it would keep -/// a deleted session's content on disk. (A segmented session keeps its MCP App sidecars inside its -/// own directory, which moves whole.) +/// a deleted session's content on disk. The MCP task journal's key ([`journal_key_beside`]) goes +/// with them too, so its journal resumes wherever the session goes. (A segmented session keeps its +/// MCP App sidecars and its journal key inside its own directory, which moves whole.) const SESSION_SIDECARS: [&str; 4] = [ "mcp-events.json", "mcp-events.log", @@ -2744,6 +2914,16 @@ const SESSION_SIDECARS: [&str; 4] = [ "mcp-app-views.json", ]; +/// Where a single-file session's MCP task journal key lives: its whole file name plus a suffix +/// (`.jsonl.mcp-task-journal.json`), not `with_extension`, so sessions whose names differ +/// only after their last dot (`work.1`, `work.2`) never share one key. +fn journal_key_beside(session: &Path) -> PathBuf { + let mut name = session.as_os_str().to_owned(); + name.push("."); + name.push(crate::mcp_resume::JournalAuth::SIDECAR); + PathBuf::from(name) +} + /// The MCP Events subset — the only sidecars a segmented session keeps *beside* its directory. const MCP_EVENTS_SIDECARS: [&str; 2] = ["mcp-events.json", "mcp-events.log"]; @@ -2757,6 +2937,7 @@ fn remove_sibling_memory(session_jsonl: &Path) { for ext in SESSION_SIDECARS { let _ = fs::remove_file(session_jsonl.with_extension(ext)); } + let _ = fs::remove_file(journal_key_beside(session_jsonl)); } /// A directory of session files. `Clone` is just a `PathBuf` copy — cheap, and lets a caller move an @@ -4651,6 +4832,24 @@ pub(crate) fn remove_durably(path: &Path) -> std::io::Result<()> { } } +/// Tighten an existing session file to `0600` before writing more into it. `create_private` sets +/// the mode when a file is created, but a file that already exists keeps whatever mode it has (an +/// older version's, a restore's, or a `chmod`), and every append adds transcript, MCP task ids +/// (which a server may treat as bearer tokens, see `tools::mcp::McpTaskRecord`) included. One +/// `fstat` per append, beside the append's own `fsync`. +fn keep_private(file: &File) -> std::io::Result<()> { + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + if file.metadata()?.permissions().mode() & 0o077 != 0 { + file.set_permissions(fs::Permissions::from_mode(0o600))?; + } + } + #[cfg(not(unix))] + let _ = file; + Ok(()) +} + fn fsync_dir(path: &Path) -> std::io::Result<()> { if let Some(parent) = path.parent() { // Opening the directory read-only and `sync_all`ing it flushes its updated entries. Linux/macOS @@ -5288,6 +5487,29 @@ impl Log { } } + /// Where this session's MCP task journal key lives (see `crate::mcp_resume::JournalAuth`): + /// beside a single-file session, inside a segmented session's directory, and nowhere for one + /// sealed under a tenant codec, whose storage already authenticates every line. + fn journal_key_path(&self) -> Option { + match self { + Log::File { path, .. } => Some(journal_key_beside(path)), + Log::Segmented(s) if s.codec.is_some() => None, + Log::Segmented(s) => Some(s.dir.join(crate::mcp_resume::JournalAuth::SIDECAR)), + } + } + + /// The name a single-file session's journal key had before it became a suffix: + /// `with_extension`, which gave `work.1` and `work.2` one key file. Read, and carried over by the + /// first journal write, never written. + fn legacy_journal_key_path(&self) -> Option { + match self { + Log::File { path, .. } => { + Some(path.with_extension(crate::mcp_resume::JournalAuth::SIDECAR)) + } + Log::Segmented(_) => None, + } + } + /// Where this session's `/session` working memory lives. /// /// For the segmented layout that is `/memory`, *inside* the session directory — not a @@ -5463,6 +5685,7 @@ impl Log { match self { Log::File { path, .. } => { let mut f = OpenOptions::new().append(true).open(path)?; + keep_private(&f)?; f.write_all(bytes)?; // `flush` only pushes past our buffer into the OS; `sync_all` forces the bytes to // disk, which is what the module's crash-safety claim actually requires. The parent @@ -5816,6 +6039,7 @@ impl SegLog { let before = self.target_len; let write = (|| -> std::io::Result<()> { let mut f = OpenOptions::new().append(true).open(&path)?; + keep_private(&f)?; f.write_all(&buf)?; f.flush()?; f.sync_all() @@ -6267,6 +6491,311 @@ mod tests { ); } + /// A session's first journal write makes its MCP task journal key, beside it (`0600`); a + /// session that never journals gets no file. Entries it seals read back after a reopen; a line + /// planted into the file without the key does not. + #[test] + fn a_session_keeps_its_journal_key_beside_it_and_refuses_planted_entries() { + use crate::mcp_resume::TASK_ENTRY_KIND; + let dir = tmpdir(); + let path = dir.path().join("s.jsonl"); + let mut store = SessionStore::create(path.clone(), SessionMeta::new("/w", "m")).unwrap(); + let key = journal_key_beside(&path); + assert!(!key.exists(), "no key until the session journals something"); + store + .append_journal(TASK_ENTRY_KIND, serde_json::json!({ "taskId": "real" })) + .unwrap(); + assert!(key.is_file()); + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + assert_eq!( + fs::metadata(&key).unwrap().permissions().mode() & 0o777, + 0o600 + ); + } + store + .append_custom(TASK_ENTRY_KIND, serde_json::json!({ "taskId": "planted" })) + .unwrap(); + drop(store); + let (store, _) = SessionStore::open(path).unwrap(); + let tasks = store.active_journal(TASK_ENTRY_KIND); + assert_eq!(tasks.len(), 1, "{tasks:?}"); + assert_eq!(tasks[0]["taskId"], "real"); + assert_eq!(store.active_custom(TASK_ENTRY_KIND).len(), 2); + } + + /// A session from before per-session keys (journal entries, no key file) keeps the entries it + /// already has: opening it writes nothing, and its next journal write makes the key with those + /// entries accepted. A line planted after that is refused. + #[test] + fn a_session_without_a_journal_key_keeps_its_existing_entries() { + use crate::mcp_resume::{RESULT_ENTRY_KIND, TASK_ENTRY_KIND}; + let dir = tmpdir(); + let path = dir.path().join("s.jsonl"); + let mut store = SessionStore::create(path.clone(), SessionMeta::new("/w", "m")).unwrap(); + store + .append_custom(TASK_ENTRY_KIND, serde_json::json!({ "taskId": "old" })) + .unwrap(); + store + .append_custom(RESULT_ENTRY_KIND, serde_json::json!({ "toolUseId": "tu" })) + .unwrap(); + drop(store); + let key = journal_key_beside(&path); + + let (mut store, _) = SessionStore::open(path.clone()).unwrap(); + assert!(!key.exists(), "opening writes no key"); + assert_eq!(store.active_journal(TASK_ENTRY_KIND).len(), 1); + assert_eq!(store.active_journal(RESULT_ENTRY_KIND).len(), 1); + store + .append_journal(TASK_ENTRY_KIND, serde_json::json!({ "taskId": "new" })) + .unwrap(); + assert!(key.is_file()); + store + .append_custom(TASK_ENTRY_KIND, serde_json::json!({ "taskId": "planted" })) + .unwrap(); + drop(store); + let (store, _) = SessionStore::open(path).unwrap(); + let tasks: Vec<_> = store + .active_journal(TASK_ENTRY_KIND) + .into_iter() + .map(|t| t["taskId"].as_str().unwrap().to_owned()) + .collect(); + assert_eq!(tasks, ["old", "new"]); + assert_eq!(store.active_journal(RESULT_ENTRY_KIND).len(), 1); + } + + /// Deleting a session moves its journal key out of the session directory into `.trash` with + /// it; restoring brings it back, and the journal still verifies. + #[test] + fn a_trashed_session_takes_its_journal_key_along_and_brings_it_back() { + use crate::mcp_resume::TASK_ENTRY_KIND; + let dir = tmpdir(); + let repo = SessionRepo::open(dir.path()).unwrap(); + let mut store = repo.create(SessionMeta::new("/w", "m")).unwrap(); + store + .append_journal(TASK_ENTRY_KIND, serde_json::json!({ "taskId": "t" })) + .unwrap(); + let id = store.meta().id.clone(); + let path = store.path().to_path_buf(); + drop(store); + let key = journal_key_beside(&path); + assert!(key.is_file()); + + repo.delete(&id).unwrap(); + assert!(!key.exists(), "the key left the session directory"); + assert!( + dir.path() + .join(".trash") + .join(key.file_name().unwrap()) + .is_file(), + "the key is in the trash beside its session" + ); + + assert!(repo.restore_session(&id).unwrap()); + assert!(key.is_file(), "the key came back with its session"); + let (store, _) = SessionStore::open(path).unwrap(); + assert_eq!(store.active_journal(TASK_ENTRY_KIND).len(), 1); + } + + /// Sessions whose names differ only after their last dot each get their own key: the key's + /// name is a suffix on the whole file name, not a replaced extension. + #[test] + fn sessions_named_alike_keep_separate_journal_keys() { + use crate::mcp_resume::TASK_ENTRY_KIND; + let dir = tmpdir(); + let mut keys = Vec::new(); + for name in ["work.1", "work.2"] { + let path = dir.path().join(name); + let mut store = + SessionStore::create(path.clone(), SessionMeta::new("/w", "m")).unwrap(); + store + .append_journal(TASK_ENTRY_KIND, serde_json::json!({ "taskId": name })) + .unwrap(); + keys.push(fs::read(journal_key_beside(&path)).unwrap()); + } + assert_ne!(keys[0], keys[1], "one key file shared by two sessions"); + for name in ["work.1", "work.2"] { + let (store, _) = SessionStore::open(dir.path().join(name)).unwrap(); + let tasks = store.active_journal(TASK_ENTRY_KIND); + assert_eq!(tasks.len(), 1, "{name}: {tasks:?}"); + assert_eq!(tasks[0]["taskId"], name); + } + } + + /// A key under the old `with_extension` name still verifies the entries it sealed, and the next + /// journal write carries it over to the suffixed name rather than making a new one. + #[test] + fn a_journal_key_under_the_old_name_is_carried_over() { + use crate::mcp_resume::{JournalAuth, TASK_ENTRY_KIND}; + let dir = tmpdir(); + let path = dir.path().join("s.jsonl"); + let mut store = SessionStore::create(path.clone(), SessionMeta::new("/w", "m")).unwrap(); + store + .append_journal(TASK_ENTRY_KIND, serde_json::json!({ "taskId": "old" })) + .unwrap(); + drop(store); + let legacy = path.with_extension(JournalAuth::SIDECAR); + fs::rename(journal_key_beside(&path), &legacy).unwrap(); + + let (mut store, _) = SessionStore::open(path.clone()).unwrap(); + assert_eq!(store.active_journal(TASK_ENTRY_KIND).len(), 1); + store + .append_custom(TASK_ENTRY_KIND, serde_json::json!({ "taskId": "planted" })) + .unwrap(); + store + .append_journal(TASK_ENTRY_KIND, serde_json::json!({ "taskId": "new" })) + .unwrap(); + assert_eq!( + fs::read(journal_key_beside(&path)).unwrap(), + fs::read(&legacy).unwrap(), + "the old key is carried over, not replaced" + ); + drop(store); + let (store, _) = SessionStore::open(path).unwrap(); + let tasks: Vec<_> = store + .active_journal(TASK_ENTRY_KIND) + .into_iter() + .map(|t| t["taskId"].as_str().unwrap().to_owned()) + .collect(); + assert_eq!(tasks, ["old", "new"]); + } + + /// Set by [`concurrent_first_journal_writers_all_use_one_key`] for the child processes it + /// runs this test binary as. + const WRITER_SESSION: &str = "JOURNAL_KEY_WRITER_SESSION"; + + /// One of [`concurrent_first_journal_writers_all_use_one_key`]'s writer processes: waits for + /// the go file, opens the session and journals one entry. A no-op when run on its own. + #[test] + #[ignore = "a child process of concurrent_first_journal_writers_all_use_one_key"] + fn journal_key_writer_child() { + let Ok(session) = std::env::var(WRITER_SESSION) else { + return; + }; + let session = PathBuf::from(session); + let id = std::env::var("JOURNAL_KEY_WRITER_ID").unwrap(); + journal_one_after_go(&session, &id); + } + + fn journal_one_after_go(session: &Path, id: &str) { + let go = session.with_extension("go"); + fs::write(session.with_extension(format!("ready.{id}")), b"").unwrap(); + let deadline = std::time::Instant::now() + std::time::Duration::from_secs(30); + while !go.exists() { + assert!(std::time::Instant::now() < deadline, "no go"); + std::thread::sleep(std::time::Duration::from_millis(1)); + } + let (mut store, _) = SessionStore::open(session.to_path_buf()).unwrap(); + store + .append_journal( + crate::mcp_resume::TASK_ENTRY_KIND, + serde_json::json!({ "taskId": id }), + ) + .unwrap(); + } + + /// Every first journal writer, threads of this process and other processes alike, ends up + /// using the one key on disk, so every entry any of them journals verifies against it. A writer + /// that believed in a key that lost the race would have its entries refused on replay. + #[test] + fn concurrent_first_journal_writers_all_use_one_key() { + use crate::mcp_resume::{JournalAuth, TASK_ENTRY_KIND}; + const PROCESSES: usize = 4; + const THREADS: usize = 8; + let dir = tmpdir(); + let path = dir.path().join("s.jsonl"); + drop(SessionStore::create(path.clone(), SessionMeta::new("/w", "m")).unwrap()); + + let exe = std::env::current_exe().unwrap(); + let children: Vec<_> = (0..PROCESSES) + .map(|i| { + std::process::Command::new(&exe) + .args([ + "session_store::tests::journal_key_writer_child", + "--exact", + "--ignored", + "--test-threads=1", + ]) + .env(WRITER_SESSION, &path) + .env("JOURNAL_KEY_WRITER_ID", format!("process-{i}")) + .stdout(std::process::Stdio::null()) + .stderr(std::process::Stdio::null()) + .spawn() + .unwrap() + }) + .collect(); + let threads: Vec<_> = (0..THREADS) + .map(|i| { + let path = path.clone(); + std::thread::spawn(move || journal_one_after_go(&path, &format!("thread-{i}"))) + }) + .collect(); + let deadline = std::time::Instant::now() + std::time::Duration::from_secs(60); + while fs::read_dir(dir.path()) + .unwrap() + .flatten() + .filter(|e| e.file_name().to_string_lossy().contains(".ready.")) + .count() + < PROCESSES + THREADS + { + assert!( + std::time::Instant::now() < deadline, + "writers never got ready" + ); + std::thread::sleep(std::time::Duration::from_millis(5)); + } + fs::write(path.with_extension("go"), b"").unwrap(); + for thread in threads { + thread.join().unwrap(); + } + for mut child in children { + assert!(child.wait().unwrap().success(), "a writer process failed"); + } + + let auth = JournalAuth::from_sidecar(&fs::read(journal_key_beside(&path)).unwrap()) + .expect("one key on disk"); + let entries: Vec = fs::read_to_string(&path) + .unwrap() + .lines() + .filter_map(|l| serde_json::from_str::(l).ok()) + .filter(|v| v["type"] == "custom" && v["kind"] == TASK_ENTRY_KIND) + .map(|v| v["data"].clone()) + .collect(); + assert_eq!(entries.len(), PROCESSES + THREADS); + for entry in &entries { + assert!( + entry.get("mac").is_some() && auth.accepts(TASK_ENTRY_KIND, entry), + "sealed under a key that is not the one on disk: {entry}" + ); + } + } + + /// A session file that exists with a looser mode (an older version's, a restore's, a + /// `chmod`) is tightened to `0600` by the next append: the transcript, and the MCP task ids it + /// journals (possible bearer tokens), are never written into a group/world-readable file. + #[test] + #[cfg(unix)] + fn an_append_tightens_a_loosened_session_file_to_private() { + use std::os::unix::fs::PermissionsExt; + + let dir = tmpdir(); + let path = dir.path().join("s.jsonl"); + let mut store = SessionStore::create(path.clone(), SessionMeta::new("/w", "m")).unwrap(); + fs::set_permissions(&path, fs::Permissions::from_mode(0o644)).unwrap(); + store + .append_custom( + "mcp_task", + serde_json::json!({ "taskId": "SECRET-TASK-ID", "sessionId": "s" }), + ) + .unwrap(); + assert_eq!( + fs::metadata(&path).unwrap().permissions().mode() & 0o777, + 0o600, + "the append must tighten the file before writing a task id into it" + ); + } + #[test] fn create_initializes_an_existing_empty_file_in_place_instead_of_failing() { // Track L8: a zero-byte file at the target path (e.g. `touch`'d ahead of time, or left over diff --git a/crates/agent/src/tools/mcp.rs b/crates/agent/src/tools/mcp.rs index 8f9bdf81..a68de292 100644 --- a/crates/agent/src/tools/mcp.rs +++ b/crates/agent/src/tools/mcp.rs @@ -1269,7 +1269,9 @@ async fn drive_tool_call( cancel_on_drop: true, fresh: true, }; - return await_task(conn, client, host, record, seed, progress).await; + return await_task(conn, client, host, record, seed, progress) + .await + .map_err(ToolError::from); } other => { return Err(ToolError::Execution(format!( @@ -1485,6 +1487,48 @@ impl Drop for TaskCancelOnDrop { } } +/// Why polling a task stopped short of a result. +enum TaskFailure { + /// The connection could not be recovered: the task may well still be alive on the server. + ConnectionLost(String), + /// A definitive outcome: a terminal status, a JSON-RPC error, an expired TTL. + Final(ToolError), +} + +impl From for TaskFailure { + fn from(e: ToolError) -> Self { + Self::Final(e) + } +} + +impl From for ToolError { + fn from(f: TaskFailure) -> Self { + match f { + TaskFailure::ConnectionLost(why) => ToolError::Execution(why), + TaskFailure::Final(e) => e, + } + } +} + +/// Why a journaled task could not be resolved now (see [`McpCatalog::resume_task`]). +#[derive(Debug)] +pub enum ResumeError { + /// The server could not be reached. Not an answer: the task stays pending and journaled, and + /// is tried again on a later prompt or start. + Unreachable(String), + /// The call's definitive answer (an error result the model is told). + Final(ToolError), +} + +impl std::fmt::Display for ResumeError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Unreachable(why) => f.write_str(why), + Self::Final(e) => write!(f, "{e}"), + } + } +} + /// How polling starts: from a fresh `CreateTaskResult`, or (resume) from a journaled record. struct TaskSeed { poll_interval_ms: Option, @@ -1583,7 +1627,7 @@ async fn await_task( mut record: McpTaskRecord, seed: TaskSeed, progress: Option<&ToolProgress>, -) -> Result { +) -> Result { let task_id = record.task_id.clone(); let server_name = record.server.clone(); let remote_name = record.tool.clone(); @@ -1604,7 +1648,7 @@ async fn await_task( let mut pending_update: Option<(Vec, InputResponses)> = None; let mut losses: u32 = 0; let lost = |losses: u32, e: &ServiceError| { - ToolError::Execution(format!( + TaskFailure::ConnectionLost(format!( "mcp task `{task_id}` on `{server_name}`/`{remote_name}`: connection lost {losses} times in a row, giving up: {e}" )) }; @@ -1642,7 +1686,7 @@ async fn await_task( client = recover(conn, client, loss, losses).await; continue; } - None => return Err(tool_call_err(server_name, remote_name, e)), + None => return Err(tool_call_err(server_name, remote_name, e).into()), }, } } @@ -1657,7 +1701,8 @@ async fn await_task( return Err(ToolError::Execution(format!( "mcp task `{task_id}` on `{server_name}`/`{remote_name}` did not finish within its ttlMs ({} ms)", ttl.as_millis() - ))); + )) + .into()); } wait = wait.min(ttl - elapsed); } @@ -1681,7 +1726,7 @@ async fn await_task( client = recover(conn, client, loss, losses).await; continue; } - None => return Err(tool_call_err(server_name, remote_name, e)), + None => return Err(tool_call_err(server_name, remote_name, e).into()), }, }; let detailed = info.task; @@ -1748,18 +1793,21 @@ async fn await_task( return Err(ToolError::Execution(format!( "mcp task `{task_id}` on `{server_name}`/`{remote_name}` failed: {}", Value::Object(error) - ))); + )) + .into()); } TaskPayload::Cancelled => { cancel.finish(); return Err(ToolError::Execution(format!( "mcp task `{task_id}` on `{server_name}`/`{remote_name}` was cancelled" - ))); + )) + .into()); } other => { return Err(ToolError::Execution(format!( "mcp task `{task_id}` on `{server_name}`/`{remote_name}` returned unknown payload: {other:?}" - ))); + )) + .into()); } } } @@ -1767,16 +1815,20 @@ async fn await_task( impl McpCatalog { /// Resume polling a journaled task after a restart (see [`McpTaskRecord`]) through this - /// catalog's live (or redialed) connection to `record.server`, to its terminal result. Never - /// cancels the task when dropped. A server that no longer knows the task answers `-32602`, - /// and that is the call's answer. `host` answers any in-task input; `progress` receives the - /// task's status and record updates. + /// catalog's connection to `record.server` (live, redialed, or one that could not be dialed at + /// startup), to its terminal result. Never cancels the task when dropped. + /// + /// Only a definitive outcome answers the call: the task's terminal status, a JSON-RPC error + /// such as `-32602` (the server no longer knows it), or its `ttlMs` running out. A server that + /// cannot be reached, after a bounded few redials, is [`ResumeError::Unreachable`]: the task + /// may be alive, so it stays pending and is tried again later. `host` answers any in-task + /// input; `progress` receives the task's status and record updates. pub async fn resume_task( &self, record: McpTaskRecord, host: &McpHost, progress: Option<&ToolProgress>, - ) -> Result { + ) -> Result { let server = record.server.clone(); let tool = record.tool.clone(); // A TTL that ran out while the process was down needs no server (nor its configuration) @@ -1784,20 +1836,15 @@ impl McpCatalog { if let Some(ttl) = record.ttl_ms && unix_ms().saturating_sub(record.created_at_ms) >= ttl { - return Err(ToolError::Execution(format!( + return Err(ResumeError::Final(ToolError::Execution(format!( "mcp task `{}` on `{server}`/`{tool}` did not finish within its ttlMs ({ttl} ms)", record.task_id - ))); + )))); } - let entry = self - .snapshot() - .into_iter() - .find(|s| s.name == server) - .ok_or_else(|| { - ToolError::Execution(format!("mcp server `{server}` is not configured any more")) - })?; - let conn = entry.conn.upgrade().ok_or_else(|| { - ToolError::Execution(format!("mcp server `{server}` is no longer connected")) + let conn = self.connection(&server).ok_or_else(|| { + ResumeError::Final(ToolError::Execution(format!( + "mcp server `{server}` is not configured any more" + ))) })?; let seed = TaskSeed { // Unknown until the first poll answers; ask promptly, then follow the server. @@ -1806,14 +1853,52 @@ impl McpCatalog { cancel_on_drop: false, fresh: false, }; - let client = conn.client().await.map_err(|e| { - ToolError::Execution(format!("mcp server `{server}` is not reachable: {e}")) + let client = dial_with_backoff(&conn).await.map_err(|e| { + ResumeError::Unreachable(format!("mcp server `{server}` is not reachable: {e}")) })?; - let result = await_task(&conn, client, host, record, seed, progress).await?; - tool_output_from_result(&server, &tool, result) + match await_task(&conn, client, host, record, seed, progress).await { + Ok(result) => { + tool_output_from_result(&server, &tool, result).map_err(ResumeError::Final) + } + Err(TaskFailure::ConnectionLost(why)) => Err(ResumeError::Unreachable(why)), + Err(TaskFailure::Final(e)) => Err(ResumeError::Final(e)), + } + } + + /// The connection to `server`: a connected one's, or the dormant one kept for a configured + /// server that could not be dialed at startup (so a resume can still reach it once it is back). + fn connection(&self, server: &str) -> Option> { + self.snapshot() + .into_iter() + .find(|s| s.name == server) + .and_then(|s| s.conn.upgrade()) + .or_else(|| { + self.unconnected + .lock() + .ok() + .and_then(|u| u.get(server).cloned()) + }) } } +/// Dial a connection, retrying a few times with backoff (250 ms doubling, about 5 s in all): a +/// server that is restarting is not gone. +async fn dial_with_backoff(conn: &McpConnection) -> Result, String> { + let mut delay = Duration::from_millis(250); + let mut last = String::new(); + for attempt in 0..5 { + if attempt > 0 { + tokio::time::sleep(delay).await; + delay = (delay * 2).min(Duration::from_secs(2)); + } + match conn.client().await { + Ok(client) => return Ok(client), + Err(e) => last = e, + } + } + Err(last) +} + pub(crate) fn tool_output_from_result( server_name: &str, remote_name: &str, @@ -2048,6 +2133,7 @@ async fn connect_many( let mut tools: Vec> = Vec::new(); let mut catalogs = Vec::new(); let mut warnings = Vec::new(); + let mut unconnected = HashMap::new(); for (name, result) in connect_each(jobs, idle_reap_after, manifest_dir).await { match result { Ok((server_tools, catalog)) => { @@ -2061,10 +2147,24 @@ async fn connect_many( "failed to connect to MCP server; its tools will not be available" ); warnings.push(format!("mcp server `{name}`: {e}")); + if let Some((config, dial)) = jobs.iter().find(|(c, _)| c.name == name) { + unconnected.insert( + name.clone(), + Arc::new(McpConnection::dormant( + config.clone(), + dial.clone(), + idle_reap_after, + )), + ); + } } } } - (tools, McpCatalog::new(catalogs), warnings) + let catalog = McpCatalog::new(catalogs); + if let Ok(mut u) = catalog.unconnected.lock() { + *u = unconnected; + } + (tools, catalog, warnings) } /// Dial every job concurrently, keeping each server's outcome separate. @@ -2728,6 +2828,10 @@ pub struct McpPromptInfo { #[derive(Clone, Default)] pub struct McpCatalog { servers: Arc>>, + /// Configured servers that could not be dialed at startup, kept dormant: they offer no tools, + /// but a journaled task on one can still be resumed once it is back (see + /// [`Self::resume_task`]). + unconnected: Arc>>>, /// The session this view of the catalog belongs to (see [`Self::for_session`]): the host its /// non-tool requests (`events/*`, `completion/complete`) are registered under, so a nested /// request raised during one is attributed to that session or refused, never guessed. @@ -2738,14 +2842,25 @@ impl McpCatalog { pub fn new(servers: Vec) -> Self { Self { servers: Arc::new(std::sync::Mutex::new(servers)), + unconnected: Arc::default(), session_host: None, } } + /// Whether `server` is one of this catalog's configured servers (connected or not). + pub fn has_server(&self, server: &str) -> bool { + self.snapshot().iter().any(|s| s.name == server) + || self + .unconnected + .lock() + .is_ok_and(|u| u.contains_key(server)) + } + /// This catalog (the same servers, shared) as session `host` uses it. pub fn for_session(&self, host: Arc) -> Self { Self { servers: self.servers.clone(), + unconnected: self.unconnected.clone(), session_host: Some(host), } } diff --git a/crates/agent/tests/mcp_tasks.rs b/crates/agent/tests/mcp_tasks.rs index 869c9d3c..dfefe567 100644 --- a/crates/agent/tests/mcp_tasks.rs +++ b/crates/agent/tests/mcp_tasks.rs @@ -416,8 +416,12 @@ fn serve_abort_sends_tasks_cancel_for_sticky_task() { let cancel_deadline = Instant::now() + Duration::from_secs(3); let mut cancelled_id = None; while Instant::now() < cancel_deadline { - if cancel_flag.exists() { - cancelled_id = Some(std::fs::read_to_string(&cancel_flag).unwrap()); + // The fixture renames the flag into place whole; reading until it has content also + // covers a fixture that ever wrote it in place. + if let Ok(id) = std::fs::read_to_string(&cancel_flag) + && !id.is_empty() + { + cancelled_id = Some(id); break; } std::thread::sleep(Duration::from_millis(50)); diff --git a/crates/agent/tests/mcp_tasks_env/mod.rs b/crates/agent/tests/mcp_tasks_env/mod.rs index dee58258..b6f74886 100644 --- a/crates/agent/tests/mcp_tasks_env/mod.rs +++ b/crates/agent/tests/mcp_tasks_env/mod.rs @@ -64,7 +64,9 @@ impl Env { self.settings(json!([self.stdio_with(json!({}))])); } - /// Start the fixture as a standalone HTTP server; returns it and its settings entry. + /// Start the fixture as a standalone HTTP server, on a port it picks; returns it and its settings + /// entry. Gated tasks survive a restart (`MCP_TASKS_FIXTURE_KEEP`); a restarted server keeps its + /// URL behind a [`Relay`]. pub fn http_server(&self) -> (ChildGuard, Value) { let port_file = self.dir.path().join("port"); let _ = std::fs::remove_file(&port_file); @@ -72,6 +74,7 @@ impl Env { .env("MCP_TASKS_FIXTURE_HTTP_PORT_FILE", &port_file) .env("MCP_TASKS_FIXTURE_LOG", &self.log) .env("MCP_TASKS_FIXTURE_GATE", &self.gate) + .env("MCP_TASKS_FIXTURE_KEEP", self.dir.path().join("kept")) .stdin(Stdio::null()) .stdout(Stdio::null()) .stderr(Stdio::null()) @@ -247,3 +250,58 @@ pub fn tool_result_sent(request: &str, tool_use_id: &str) -> Value { .find(|b| b["type"] == "tool_result" && b["tool_use_id"] == tool_use_id) .unwrap_or_else(|| panic!("no tool_result for {tool_use_id} in {body}")) } + +/// A URL that outlives the HTTP fixture behind it: a loopback listener the test binds and holds for +/// its whole life, relaying each connection to whichever fixture is up, and closing it at once while +/// none is. A server "restarted on the same URL" without releasing a port for anything else to take. +pub struct Relay { + pub port: u16, + backend: std::sync::Arc>>, +} + +impl Relay { + pub fn start() -> Self { + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let port = listener.local_addr().unwrap().port(); + let backend = std::sync::Arc::new(std::sync::Mutex::new(None::)); + let routes = backend.clone(); + std::thread::spawn(move || { + for client in listener.incoming() { + let Ok(client) = client else { continue }; + let target = *routes.lock().unwrap(); + let Some(server) = + target.and_then(|p| std::net::TcpStream::connect(("127.0.0.1", p)).ok()) + else { + continue; // nothing up: the connection closes as it drops + }; + for (mut from, mut to) in [ + (client.try_clone().unwrap(), server.try_clone().unwrap()), + (server, client), + ] { + std::thread::spawn(move || { + let _ = std::io::copy(&mut from, &mut to); + let _ = to.shutdown(std::net::Shutdown::Write); + }); + } + } + }); + Self { port, backend } + } + + /// Relay to the fixture `entry` (from [`Env::http_server`]) from now on, or to nothing. + pub fn route_to(&self, entry: Option<&Value>) { + let port = entry.map(|e| { + let url = e["url"].as_str().unwrap(); + let port = url.rsplit(':').next().unwrap(); + port.trim_end_matches("/mcp").parse::().unwrap() + }); + *self.backend.lock().unwrap() = port; + } + + /// `entry` with its URL pointed at the relay. + pub fn entry_for(&self, entry: &Value) -> Value { + let mut entry = entry.clone(); + entry["url"] = json!(format!("http://127.0.0.1:{}/mcp", self.port)); + entry + } +} diff --git a/crates/agent/tests/mcp_tasks_resume.rs b/crates/agent/tests/mcp_tasks_resume.rs index a8fd4a0b..d2360a95 100644 --- a/crates/agent/tests/mcp_tasks_resume.rs +++ b/crates/agent/tests/mcp_tasks_resume.rs @@ -7,6 +7,7 @@ mod common; mod mcp_tasks_env; +use std::io::Write as _; use std::process::ChildStdin; use std::sync::{Arc, Mutex}; use std::time::Duration; @@ -637,3 +638,581 @@ fn a_clone_does_not_resume_its_parents_task() { ); s.close(); } + +/// Item 1: only the session that journaled a task may resume it. No fork path copies the journal +/// (forks copy messages only), so the filter is exercised the way a journal really does end up +/// under another session id: the session file copied (a restore, a migration) and given a new id. +/// The copy carries the unanswered call and the task record; opening it must not resume (or wait +/// on) the original session's task. +#[test] +fn a_journal_carried_under_another_session_id_is_not_resumed() { + let env = Env::new(); + let _server = env.http(); + let task = start_and_kill(&env, "mcp__t__gated_task", None, ""); + let polls = gets_for(&env, &task); + + let copy = env.dir.path().join("copy.jsonl"); + let original = std::fs::read_to_string(env.session_file()).unwrap(); + let mut lines = original.lines(); + let mut header: Value = serde_json::from_str(lines.next().unwrap()).unwrap(); + header["id"] = json!("restored-under-a-new-id"); + let mut text = header.to_string(); + for line in lines { + text.push('\n'); + text.push_str(line); + } + text.push('\n'); + std::fs::write(©, text).unwrap(); + assert!( + std::fs::read_to_string(©).unwrap().contains(&task), + "the copy carries the task record" + ); + + let (base, bodies) = spawn_model_server(vec![turn_text("carrying on")]); + let mut cmd = common::serve_cmd(common::BIN, &base, copy.to_str().unwrap()); + cmd.env("HOME", &env.home) + .env("BEYOND_AI_AGENT_MCP_IDLE_SECS", "0"); + let mut child = common::SpawnGuarded::spawn_guarded(&mut cmd); + let mut stdin = child.stdin.take().unwrap(); + let mut stdout = common::child_frames(&mut child); + prompt(&mut stdin, "what happened?"); + let frames = read_until_or_fail( + &mut stdout, + "the copy's prompt", + Duration::from_secs(20), + |f| is_event(f, "tool_progress", "toolu_g"), + |f| f["type"] == "response" && f["command"] == "prompt", + ); + assert_eq!(frames.last().unwrap()["success"], true, "{frames:?}"); + drop(stdin); + child.wait().unwrap(); + let requests = bodies.lock().unwrap().clone(); + let sent = tool_result_sent(requests.last().unwrap(), "toolu_g"); + assert!(!sent.to_string().contains("gated-done"), "{sent}"); + assert!( + gets_for(&env, &task) <= polls, + "the copy never polls the original's task" + ); +} + +/// Every `.jsonl` file under `dir`, recursively, with its text. +fn jsonl_files(dir: &std::path::Path) -> Vec { + session_files(dir) + .into_iter() + .map(|(_, text)| text) + .collect() +} + +fn run_once(env: &Env, base: &str, args: &[&str]) -> ChildGuard { + let mut cmd = common::run_cmd(common::BIN); + cmd.arg("run") + .args(args) + .args([ + "--gateway-url", + base, + "--key", + "bai_v1.test", + "--model", + "claude-test", + ]) + .env("HOME", &env.home) + .env("BEYOND_AI_AGENT_MCP_IDLE_SECS", "0") + .current_dir(env.dir.path()) + .stdin(std::process::Stdio::null()) + .stdout(std::process::Stdio::piped()) + .stderr(std::process::Stdio::piped()); + common::SpawnGuarded::spawn_guarded(&mut cmd) +} + +/// Item 3: `run` journals the tasks it starts, and `run --continue` resumes the ones a killed run +/// left in flight, answering the call with the real result before its own turn (not the generic +/// interrupted placeholder), with roles alternating. +#[test] +fn run_continue_resumes_a_task_a_killed_run_left_in_flight() { + let env = Env::new(); + let _server = env.http(); + let (base, _bodies) = + spawn_model_server(vec![turn_tool_use("toolu_g", "mcp__t__gated_task", "{}")]); + let mut first = run_once(&env, &base, &["start the job"]); + wait_for("the task being polled", || { + !env.methods("tasks/get").is_empty() + }); + let task = env.methods("tasks/get")[0]["params"]["taskId"] + .as_str() + .unwrap() + .to_owned(); + wait_for("run's journal entry", || { + jsonl_files(&env.home) + .iter() + .any(|text| text.contains(&task)) + }); + first.kill().unwrap(); + first.wait().unwrap(); + std::fs::write(&env.gate, b"open").unwrap(); + + let (base, bodies) = spawn_model_server(vec![turn_text("carried on"), turn_text("t")]); + let second = run_once(&env, &base, &["--continue", "what happened?"]); + let output = second.wait_with_output().unwrap(); + assert!(output.status.success(), "{output:?}"); + let requests: Vec = bodies + .lock() + .unwrap() + .iter() + .filter(|r| !r.contains("You write short titles")) + .cloned() + .collect(); + assert_eq!(requests.len(), 1, "{requests:#?}"); + assert_alternates(&requests[0]); + let sent = tool_result_sent(&requests[0], "toolu_g"); + assert_eq!(sent["is_error"], json!(false), "{sent}"); + assert!(sent.to_string().contains("gated-done"), "{sent}"); + assert_eq!( + env.methods("tools/call").len(), + 1, + "resumed, not re-invoked" + ); +} + +/// A client cannot forge the host's task journal. `append_custom` with an `mcp_task` (or +/// `mcp_task_result`) kind is refused; had it landed, the resumer would poll a task id of the +/// client's choosing on the next start (the latest record for a call wins). +#[test] +fn a_forged_task_record_from_a_client_is_refused_and_never_polled() { + let env = Env::new(); + let _server = env.http(); + start_and_kill(&env, "mcp__t__gated_task", None, ""); + + let mut s = serve(&env, vec![turn_text("unused")]); + send(&mut s.stdin, json!({ "type": "get_state", "id": "st" })); + let state = read_until(&mut s.stdout, "state", |f| { + f["type"] == "response" && f["command"] == "get_state" + }); + let session_id = state.last().unwrap()["data"]["session_id"] + .as_str() + .unwrap() + .to_owned(); + for kind in ["mcp_task", "mcp_task_result"] { + send( + &mut s.stdin, + json!({ + "type": "append_custom", + "id": kind, + "kind": kind, + "data": { + "server": "t", "tool": "gated_task", "taskId": "FORGED-TASK", + "createdAtMs": 0, "toolUseId": "toolu_g", "sessionId": session_id, + "name": "mcp__t__gated_task", "content": "FORGED-RESULT", "isError": false, + }, + }), + ); + let resp = read_until(&mut s.stdout, "append_custom", |f| { + f["type"] == "response" && f["command"] == "append_custom" + }); + let resp = resp.last().unwrap(); + assert_eq!(resp["success"], false, "{kind}: {resp}"); + assert!( + resp["error"].as_str().unwrap().contains("reserved"), + "{resp}" + ); + } + s.close(); + assert!( + !env.session_text().contains("FORGED"), + "nothing forged reached the file" + ); + + // And a restart polls only the real task. + let mut again = serve(&env, vec![turn_text("unused")]); + read_until(&mut again.stdout, "the real task's resume", |f| { + is_event(f, "tool_progress", "toolu_g") + }); + again.close(); + assert!( + env.methods("tasks/get") + .iter() + .all(|g| g["params"]["taskId"] != "FORGED-TASK"), + "the forged task id is never polled" + ); +} + +/// The MCP server behind `relay` (so a restarted one keeps its URL), configured as the only server. +fn http_on(env: &Env, relay: &mcp_tasks_env::Relay) -> ChildGuard { + let (child, entry) = env.http_server(); + relay.route_to(Some(&entry)); + env.settings(json!([relay.entry_for(&entry)])); + child +} + +const PENDING: &str = "[MCP task result pending]"; + +/// A server that is down is not an answer. `run --continue` while it is down tells the turn the +/// result is pending and journals nothing; once the server is back (its task still alive), the +/// next `run --continue` delivers the real result. +#[test] +fn run_continue_while_the_server_is_down_leaves_the_task_pending_then_delivers_it() { + let env = Env::new(); + let relay = mcp_tasks_env::Relay::start(); + let server = http_on(&env, &relay); + let (base, _bodies) = + spawn_model_server(vec![turn_tool_use("toolu_g", "mcp__t__gated_task", "{}")]); + let mut first = run_once(&env, &base, &["start the job"]); + wait_for("the task being polled", || { + !env.methods("tasks/get").is_empty() + }); + let task = env.methods("tasks/get")[0]["params"]["taskId"] + .as_str() + .unwrap() + .to_owned(); + wait_for("run's journal entry", || { + jsonl_files(&env.home) + .iter() + .any(|text| text.contains(&task)) + }); + first.kill().unwrap(); + first.wait().unwrap(); + drop(server); + relay.route_to(None); + + // Down: the turn is told the result is pending; nothing is journaled as the answer. + let (base, bodies) = spawn_model_server(vec![turn_text("noted"), turn_text("t")]); + let output = run_once(&env, &base, &["--continue", "what happened?"]) + .wait_with_output() + .unwrap(); + assert!(output.status.success(), "{output:?}"); + let sent = tool_result_sent(&bodies.lock().unwrap()[0], "toolu_g"); + assert!(sent.to_string().contains(PENDING), "{sent}"); + assert!( + !jsonl_files(&env.home) + .iter() + .any(|text| text.contains("mcp_task_result")), + "an unreachable server is not a result" + ); + + // Back, the task still alive: the real result reaches the model. + let _server = http_on(&env, &relay); + std::fs::write(&env.gate, b"open").unwrap(); + let (base, bodies) = spawn_model_server(vec![turn_text("carried on"), turn_text("t")]); + let output = run_once(&env, &base, &["--continue", "and now?"]) + .wait_with_output() + .unwrap(); + assert!(output.status.success(), "{output:?}"); + let request = bodies.lock().unwrap()[0].clone(); + assert_alternates(&request); + let sent = tool_result_sent(&request, "toolu_g"); + assert!(sent.to_string().contains("gated-done"), "{sent}"); + assert!( + !request.contains(PENDING), + "the placeholder is replaced: {request}" + ); +} + +/// The same in `serve`: with the server down at start, the prompt is told the result is pending +/// (nothing journaled), and a later prompt, with the server back on its URL, gets the real result. +#[test] +fn serve_with_the_server_down_leaves_the_task_pending_then_delivers_it() { + let env = Env::new(); + let relay = mcp_tasks_env::Relay::start(); + let server = http_on(&env, &relay); + start_and_kill(&env, "mcp__t__gated_task", None, ""); + drop(server); + relay.route_to(None); + + let mut s = serve( + &env, + vec![turn_text("noted"), turn_text("carried on"), turn_text("t")], + ); + prompt(&mut s.stdin, "what happened?"); + let frames = read_until_prompt_done(&mut s.stdout); + let end = tool_end(&frames, "toolu_g"); + assert!(end["result"].as_str().unwrap().contains(PENDING), "{end}"); + assert!( + tool_result_sent(&s.requests()[0], "toolu_g") + .to_string() + .contains(PENDING) + ); + assert!(!env.session_text().contains("mcp_task_result")); + + let _server = http_on(&env, &relay); + std::fs::write(&env.gate, b"open").unwrap(); + prompt(&mut s.stdin, "and now?"); + let frames = read_until_prompt_done(&mut s.stdout); + let end = tool_end(&frames, "toolu_g"); + assert!( + end["result"].as_str().unwrap().contains("gated-done"), + "{end}" + ); + let requests = s.requests(); + let last = requests.last().unwrap(); + assert_alternates(last); + assert!( + tool_result_sent(last, "toolu_g") + .to_string() + .contains("gated-done") + ); + s.close(); + assert!(env.session_text().contains("mcp_task_result")); +} + +/// A journaled task on a server that is no longer configured is never resumed (nor answered with a +/// made-up error): the call keeps the generic interrupted repair, nothing is polled or journaled. +#[test] +fn a_task_on_a_server_no_longer_configured_is_not_resumed() { + let env = Env::new(); + let (_server, entry) = env.http_server(); + env.settings(json!([entry.clone()])); + start_and_kill(&env, "mcp__t__gated_task", None, ""); + let mut renamed = entry; + renamed["name"] = json!("other"); + env.settings(json!([renamed])); + + let mut s = serve(&env, vec![turn_text("carrying on")]); + prompt(&mut s.stdin, "what happened?"); + let frames = read_until_or_fail( + &mut s.stdout, + "the prompt", + Duration::from_secs(20), + |f| is_event(f, "tool_progress", "toolu_g"), + |f| f["type"] == "response" && f["command"] == "prompt", + ); + assert_eq!(frames.last().unwrap()["success"], true, "{frames:?}"); + s.close(); + assert!(!env.session_text().contains("mcp_task_result")); +} + +/// `run --continue` applies the same rule: a task on a server no longer configured is not +/// resumed (no attempt, no made-up "not configured" answer journaled). +#[test] +fn run_continue_does_not_resume_a_task_on_a_server_no_longer_configured() { + let env = Env::new(); + let (_server, entry) = env.http_server(); + env.settings(json!([entry.clone()])); + let (base, _bodies) = + spawn_model_server(vec![turn_tool_use("toolu_g", "mcp__t__gated_task", "{}")]); + let mut first = run_once(&env, &base, &["start the job"]); + wait_for("the task being polled", || { + !env.methods("tasks/get").is_empty() + }); + let task = env.methods("tasks/get")[0]["params"]["taskId"] + .as_str() + .unwrap() + .to_owned(); + wait_for("run's journal entry", || { + jsonl_files(&env.home) + .iter() + .any(|text| text.contains(&task)) + }); + first.kill().unwrap(); + first.wait().unwrap(); + let mut renamed = entry; + renamed["name"] = json!("other"); + env.settings(json!([renamed])); + + let (base, _bodies) = spawn_model_server(vec![turn_text("carried on"), turn_text("t")]); + let output = run_once(&env, &base, &["--continue", "what happened?"]) + .wait_with_output() + .unwrap(); + assert!(output.status.success(), "{output:?}"); + assert!( + !String::from_utf8_lossy(&output.stderr).contains("resuming MCP task"), + "{output:?}" + ); + assert!( + !jsonl_files(&env.home) + .iter() + .any(|text| text.contains("mcp_task_result")), + "nothing is journaled for a server that is gone from the configuration" + ); +} + +/// Lines a model could append to the session `.jsonl` with the ungated write/edit tools: a planted +/// `mcp_task` (pointing the real unanswered call at another task id) and a planted +/// `mcp_task_result` (a made-up answer). Neither carries the host's seal, so on replay the task id +/// is never polled and the made-up result is never delivered, to the model or to `get_messages`. +#[test] +fn journal_lines_planted_in_the_session_file_are_ignored_on_resume() { + let env = Env::new(); + let _server = env.http(); + let real = start_and_kill(&env, "mcp__t__gated_task", None, ""); + let text = env.session_text(); + let header: Value = serde_json::from_str(text.lines().next().unwrap()).unwrap(); + let session_id = header["id"].as_str().unwrap().to_owned(); + let tip = text + .lines() + .filter_map(|l| serde_json::from_str::(l).ok()) + .filter_map(|v| v["id"].as_str().map(str::to_owned)) + .next_back() + .unwrap(); + let planted = [ + json!({ + "type": "custom", "id": "planted-task", "parent_id": tip, "timestamp": 1, + "kind": "mcp_task", + "data": { + "server": "t", "tool": "gated_task", "taskId": "PLANTED-TASK", + "createdAtMs": 0, "toolUseId": "toolu_g", "sessionId": session_id, + }, + }), + json!({ + "type": "custom", "id": "planted-result", "parent_id": "planted-task", "timestamp": 1, + "kind": "mcp_task_result", + "data": { + "toolUseId": "toolu_g", "name": "mcp__t__gated_task", + "content": "PLANTED-RESULT", "isError": false, "sessionId": session_id, + }, + }), + ]; + let mut file = std::fs::OpenOptions::new() + .append(true) + .open(env.session_file()) + .unwrap(); + for line in planted { + writeln!(file, "{line}").unwrap(); + } + drop(file); + + std::fs::write(&env.gate, b"open").unwrap(); + let mut s = serve(&env, vec![turn_text("carrying on"), turn_text("t")]); + send(&mut s.stdin, json!({ "type": "get_messages", "id": "m" })); + let got = read_until(&mut s.stdout, "the transcript", |f| { + f["type"] == "response" && f["command"] == "get_messages" + }); + assert!( + !got.last().unwrap().to_string().contains("PLANTED-RESULT"), + "a planted result is not part of the transcript" + ); + prompt(&mut s.stdin, "what happened?"); + let frames = read_until_prompt_done(&mut s.stdout); + let end = tool_end(&frames, "toolu_g"); + assert!( + end["result"].as_str().unwrap().contains("gated-done"), + "{end}" + ); + let requests = s.requests(); + let sent = tool_result_sent(requests.last().unwrap(), "toolu_g"); + assert!(sent.to_string().contains("gated-done"), "{sent}"); + assert!(!requests.last().unwrap().contains("PLANTED-RESULT")); + s.close(); + let polled: Vec = env.methods("tasks/get"); + assert!( + polled + .iter() + .all(|g| g["params"]["taskId"] == real.as_str()), + "only the real task is ever polled: {polled:?}" + ); +} + +/// Copy a single-file session (its `.jsonl` and every sidecar beside it) from `from` to `to`, the +/// way a user moves a session to another machine or a fresh `$HOME`. +fn copy_session(from: &Env, to: &Env) { + for entry in std::fs::read_dir(from.dir.path()).unwrap().flatten() { + let name = entry.file_name(); + if name.to_string_lossy().starts_with("s.") && entry.path().is_file() { + std::fs::copy(entry.path(), to.dir.path().join(&name)).unwrap(); + } + } +} + +/// The journal's key lives with the session, not with the machine: a session copied to a fresh +/// `$HOME` (no agent config there at all) still resumes the task it left in flight, and copied on +/// again once that task's result is journaled, still shows the result. +#[test] +fn a_session_copied_to_a_fresh_home_still_resumes_and_shows_its_journal() { + let first = Env::new(); + let (_server, entry) = first.http_server(); + first.settings(json!([entry.clone()])); + start_and_kill(&first, "mcp__t__gated_task", None, ""); + + let second = Env::new(); + second.settings(json!([entry.clone()])); + copy_session(&first, &second); + let mut s = serve(&second, vec![turn_text("t")]); + read_until(&mut s.stdout, "the copied session's task resumed", |f| { + is_event(f, "tool_progress", "toolu_g") + }); + std::fs::write(&first.gate, b"open").unwrap(); + wait_for("the result journaled in the copy", || { + second.session_text().contains("mcp_task_result") + }); + s.close(); + + let third = Env::new(); + third.settings(json!([entry])); + copy_session(&second, &third); + let gets = first.methods("tasks/get").len(); + let mut s = serve(&third, vec![turn_text("carrying on"), turn_text("t")]); + prompt(&mut s.stdin, "what happened?"); + let frames = read_until_prompt_done(&mut s.stdout); + let end = tool_end(&frames, "toolu_g"); + assert!( + end["result"].as_str().unwrap().contains("gated-done"), + "{end}" + ); + let requests = s.requests(); + assert_alternates(&requests[0]); + let sent = tool_result_sent(&requests[0], "toolu_g"); + assert!(sent.to_string().contains("gated-done"), "{sent}"); + send(&mut s.stdin, json!({ "type": "get_messages", "id": "m" })); + let got = read_until(&mut s.stdout, "the transcript", |f| { + f["type"] == "response" && f["command"] == "get_messages" + }); + let transcript = got.last().unwrap().to_string(); + assert!(transcript.contains("gated-done"), "{transcript}"); + s.close(); + assert_eq!( + first.methods("tasks/get").len(), + gets, + "a journaled result is not polled again" + ); +} + +/// A "result pending" placeholder is marked in the journal, never recognised by its text: a task +/// that completed in its own run with a result that merely begins like a placeholder is answered, +/// so a restart neither polls it again nor answers it a second time. +#[test] +fn a_real_result_that_reads_like_a_placeholder_is_still_an_answer() { + let env = Env::new(); + let lookalike = format!("{PENDING} no, this is the real answer"); + env.settings(json!([env.stdio_with(json!({ + "MCP_TASKS_FIXTURE_GATED_TEXT": lookalike, + }))])); + std::fs::write(&env.gate, b"open").unwrap(); + let mut s = serve( + &env, + vec![ + turn_tool_use("toolu_g", "mcp__t__gated_task", "{}"), + turn_text("done"), + turn_text("t"), + ], + ); + prompt(&mut s.stdin, "start the job"); + let frames = read_until_prompt_done(&mut s.stdout); + let end = tool_end(&frames, "toolu_g"); + assert!( + end["result"].as_str().unwrap().contains("the real answer"), + "{end}" + ); + s.close(); + assert!(env.session_text().contains("\"kind\":\"mcp_task\"")); + let gets = env.methods("tasks/get").len(); + + let mut s = serve(&env, vec![turn_text("carrying on"), turn_text("t")]); + prompt(&mut s.stdin, "what happened?"); + let frames = read_until_prompt_done(&mut s.stdout); + assert!( + !frames + .iter() + .any(|f| is_event(f, "tool_progress", "toolu_g") || is_event(f, "tool_end", "toolu_g")), + "an answered call is not resumed: {frames:?}" + ); + let request = s.requests()[0].clone(); + s.close(); + assert_eq!(env.methods("tasks/get").len(), gets, "never polled again"); + let answers = common::body_json(&request)["messages"] + .as_array() + .unwrap() + .iter() + .flat_map(|m| m["content"].as_array().cloned().unwrap_or_default()) + .filter(|b| b["type"] == "tool_result" && b["tool_use_id"] == "toolu_g") + .count(); + assert_eq!(answers, 1, "{request}"); + assert!(!env.session_text().contains("mcp_task_result")); +} diff --git a/crates/agent/tests/mcp_tasks_service.rs b/crates/agent/tests/mcp_tasks_service.rs new file mode 100644 index 00000000..82fc8924 --- /dev/null +++ b/crates/agent/tests/mcp_tasks_service.rs @@ -0,0 +1,148 @@ +//! SEP-2663 tasks in service mode: a session's journal lives in its sealed segments, which already +//! authenticate every line under the tenant's key, so whichever replica owns the session next +//! resumes its tasks (see `crate::mcp_resume::JournalAuth`). +#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)] + +mod common; +mod mcp_tasks_env; + +use std::time::{Duration, Instant}; + +use common::service::{Options, Service}; +use common::{ + TestWs, spawn_model_server, turn_text, turn_tool_use, ws_connect_with_headers, ws_next_frame, + ws_read_until_response, ws_send, +}; +use mcp_tasks_env::Env; +use serde_json::{Value, json}; + +/// Attach to `session_id` on `port` with a grant naming the tasks fixture as connector `t`, retrying +/// while another replica still holds it. +async fn attach(svc: &Service, port: u16, session_id: &str, url: &str) -> TestWs { + let mut claims = svc.claims("tenant-a", session_id, &svc.shards[0].0); + claims.mcp = vec![("t".to_string(), url.to_string())]; + let token = svc.minter.mint(&claims, &svc.secrets()); + let deadline = Instant::now() + Duration::from_secs(30); + loop { + match ws_connect_with_headers(port, Some(session_id), &svc.header(&token)).await { + Ok(ws) => return ws, + Err(status) => assert!( + Instant::now() < deadline, + "{session_id} still refused with HTTP {status}" + ), + } + tokio::time::sleep(Duration::from_millis(100)).await; + } +} + +/// Frames until `stop` matches, failing after 30 s however busy the socket is (its keepalives +/// included). +async fn read_until(ws: &mut TestWs, what: &str, stop: impl Fn(&Value) -> bool) -> Vec { + let mut frames = Vec::new(); + let read = tokio::time::timeout(Duration::from_secs(30), async { + loop { + let Some(frame) = ws_next_frame(ws).await else { + return false; + }; + let done = stop(&frame); + frames.push(frame); + if done { + return true; + } + } + }) + .await; + match read { + Ok(true) => frames, + Ok(false) => panic!("socket closed before {what}: {frames:#?}"), + Err(_) => panic!("{what}: not within 30s: {frames:#?}"), + } +} + +fn event(frame: &Value, kind: &str) -> bool { + frame["type"] == "event" && frame["event"]["kind"] == kind && frame["event"]["id"] == "toolu_g" +} + +/// A task left in flight by one replica (killed outright) is resumed by the next owner of the +/// session, another process on the same store: polled in the background on attach, its result +/// journaled, and the next prompt answered with it. +#[tokio::test] +async fn a_task_left_in_flight_on_one_replica_resumes_on_another() { + let env = Env::new(); + let (_server, entry) = env.http_server(); + let url = entry["url"].as_str().unwrap().to_owned(); + let (base, requests) = spawn_model_server(vec![ + turn_tool_use("toolu_g", "mcp__t__gated_task", "{}"), + turn_text("carrying on"), + turn_text("t"), + turn_text("t"), + ]); + let opts = Options { + extra_args: vec!["--mcp-allow-private".into()], + ..Default::default() + }; + let mut svc = Service::start_with(&base, &["s1"], opts).await; + let peer = svc.start_peer(&["--mcp-allow-private"]); + + let mut ws = attach(&svc, svc.port, "s1.alpha", &url).await; + ws_send( + &mut ws, + json!({ "type": "prompt", "id": "p1", "message": "start the job" }), + ) + .await; + read_until(&mut ws, "the task's record", |f| { + event(f, "tool_progress") && f["event"]["details"]["mcpTask"].is_object() + }) + .await; + // The record is journaled as it is emitted; give the append a moment, then kill the owner. + let gets = env.methods("tasks/get").len(); + let deadline = Instant::now() + Duration::from_secs(20); + while env.methods("tasks/get").len() == gets { + assert!(Instant::now() < deadline, "the task was never polled"); + tokio::time::sleep(Duration::from_millis(50)).await; + } + svc.child.kill().unwrap(); + let _ = svc.child.wait(); + drop(ws); + + let mut ws = attach(&svc, peer.port, "s1.alpha", &url).await; + read_until(&mut ws, "the task resumed on the other replica", |f| { + event(f, "tool_progress") + }) + .await; + std::fs::write(&env.gate, b"open").unwrap(); + ws_send( + &mut ws, + json!({ "type": "prompt", "id": "p2", "message": "what happened?" }), + ) + .await; + let frames = ws_read_until_response(&mut ws, "prompt").await; + assert_eq!(frames.last().unwrap()["success"], true, "{frames:#?}"); + let end = frames + .iter() + .find(|f| event(f, "tool_end")) + .unwrap_or_else(|| panic!("no tool_end for the resumed call: {frames:#?}")); + assert_eq!(end["event"]["is_error"], false, "{end}"); + assert!( + end["event"]["result"] + .as_str() + .unwrap() + .contains("gated-done"), + "{end}" + ); + let sent = requests + .lock() + .unwrap() + .iter() + .rfind(|r| !r.contains("You write short titles")) + .cloned() + .unwrap(); + mcp_tasks_env::assert_alternates(&sent); + let result = mcp_tasks_env::tool_result_sent(&sent, "toolu_g"); + assert!(result.to_string().contains("gated-done"), "{result}"); + assert_eq!( + env.methods("tools/call").len(), + 1, + "resumed, not re-invoked" + ); +} diff --git a/crates/agent/tests/session_segments_sealing.rs b/crates/agent/tests/session_segments_sealing.rs index 7e25a934..5002e327 100644 --- a/crates/agent/tests/session_segments_sealing.rs +++ b/crates/agent/tests/session_segments_sealing.rs @@ -89,6 +89,67 @@ fn transcripts_listings_and_memory_hold_no_plaintext() { assert_eq!(session.messages.len(), 1); } +/// MCP task handles journaled with a service session (task ids, which a server may treat as bearer +/// tokens, and resolved results) are sealed with the tenant's key like the rest of the transcript, +/// and read back intact with it, so a resume works on sealed storage. +#[test] +fn mcp_task_journal_entries_are_sealed_and_read_back() { + let dir = TempDir::new().unwrap(); + let repo = sealed_repo(dir.path(), "tenant-a", &DEK); + let mut store = repo.create(SessionMeta::with_id("s1", "/w", "m")).unwrap(); + store.append_new(&[Message::user("start the job")]).unwrap(); + // Through the journal path, which is where a key would be made if this storage were given one. + store + .append_journal( + "mcp_task", + serde_json::json!({ + "server": "srv", "tool": "job", "taskId": "SECRET-TASK-ID", + "createdAtMs": 1, "toolUseId": "toolu_1", "sessionId": "s1", + }), + ) + .unwrap(); + store + .append_journal( + "mcp_task_result", + serde_json::json!({ + "toolUseId": "toolu_1", "name": "mcp__srv__job", + "content": "SECRET-RESULT", "isError": false, "sessionId": "s1", + }), + ) + .unwrap(); + drop(store); + + assert_no_plaintext(dir.path(), &["SECRET-TASK-ID", "SECRET-RESULT"]); + + let (store, _session) = repo.open_or_create_id("s1", "/w", "m").unwrap(); + // The sealed store authenticates every line itself: the journal needs no key of its own (so no + // key file sits beside the session, and any replica holding the tenant key reads it). + let tasks = store.active_journal("mcp_task"); + assert_eq!(tasks.len(), 1); + assert_eq!(tasks[0]["taskId"], "SECRET-TASK-ID"); + assert_eq!(store.active_journal("mcp_task_result").len(), 1); + let names: Vec = std::fs::read_dir(dir.path().join("s1")) + .unwrap() + .flatten() + .map(|e| e.file_name().to_string_lossy().into_owned()) + .collect(); + assert!( + names.iter().all(|n| !n.contains("mcp-task-journal")), + "{names:?}" + ); + let beside: Vec = std::fs::read_dir(dir.path()) + .unwrap() + .flatten() + .map(|e| e.file_name().to_string_lossy().into_owned()) + .collect(); + assert!( + beside.iter().all(|n| !n.contains("mcp-task-journal")), + "{beside:?}" + ); + // Nor is any entry given a MAC: nothing on top of the storage's own seal. + assert!(tasks.iter().all(|t| t.get("mac").is_none()), "{tasks:?}"); +} + #[test] fn a_base_seals_its_content_too() { // Compaction replaces the whole transcript at once. That content is still transcript.