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.