From 14ec00c881d6591590502473f0e9ca3701728da4 Mon Sep 17 00:00:00 2001 From: Jared Lunde Date: Wed, 7 Oct 2026 03:24:50 -0700 Subject: [PATCH] agent: MCP OAuth fixture coverage, HTTP sessions ended on exit, OFD session locks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit One commit for #144, rebased onto #137/#148/#150 (0602cac). - OAuth fixture: full streamable-HTTP server (sessions + 404 expiry, SSE streamed event by event with an in-POST elicitation, GET stream, DELETE). Routing tests: DELETE with bearer + session on reap and on exit, in-POST question routed and answered mid-stream. - HTTP MCP sessions get a bounded, parallel, authenticated DELETE on every graceful exit (mcp_http_exit); sessions held weakly, orphan DELETEs dropped once their attempt is over; the registry is a value, tests use their own. - file_lock: open file description locks (F_OFD_SETLK via nix), released by an explicit unlock — no other descriptor releases them, a forked child's copy cannot keep them, two descriptions conflict in one process. Off Linux, flock with explicit LOCK_UN. Legacy flock half for mixed rollout. Target::Itself for #137's journal key: locked itself, written through the locked descriptor, no legacy flock (it would self-conflict on NFS). Verified on loopback NFS v4.2/v4.0/v3. - A failed start's session dir: lock released and closed before its files are unlinked (NFS silly-renames open files and rmdir then fails). - write_atomic never renames over a lock file, record or legacy (only a legacy one beside a record one, so Cargo.lock is untouched). - apps cache test waits for list_changed deterministically; output.rs sweep test uses a unique prefix. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01JimHGjsfk2Ktm5GxyZJKKk --- Cargo.lock | 1 + crates/agent/ARCHITECTURE.md | 109 +++- crates/agent/Cargo.toml | 5 + crates/agent/src/file_lock.rs | 600 ++++++++++++++---- crates/agent/src/serve_ws.rs | 24 +- crates/agent/src/session_store.rs | 197 +++++- crates/agent/src/tools/mcp.rs | 18 +- crates/agent/src/tools/mcp_http_exit.rs | 422 ++++++++++++ crates/agent/src/tools/mcp_manifest.rs | 15 +- crates/agent/src/tools/mcp_oauth.rs | 85 ++- crates/agent/src/tools/mcp_stdio.rs | 7 + crates/agent/src/tools/mod.rs | 40 ++ crates/agent/src/tools/output.rs | 8 +- crates/agent/src/worktree.rs | 51 +- .../agent/tests/common/mcp_oauth_fixture.rs | 94 +++ crates/agent/tests/mcp_apps_session.rs | 37 +- crates/agent/tests/mcp_oauth_routing.rs | 173 ++++- 17 files changed, 1683 insertions(+), 203 deletions(-) create mode 100644 crates/agent/src/tools/mcp_http_exit.rs diff --git a/Cargo.lock b/Cargo.lock index 66c4b2b1..d37a442e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -432,6 +432,7 @@ dependencies = [ "libc", "listenfd", "mimalloc", + "nix", "percent-encoding", "prometheus", "pulldown-cmark", diff --git a/crates/agent/ARCHITECTURE.md b/crates/agent/ARCHITECTURE.md index 7933d986..a3879ead 100644 --- a/crates/agent/ARCHITECTURE.md +++ b/crates/agent/ARCHITECTURE.md @@ -3085,16 +3085,72 @@ coherent across a tenant's sessions. ### Locks -`acquire_session_lock(path)` opens `/lock` (or `.lock`) read+write+create and takes -`File::try_lock`. `Ok(None)` means someone holds it right now. It is **liveness only** — correctness -is the epoch fence, which needs no lock — so a lock lost to a crash or a partition costs a retry, -never history. Three details that matter on EFS: - -- read+write, because NFS emulates `flock` with POSIX record locks, which need a writable descriptor; -- after locking, `fstat` on the held descriptor is compared with `stat` of the path, catching a - directory renamed into `.trash/` between the open and the lock; -- a session is opened **once per process**, because closing any descriptor to a POSIX-locked file - drops the lock. +`acquire_session_lock(path)` is `file_lock::try_lock` on the session — `Target::Dir` for a session +directory, `Target::File` for a `.jsonl` (the MCP manifest's write lock is the same primitive on the +manifest file). `Ok(None)` means someone holds it right now. It is **liveness only** — correctness is +the epoch fence, which needs no lock — so a lock lost to a crash or a partition costs a retry, never +history. + +**An open file description (OFD) lock** on Linux: exclusive, non-blocking, whole file +(`fcntl(F_OFD_SETLK)`, through `nix`), on a lock file only `file_lock` names (`/.beyond-lock`, +`.beyond-lock`). It belongs to the one `open` that took it, which settles the two ways a lock used +to be lost: + +- **Another descriptor of the same file** — a file tool, the memory tool, an `@file` read, reaching it + by name, through a symlink or a hard link — is a different description: opening and closing it + neither releases this lock nor shares it. (A POSIX record lock belongs to the process, so closing + any of its descriptors of the file dropped it; that made every in-process opener a hazard.) Two + descriptions conflict even inside one process, so a second acquire here, under any spelling of the + path, is refused by the kernel — there is no in-process registry. +- **A forked child.** A child forked by any thread (a `bash` call, an MCP stdio server) shares the + description until it execs. The lock is released with an explicit `F_UNLCK` on its descriptor before + that descriptor closes, which drops it whatever copies a child holds. (`flock` released by closing + is what made a just-released lock look held 180 times in 3000 while other threads spawned + processes — the root cause of `serve_ws::tests::a_session_directory_with_a_segment_is_left_alone` + flaking under load.) + +Off Linux there are no OFD locks; there it is `flock`, which also belongs to the description, released +the same way with an explicit `LOCK_UN`. + +**On NFS/EFS** the Linux client sends an OFD lock to the server as a POSIX byte-range lock owned by the +open file description, so the same two properties hold there. Verified against a loopback NFSv4.2 +mount (`local_lock=none`, so every lock goes to the server): the `file_lock` tests and the session-lock +probes below all pass with their files on the mount. (`flock` on NFS is emulated with a process-owned +POSIX lock, which is why it is not the authoritative half.) + +Nothing else needs guarding: the only remaining hazard is _replacing_ a lock file, which an atomic +write (`tools::write_atomic`) could do by renaming over it — it refuses a lock file as its target +(`file_lock::is_lock_file`: the record lock file, or a legacy `lock` / `.lock` sitting beside one, +so an unrelated `Cargo.lock` is written as ever; applied after resolving a symlink; a hard link's +rename replaces only the link). A worktree seed leaves lock files behind (a copy of a lock is no +lock). A failed session start's directory is taken back through `file_lock::dir_holds_only_lock_files` +and `FileLock::release_and_remove_files`, which unlocks and **closes** the lock files before unlinking +them: an NFS client silly-renames a file it still holds open to `.nfs*`, which kept the directory from +being removed (`serve_ws::tests::an_empty_session_directory_is_taken_back_with_its_lock` fails on a +loopback NFS mount with the old unlink-then-close order, and passes with this one). + +A small file whose holder writes it under its own lock — the MCP task journal key — is locked +**itself** (`Target::Itself`): no lock file beside it, and the key is re-read and written through the +locked descriptor (`FileLock::file`). No legacy `flock` there: no released binary took one, and an +`flock` on the same file would conflict with the process's own OFD lock on NFS, where `flock` is a +process-owned POSIX lock. + +**Mixed rollout.** Binaries before this one took only an `flock` on `/lock` / `.lock`. So +`try_lock` also takes that `flock` (its own descriptor, close-on-exec, released with an explicit +`LOCK_UN`), and a held one means held: an old binary still excludes this one and the other way round. +It is retried for up to 250 ms, since an old binary's own fork→exec window can hold it briefly. **The +legacy `flock` half is removed in a later release, once no older binary can run against the same +files.** + +Pinned by `session_store::tests::a_released_session_lock_is_free_at_once_while_other_threads_spawn_processes` +(0 of 3000 with 8 threads spawning), +`every_in_process_reader_reaching_the_lock_file_leaves_it_held` (read, write, edit, search, the memory +tool's view and an `@file` read, through a symlink and a hard link: still held for another process, +refused a second time here, the file untouched), `reading_a_locked_session_through_any_public_api_leaves_it_locked`, +and `file_lock`'s own tests (another descriptor neither releases nor shares it, a child holding the +description cannot keep a released lock, one holder per process, an old binary's `flock` excluding this +one). After locking, `fstat` on the held descriptor is compared with `stat` of the path, catching a +directory renamed into `.trash/` between the open and the lock. `SessionRepo::delete` on a segmented repo takes that lock first and refuses if it is held, then renames `/` → `.trash/.`. The timestamp suffix is what makes repeated deletes of one @@ -3758,7 +3814,7 @@ context. Only host-authored records count: a skill-tool result (matched to its ` its `skill://` resources are ordinary resources. - **Manifest cache.** `mcp_manifest` (version 3) records the entries (`None` = extension not declared) and their diagnostics, so a reaped or never-started server still lists its skills; the - first load dials it (and re-fetches the entry, whose freshness a cache cannot vouch for). Every write is a read-modify-write of one file shared by all connections (a connect, a forget, a refreshed listing), so each runs under `.lock` — `file_lock::try_lock` — the kernel-lock-plus-registry primitive the session lock (`session_store::acquire_session_lock`) is also built on, a kernel file lock that dies with its holder, so a crashed writer's leftover file is simply unlocked and no two waiters can both take it — and renames a temporary of its own into place; concurrent refreshes lose nobody's update. Waiting for the lock blocks, so the public writers (`store`, `forget`, `store_skills`) are `async` and run it with `spawn_blocking`: a connection task never holds an async worker while another writer finishes. A write that panics leaves the file as it was and is logged at warn. + first load dials it (and re-fetches the entry, whose freshness a cache cannot vouch for). Every write is a read-modify-write of one file shared by all connections (a connect, a forget, a refreshed listing), so each runs under `file_lock::try_lock` on the manifest file — the OFD-lock primitive the session lock (`session_store::acquire_session_lock`) is also built on, a kernel file lock that dies with its holder, so a crashed writer's leftover file is simply unlocked and no two waiters can both take it — and renames a temporary of its own into place; concurrent refreshes lose nobody's update. Waiting for the lock blocks, so the public writers (`store`, `forget`, `store_skills`) are `async` and run it with `spawn_blocking`: a connection task never holds an async worker while another writer finishes. A write that panics leaves the file as it was and is logged at warn. - **Nested-request attribution.** Every skills request (`skills/list`, `skills/get`, `resources/read`) goes through `mcp::request_tracked`, registered with `track_call` under the calling session's host for as long as it is outstanding, so a nested request the server raises @@ -4170,6 +4226,7 @@ backend an operator points it at. | `src/tools/mcp_events/` | MCP Events client (**draft** extension, pinned design-sketch commit): `mod.rs` (hub, ownership, configured subscriptions kept up, durable delivery, coalescer), `webhook.rs` (receiver behind `/_beyond/mcp-events/`, Standard Webhooks `v1,`/`v1a,`, fail-closed key policy, persisted callback, subscribe/refresh/rotate loop), `state.rs` (the session's snapshot and append-only pending log), `wire.rs` (notification router, direct HTTP/SSE, RFC 3339). See "MCP Events (draft extension)" | | `src/tools/mcp_stdio.rs` | The transport for **every** stdio MCP server: spawns it (process group, inherited stderr), skips non-UTF-8 stdout lines, gives it a 3 s grace after stdin closes, then `retire` kills what is left of its group on a tracked thread — on a reap, a registry rebuild and every `run`/`serve` exit alike (`retire_all`) — and rescues custom results rmcp 3.x would decode as empty tool results (`rescue`/`unwrap_rescued`, one boundary); caps each inbound message (`read_capped`: 64 MiB, and `MAX_VIEW_BYTES` for the answer to a `ui://` `resources/read` it wrote) without holding it whole, answering an over-cap response with a JSON-RPC error only when its first 4 KiB prove which host request it answers (`scan_head`), and dropping anything else over the cap | | `src/tools/mcp_oauth.rs` | Mid-session MCP OAuth refresh: `ServerAuth` (one configured server's bearer token, shared by every connection to it, reloaded from the `mcp-login` store at each dial; `after_rejection` refreshes once per rejected token, single-flight, persisted, a failure naming `agent mcp-login`) and `OAuthHttp` (the outermost streamable-HTTP client layer: attaches the current token and retries a 401 once). See "MCP Server OAuth" | +| `src/tools/mcp_http_exit.rs` | Ends every streamable-HTTP MCP session still open on the way out of the process: `HttpSession` (the session a connection opened, recorded by `OAuthHttp`, cleared when rmcp sends its own `DELETE`) and `begin_close_all` (parallel authenticated `DELETE`s on a thread with its own runtime and a fresh pool-less client, bounded by `CLOSE_DEADLINE`), run from `mcp_stdio::sweep_before_exit`. See "MCP Server OAuth" | | `src/bin/mcp_fixture_stdio_server.rs` | Test fixture only, not a real MCP server: a ~150-line hand-rolled MCP server speaking newline-delimited JSON-RPC over stdio, with zero new dependencies (`tokio`/`serde_json`, already ordinary deps of this crate) rather than pulling in `rmcp`'s own server-side machinery just for a test double. Auto-discovered by Cargo as a sibling `[[bin]]` of this same package (any `.rs` file under `src/bin/` becomes its own binary target), so `tests/mcp_client.rs` locates it via `env!("CARGO_BIN_EXE_mcp_fixture_stdio_server")` exactly like every other e2e test locates the real `beyond-ai-agent` binary. Tools include `echo`/`add`/`ping`/`fail`/`echo_env`/`image`/`slow_progress`/`ask_user`/`ask_user_mrtr`/`slow_task`, each proving one distinct thing the real client code must handle (see the file's own module doc comment). | | `src/bin/mcp_apps_fixture_server.rs` | Test fixture only: a hand-rolled MCP Apps (SEP-1865) server over stdio (`tokio`/`serde_json`, no `rmcp` server side) for `tests/mcp_apps*.rs`. Lists `app_only_always` (`visibility: ["app"]`) unconditionally, as the official SDK does, and the rest of its UI metadata (`_meta.ui.resourceUri`, `visibility`, an app-only `refresh_weather`, an oversized `huge_view`, a non-`ui://` `bad_uri_view`) only when the client advertised `io.modelcontextprotocol/ui` — read per request under `2026-07-28` or from `initialize` (`MCP_APPS_FIXTURE_LEGACY=1` forces the latter) — and serves a `ui://` view that speaks the iframe bridge in plain JavaScript, so it also renders in a real browser. `MCP_APPS_FIXTURE_LOG` records every request (with its `uri`) and whether its client advertised the extension; `change_view` bumps the view's version and sends `notifications/resources/list_changed`; `MCP_APPS_FIXTURE_FAIL_FIRST_UI`, `_UI_DELAY_MS` and `_VIEW_DELAY_MS` make the first apps dial fail, the apps dial slow, and the view read slow. | | `src/bin/mcp_fixture_events_server.rs` | Hand-rolled MCP Events **server** fixture (stdio or streamable HTTP) with a control API for tests: emits events, signs webhook deliveries (Standard Webhooks, dual-signing after a rotation, tamper modes), runs the verification challenge, grants TTLs, terminates subscriptions, and records what it observed | @@ -4428,11 +4485,33 @@ error (rmcp's unexpected content type), and a 4xx to `server/discover` as the le is an error here. rmcp calls it accepted and then waits for a response that cannot come (a request is answered on its own POST, as JSON or an SSE stream carrying it) until the request's timeout; failing at once with the body in the error is the honest answer. Pinned end to end against the OAuth fixture -— which speaks sessions (and their 404 expiry), SSE responses, the `GET` stream and `DELETE` — in +— which speaks sessions (and their 404 expiry), SSE responses streamed event by event (a progress +notification and an `elicitation/create` on the call's own POST, the response held until the client's +answer has come back on another POST), the `GET` stream and `DELETE` — in `tests/mcp_oauth_routing.rs` (a 401 with a JSON-RPC body through a real dial, the legacy -`server/discover` 4xx, an empty-200 notification, sessions with SSE, an expired session), and the -request/notification difference in `mcp_wire`'s -`a_json_success_that_is_not_json_rpc_is_accepted_for_a_notification_but_fails_a_request`. **Single flight:** the refresh runs under the token's lock, and +`server/discover` 4xx, an empty-200 notification, sessions with SSE, an expired session, an in-POST +question routed to the client and answered mid-stream, and a reaped connection's DELETE carrying the +bearer and session id), and the request/notification difference in `mcp_wire`'s +`a_json_success_that_is_not_json_rpc_is_accepted_for_a_notification_but_fails_a_request`. + +**Sessions end on the way out** (`tools/mcp_http_exit.rs`). rmcp sends a session's `DELETE` when a +connection is closed while the runtime runs (the idle reaper, a registry rebuild); a process that +exits (`agent run` finishing, a `serve` shutdown) never got that far. So every HTTP connection records +the session it opened (`OAuthHttp` sees each response's `Mcp-Session-Id` and each `DELETE` rmcp sends; +the process holds it only weakly, so a connection that goes away takes its entry with it — one that +goes with its session still open sends that `DELETE` at once on its runtime and is dropped once that +attempt is over, answered or timed out — kept for the exit only when dropped by the runtime's own +teardown), +and `mcp_stdio::sweep_before_exit` — every graceful way out, `exit_process` and `ExitSweep` alike — +first sends a `DELETE` for every session still open: in parallel, with the server's current bearer and +the connection's own headers, on a thread with its own runtime and a fresh pool-less client (the +process runtime may be gone, or blocked in the very call that is exiting; a grant connector's client +keeps its SSRF resolver), bounded by `CLOSE_DEADLINE` (1.5 s) overlapping the stdio grace. A server that +never answers costs at most the deadline. Pinned in `tests/mcp_oauth_routing.rs`: an `agent run` sends +exactly one authenticated `DELETE` before it exits, and one whose `DELETE` is held for 30 s still +exits within its deadline. + +**Single flight (refresh):** the refresh runs under the token's lock, and a caller whose rejected token has already been replaced takes the replacement without refreshing, so concurrent 401s cost one refresh (a second refresh with a rotated-away refresh token would fail). **A failed refresh is remembered:** a _definitive_ one for good, with an error naming `agent mcp-login diff --git a/crates/agent/Cargo.toml b/crates/agent/Cargo.toml index 453805ed..392aab2d 100644 --- a/crates/agent/Cargo.toml +++ b/crates/agent/Cargo.toml @@ -148,6 +148,11 @@ htmd = "0.5" # Already resolved transitively at this exact version (tokio → signal-hook-registry → errno), so # depending on it directly adds nothing to the build graph. libc = "0.2" +# The session and manifest locks are open file description locks (`fcntl(F_OFD_SETLK)`, see +# `file_lock`): owned by the one `open` that took them, so no other descriptor of the file can release +# them, and released explicitly so a forked child's copy cannot keep them. Already in the build graph +# at this version (transitively); its safe `fcntl` keeps the `unsafe` out of this crate, which forbids it. +nix = { version = "0.31", default-features = false, features = ["fs"] } # `export_html`'s markdown rendering (headers/bold/lists/code fences) for message text — pure Rust, no # vendored JS. Deliberately not paired with a syntax-highlighting crate (e.g. `syntect`): that pulls in # several MB of bundled syntax/theme data and would slow every build of this CLI, including `run`/ diff --git a/crates/agent/src/file_lock.rs b/crates/agent/src/file_lock.rs index c3bf9d72..1cdb478a 100644 --- a/crates/agent/src/file_lock.rs +++ b/crates/agent/src/file_lock.rs @@ -1,101 +1,311 @@ -//! An advisory lock on one file: a kernel lock (`flock`, emulated with POSIX record locks over NFS) -//! plus a process-wide registry of the paths this process holds. +//! An advisory lock on one resource (a session directory, a session file, the MCP manifest). //! //! The kernel releases the lock when its holder exits, however it exits, so there is no staleness to -//! judge and no lockfile to break: a crashed holder's leftover file is simply unlocked. The registry -//! is there because closing *any* descriptor to a POSIX-locked file drops the lock — a second open of -//! the same path in this process would quietly release the first one's hold — so a second attempt in -//! the same process reports "held" instead of opening it. +//! judge and no lockfile to break: a crashed holder's leftover file is simply unlocked. Used for a +//! session's liveness lock (`session_store::acquire_session_lock`) and the MCP manifest cache's write +//! lock (`tools::mcp_manifest`). Non-blocking: [`try_lock`] answers at once (bar a short retry, below), +//! and a caller that wants to wait polls it off the async runtime. //! -//! Used for a session's liveness lock (`session_store::acquire_session_lock`) and the MCP manifest -//! cache's write lock (`tools::mcp_manifest`). Non-blocking: [`try_lock`] answers at once, and a -//! caller that wants to wait polls it off the async runtime. +//! **The lock is an open file description (OFD) lock** on Linux (`fcntl(F_OFD_SETLK)`, whole file, +//! through `nix`), on a lock file only this module names. What it belongs to is the one `open` this +//! module made, which settles the two ways a lock used to be lost: +//! +//! - **Another descriptor of the same file** — a file tool or the memory tool reading it through a +//! symlink or a hard link, anything — is a different open file description: opening and closing it +//! neither releases this lock nor shares it. (A POSIX record lock belongs to the *process*, so +//! closing *any* of its descriptors of the file dropped it.) Two descriptions conflict even inside +//! one process, so a second acquire here — under any spelling of the path — is refused by the +//! kernel itself; there is no in-process registry to keep in step. +//! - **A forked child.** A child forked by any thread (a `bash` call, an MCP stdio server) shares the +//! description until it execs, and with it the lock. So the lock is released with an explicit +//! `F_UNLCK` on this descriptor before it closes, which drops it whatever copies a child holds. +//! (`flock` without that explicit unlock is what made a just-released lock look held 180 times in +//! 3000 while other threads spawned processes.) +//! +//! Off Linux there are no OFD locks: there it is `flock`, which also belongs to the open file +//! description, released the same way with an explicit `LOCK_UN`. +//! +//! **On NFS/EFS** the Linux client sends an OFD lock to the server as a POSIX byte-range lock whose +//! owner is the open file description, not the process — so another description's open/close does not +//! release it there either, and a second description on the same client conflicts with it (verified +//! against a loopback NFSv4.2 mount; see ARCHITECTURE.md). `flock` on NFS, by contrast, is emulated +//! with a process-owned POSIX lock, which is exactly why it is not the authoritative half. +//! +//! **Mixed rollout.** Binaries before this one took an `flock` on `/lock` / `.lock` and +//! nothing else. So, for now, [`try_lock`] also takes that `flock` (on its own descriptor, +//! close-on-exec, released with an explicit `LOCK_UN`), and a held one means "held": an old binary +//! still excludes this one, and this one still excludes an old one. It is taken with a short bounded +//! retry ([`LEGACY_RETRY`]), since an old binary's own fork→exec window can briefly hold it. The OFD +//! lock is the authoritative one; **the legacy `flock` half is removed in a later release, once no +//! binary older than this one can be running against the same files.** -use std::collections::HashSet; +use std::ffi::{OsStr, OsString}; use std::fs::{File, OpenOptions}; use std::path::{Path, PathBuf}; -use std::sync::{Mutex, OnceLock}; +use std::time::{Duration, Instant}; /// How many times [`try_lock`] re-opens after finding it locked a file that is no longer the one at /// the path (the file, or a directory holding it, renamed away mid-acquire). const RETRIES: usize = 5; -/// A held lock. Dropping it closes the descriptor (releasing the kernel lock) and frees the -/// registration. +/// The lock file's name: `/.beyond-lock` for a directory, `.beyond-lock` for a file. +const RECORD_SUFFIX: &str = ".beyond-lock"; + +/// How long the legacy `flock` is retried before it counts as held. +const LEGACY_RETRY: Duration = Duration::from_millis(250); +const LEGACY_RETRY_INTERVAL: Duration = Duration::from_millis(2); + +/// What is being locked: a directory (its lock files live inside it), a file (they sit beside it), or +/// a file locked **itself** — a small file whose holder writes its content through the locked +/// descriptor ([`FileLock::file`]), as the MCP task journal key is made: lock it, re-read it, write it. +#[derive(Clone, Copy, Debug)] +pub enum Target<'a> { + Dir(&'a Path), + File(&'a Path), + Itself(&'a Path), +} + +impl Target<'_> { + /// The authoritative lock file. + fn record_path(&self) -> PathBuf { + match self { + Target::Dir(dir) => dir.join(RECORD_SUFFIX), + Target::File(file) => with_suffix(file, RECORD_SUFFIX), + Target::Itself(file) => file.to_path_buf(), + } + } + + /// The legacy `flock` file older binaries lock. None for a file locked itself: no released + /// binary locked one (the journal key's lock is newer than any of them), and an `flock` on the + /// same file would conflict with this process's own OFD lock on NFS, where `flock` is emulated + /// with a process-owned POSIX lock. + fn legacy_path(&self) -> Option { + match self { + Target::Dir(dir) => Some(dir.join("lock")), + Target::File(file) => Some(with_suffix(file, ".lock")), + Target::Itself(_) => None, + } + } +} + +fn with_suffix(file: &Path, suffix: &str) -> PathBuf { + let mut name: OsString = file.as_os_str().to_os_string(); + name.push(suffix); + PathBuf::from(name) +} + +/// Whether `path` names a lock file — for copying code that should leave locks behind (a worktree +/// seed: a copy of a lock file is not a lock), and for an atomic write, which must never rename a +/// fresh file over one (the lock would stay on the old, now unnamed, inode). +pub fn is_record_lock_file(path: &Path) -> bool { + path.file_name() + .and_then(OsStr::to_str) + .is_some_and(|n| n.ends_with(RECORD_SUFFIX)) +} + +/// Whether `path` is one of this module's lock files: a record lock file, or a legacy one — `lock` / +/// `.lock` — that sits beside a record lock file (so an unrelated `Cargo.lock` is not one). An +/// atomic write must never rename a fresh file over either: the lock would stay on the old inode, +/// which the path no longer names. +pub fn is_lock_file(path: &Path) -> bool { + if is_record_lock_file(path) { + return true; + } + let Some(name) = path.file_name().and_then(OsStr::to_str) else { + return false; + }; + let record = if name == "lock" { + path.with_file_name(RECORD_SUFFIX) + } else if let Some(stem) = name.strip_suffix(".lock") { + path.with_file_name(format!("{stem}{RECORD_SUFFIX}")) + } else { + return false; + }; + record.exists() +} + +/// Whether `dir` holds nothing but its own lock files — a lock directory that never became anything +/// (a session directory whose start failed before its first segment). +pub fn dir_holds_only_lock_files(dir: &Path) -> bool { + let record = Target::Dir(dir).record_path(); + let legacy = Target::Dir(dir).legacy_path().unwrap_or_default(); + std::fs::read_dir(dir).is_ok_and(|entries| { + entries.flatten().all(|e| { + let p = e.path(); + p == record || p == legacy + }) + }) +} + +/// A held lock. Dropping it unlocks each descriptor explicitly (so a forked child's copy of the +/// description cannot keep it), then closes them. pub struct FileLock { - file: File, - _registration: Registration, + record: File, + legacy: Option, + record_path: PathBuf, + legacy_path: Option, } 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. + /// The locked file, for a holder of a [`Target::Itself`] lock that writes the file's own content. + /// (Another descriptor would not release the lock — it belongs to this one's description — but + /// writing through the descriptor that holds it keeps "re-read, then write" one step.) pub fn file(&self) -> &File { - &self.file + &self.record } -} -/// "This process holds the lock at this path." -struct Registration(PathBuf); + /// Release the lock and unlink its files — for a lock directory being taken back. Unlocked and + /// closed *first*: an NFS client cannot unlink a file it still holds open (it silly-renames it to + /// `.nfs*`, which then keeps the directory from being removed). + pub fn release_and_remove_files(self) { + let (record_path, legacy_path) = (self.record_path.clone(), self.legacy_path.clone()); + drop(self); + let _ = std::fs::remove_file(&record_path); + if let Some(legacy_path) = &legacy_path { + let _ = std::fs::remove_file(legacy_path); + } + } +} -impl Drop for Registration { +impl Drop for FileLock { fn drop(&mut self) { - held() - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .remove(&self.0); + unlock(&self.record); + if let Some(legacy) = &self.legacy { + let _ = legacy.unlock(); + } } } -fn held() -> &'static Mutex> { - static HELD: OnceLock>> = OnceLock::new(); - HELD.get_or_init(|| Mutex::new(HashSet::new())) +/// Take the lock on `target`, creating its lock files (mode 0600) if needed; the directory a lock +/// file goes in must exist. `Ok(None)` means someone — another process, or another acquire in this +/// one — holds it *right now*. +/// +/// Each descriptor is opened read+write, because a byte-range lock needs a writable one. After +/// locking, `fstat` on the held descriptor is compared with `stat` of the path, so a file renamed away +/// between the open and the lock is caught rather than silently "locked". +pub fn try_lock(target: Target<'_>) -> std::io::Result> { + let record_path = target.record_path(); + let legacy_path = target.legacy_path(); + for _ in 0..RETRIES { + let record = open_lock_file(&record_path)?; + #[cfg(test)] + tests::between_open_and_lock(&record_path); + if !try_lock_description(&record)? { + return Ok(None); + } + if !same_file(&record, &record_path)? { + // The file we locked is no longer the file at that path. Let it go and look again. + unlock(&record); + continue; + } + let legacy = match legacy_path.as_deref().map(take_legacy_flock).transpose()? { + None => None, + Some(Legacy::Taken(file)) => Some(file), + // An old binary holds it: held. + Some(Legacy::Held) => { + unlock(&record); + return Ok(None); + } + Some(Legacy::Moved) => { + unlock(&record); + continue; + } + }; + return Ok(Some(FileLock { + record, + legacy, + record_path, + legacy_path, + })); + } + Ok(None) } -/// Take the lock on `lock_path`, creating the file (mode 0600) if needed. `Ok(None)` means someone — -/// another process, or this one — holds it *right now*. -/// -/// The descriptor is opened read+write, because NFS emulates `flock` with POSIX record locks and those -/// need a writable descriptor. After locking, `fstat` on the held descriptor is compared with `stat` -/// of the path, so a file renamed away between the open and the lock is caught rather than silently -/// "locked". -pub fn try_lock(lock_path: &Path) -> std::io::Result> { +fn open_lock_file(path: &Path) -> std::io::Result { + let mut opts = OpenOptions::new(); + opts.read(true).write(true).create(true).truncate(false); + #[cfg(unix)] { - let mut held = held() - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - if !held.insert(lock_path.to_path_buf()) { - return Ok(None); + use std::os::unix::fs::OpenOptionsExt; + opts.mode(0o600); + } + opts.open(path) +} + +#[cfg(any(target_os = "linux", target_os = "android"))] +fn whole_file(l_type: libc::c_int) -> libc::flock { + libc::flock { + l_type: l_type as libc::c_short, + l_whence: libc::SEEK_SET as libc::c_short, + l_start: 0, + l_len: 0, + // Must be 0 for an OFD lock. + l_pid: 0, + } +} + +/// An exclusive, non-blocking, whole-file lock owned by `file`'s open file description; `Ok(false)` +/// when another description (in any process, this one included) holds it. +fn try_lock_description(file: &File) -> std::io::Result { + #[cfg(any(target_os = "linux", target_os = "android"))] + { + let lock = whole_file(libc::F_WRLCK); + match nix::fcntl::fcntl(file, nix::fcntl::FcntlArg::F_OFD_SETLK(&lock)) { + Ok(_) => Ok(true), + Err(nix::errno::Errno::EAGAIN | nix::errno::Errno::EACCES) => Ok(false), + Err(e) => Err(std::io::Error::from(e)), } } - // Registered from here on, so every path out of this function frees the entry. - let registration = Registration(lock_path.to_path_buf()); - for _ in 0..RETRIES { - let mut opts = OpenOptions::new(); - opts.read(true).write(true).create(true).truncate(false); - #[cfg(unix)] - { - use std::os::unix::fs::OpenOptionsExt; - opts.mode(0o600); + #[cfg(not(any(target_os = "linux", target_os = "android")))] + { + match file.try_lock() { + Ok(()) => Ok(true), + Err(std::fs::TryLockError::WouldBlock) => Ok(false), + Err(std::fs::TryLockError::Error(e)) => Err(e), } - let file = opts.open(lock_path)?; - #[cfg(test)] - tests::between_open_and_lock(lock_path); + } +} + +/// Release `file`'s lock explicitly — not by closing, which a forked child's copy of the description +/// would outlive. +fn unlock(file: &File) { + #[cfg(any(target_os = "linux", target_os = "android"))] + { + let lock = whole_file(libc::F_UNLCK); + let _ = nix::fcntl::fcntl(file, nix::fcntl::FcntlArg::F_OFD_SETLK(&lock)); + } + #[cfg(not(any(target_os = "linux", target_os = "android")))] + { + let _ = file.unlock(); + } +} + +enum Legacy { + Taken(File), + Held, + /// The file at the path was replaced mid-acquire: look again from the top. + Moved, +} + +/// The legacy `flock` older binaries take (see the module doc), retried for [`LEGACY_RETRY`]. +fn take_legacy_flock(path: &Path) -> std::io::Result { + let file = open_lock_file(path)?; + let deadline = Instant::now() + LEGACY_RETRY; + loop { match file.try_lock() { - Ok(()) => {} - Err(std::fs::TryLockError::WouldBlock) => return Ok(None), + Ok(()) => break, + Err(std::fs::TryLockError::WouldBlock) if Instant::now() < deadline => { + std::thread::sleep(LEGACY_RETRY_INTERVAL); + } + Err(std::fs::TryLockError::WouldBlock) => return Ok(Legacy::Held), Err(std::fs::TryLockError::Error(e)) => return Err(e), } - if same_file(&file, lock_path)? { - return Ok(Some(FileLock { - file, - _registration: registration, - })); - } - // The file we locked is no longer the file at that path. Drop it and look again. - drop(file); } - Ok(None) + if !same_file(&file, path)? { + let _ = file.unlock(); + return Ok(Legacy::Moved); + } + Ok(Legacy::Taken(file)) } /// Whether the open descriptor and `path` still name the same inode. @@ -117,11 +327,13 @@ fn same_file(file: &File, path: &Path) -> std::io::Result { } #[cfg(test)] -mod tests { +pub(crate) mod tests { use super::*; use std::cell::RefCell; + use std::io::{BufRead, BufReader}; + use std::process::{Command, Stdio}; - /// Run (on this thread) after `try_lock` opens the file and before it locks it. + /// Run (on this thread) after `try_lock` opens the lock file and before it locks it. type Hook = Box; thread_local! { @@ -136,15 +348,39 @@ mod tests { }); } - /// The file at the path replaced between the open and the lock — a directory holding it renamed - /// away, a crashed holder's file swapped for a fresh one — must not leave the caller "holding" a - /// lock on a file nobody else will ever open. `try_lock` sees the inode moved and starts over, and - /// the lock it returns is on the file now at the path. + /// The lock file a target's lock lives on (for a test that links to it). + pub(crate) fn record_path(target: Target<'_>) -> PathBuf { + target.record_path() + } + + /// What another process sees: `Some(true)` if it cannot take a POSIX record lock on `target`'s + /// lock file right now (which an OFD lock held anywhere else conflicts with), `Some(false)` if it + /// can (and lets go at once). `None` without `python3`. + pub(crate) fn held_for_another_process(target: Target<'_>) -> Option { + let out = Command::new("python3") + .args([ + "-c", + "import fcntl,sys\nf=open(sys.argv[1],'r+')\ntry:\n fcntl.lockf(f,fcntl.LOCK_EX|fcntl.LOCK_NB)\n print('free')\nexcept OSError:\n print('held')", + ]) + .arg(target.record_path()) + .output() + .ok()?; + match String::from_utf8_lossy(&out.stdout).trim() { + "held" => Some(true), + "free" => Some(false), + _ => None, + } + } + + /// The file at the path replaced between the open and the lock must not leave the caller + /// "holding" a lock on a file nobody else will ever open: `try_lock` sees the inode moved and + /// starts over, and the lock it returns is on the file now at the path. #[cfg(unix)] #[test] fn a_lock_file_replaced_between_open_and_lock_is_not_the_one_held() { let dir = tempfile::tempdir().unwrap(); - let path = dir.path().join("x.lock"); + let target_file = dir.path().join("x"); + let target = Target::File(&target_file); let opens = std::rc::Rc::new(std::cell::Cell::new(0)); let counted = opens.clone(); let fresh = dir.path().join("fresh"); @@ -157,83 +393,205 @@ mod tests { } })); }); - let held = try_lock(&path).unwrap().expect("taken on the second look"); + let held = try_lock(target).unwrap().expect("taken on the second look"); BETWEEN.with(|hook| *hook.borrow_mut() = None); assert_eq!(opens.get(), 2, "the moved inode sent it round again"); - // The file at the path is the locked one: another open file description cannot lock it. - let other = std::fs::OpenOptions::new() - .read(true) - .write(true) - .open(&path) - .unwrap(); + if let Some(seen) = held_for_another_process(target) { + assert!(seen, "the lock held is on the file now at the path"); + } + drop(held); + } + + /// One holder at a time — a second acquire in the same process is refused by the kernel (two open + /// file descriptions conflict) — and the lock is free again once the holder drops it. + #[test] + fn one_holder_at_a_time_in_this_process_and_free_again_on_drop() { + let dir = tempfile::tempdir().unwrap(); + let file = dir.path().join("x"); + let held = try_lock(Target::File(&file)) + .unwrap() + .expect("free at first"); assert!( - matches!(other.try_lock(), Err(std::fs::TryLockError::WouldBlock)), - "the lock held is on the file now at the path" + try_lock(Target::File(&file)).unwrap().is_none(), + "a second acquire in this process is refused" ); drop(held); + assert!( + try_lock(Target::File(&file)).unwrap().is_some(), + "free once dropped" + ); } - /// One holder at a time within the process (the registry), and the lock is free again once the - /// holder drops it. + /// Different targets are different locks. #[test] - fn one_holder_at_a_time_and_free_again_on_drop() { + fn locks_on_different_targets_are_independent() { let dir = tempfile::tempdir().unwrap(); - let path = dir.path().join("x.lock"); - let held = try_lock(&path).unwrap().expect("free at first"); - assert!(try_lock(&path).unwrap().is_none(), "held by this process"); + let _a = try_lock(Target::File(&dir.path().join("a"))) + .unwrap() + .unwrap(); + assert!( + try_lock(Target::File(&dir.path().join("b"))) + .unwrap() + .is_some() + ); + } + + /// Opening and closing the lock file through another descriptor — by its own name, a symlink or a + /// hard link — neither releases the lock nor takes it: it belongs to the description that locked + /// it. (With a process-owned record lock, any such close dropped it.) + #[cfg(unix)] + #[test] + fn another_descriptor_of_the_lock_file_neither_releases_nor_shares_it() { + let dir = tempfile::tempdir().unwrap(); + let session = dir.path().join("s1"); + std::fs::create_dir_all(&session).unwrap(); + let target = Target::Dir(&session); + let held = try_lock(target).unwrap().unwrap(); + let symlink = dir.path().join("innocent.txt"); + std::os::unix::fs::symlink(target.record_path(), &symlink).unwrap(); + let hardlink = dir.path().join("also-innocent.txt"); + std::fs::hard_link(target.record_path(), &hardlink).unwrap(); + for path in [&target.record_path(), &symlink, &hardlink] { + drop(std::fs::read(path).unwrap()); + drop( + OpenOptions::new() + .read(true) + .write(true) + .open(path) + .unwrap(), + ); + } + assert!( + try_lock(target).unwrap().is_none(), + "still held against this process" + ); + if let Some(seen) = held_for_another_process(target) { + assert!(seen, "and against another"); + } drop(held); - assert!(try_lock(&path).unwrap().is_some(), "free once dropped"); + if let Some(seen) = held_for_another_process(target) { + assert!(!seen, "released when the holder lets go"); + } } - /// A path this process holds is refused to it a second time without opening the file — so even - /// when the file at the path is replaced underneath the holder, this process never ends up - /// holding the path twice. (On NFS the kernel lock is a POSIX record lock, which closing *any* - /// descriptor drops; the registry is what keeps a second attempt from opening one.) + /// The unlock is explicit, so a forked child that still holds a copy of the description (between + /// its fork and its exec) cannot keep the lock alive after its holder lets go. + #[cfg(unix)] #[test] - fn a_path_held_here_is_refused_here_even_if_its_file_is_replaced() { + fn a_released_lock_is_free_while_a_child_still_holds_the_description() { let dir = tempfile::tempdir().unwrap(); - let path = dir.path().join("x.lock"); - let _held = try_lock(&path).unwrap().expect("free at first"); - std::fs::remove_file(&path).unwrap(); + let file = dir.path().join("x"); + let target = Target::File(&file); + let held = try_lock(target).unwrap().unwrap(); + // A child holding a copy of the lock's description for the whole test — the fork→exec + // window, made long: a `dup` (which drops close-on-exec) survives into it. + let inherited = nix::unistd::dup(&held.record).unwrap(); + let Ok(mut child) = Command::new("sh") + .args(["-c", "read _"]) + .stdin(Stdio::piped()) + .spawn() + else { + return; + }; + drop(inherited); + drop(held); + let retaken = try_lock(target).unwrap(); + drop(child.stdin.take()); + let _ = child.wait(); assert!( - try_lock(&path).unwrap().is_none(), - "this process already holds the path" + retaken.is_some(), + "the explicit unlock released it despite the child's copy" ); + } + + /// A file locked *itself* (the MCP task journal key): one holder at a time, the holder writes + /// through the locked descriptor, and no lock file appears beside it. + #[cfg(unix)] + #[test] + fn a_file_locked_itself_is_written_through_its_lock() { + use std::io::{Seek as _, Write as _}; + let dir = tempfile::tempdir().unwrap(); + let key = dir.path().join("journal.key"); + let held = try_lock(Target::Itself(&key)).unwrap().unwrap(); assert!( - !path.exists(), - "the refused attempt never opened (and so created) the file" + try_lock(Target::Itself(&key)).unwrap().is_none(), + "one holder at a time" ); + let mut file = held.file(); + file.set_len(0).unwrap(); + file.seek(std::io::SeekFrom::Start(0)).unwrap(); + file.write_all(b"secret").unwrap(); + assert_eq!(std::fs::read(&key).unwrap(), b"secret"); + let entries: Vec<_> = std::fs::read_dir(dir.path()) + .unwrap() + .flatten() + .map(|e| e.file_name()) + .collect(); + assert_eq!(entries, ["journal.key"], "no lock file beside it"); + if let Some(seen) = held_for_another_process(Target::Itself(&key)) { + assert!(seen, "held for another process"); + } + drop(held); + assert!(try_lock(Target::Itself(&key)).unwrap().is_some()); } - /// Different paths are different locks. + /// A legacy lock file is recognized only beside a record lock file, so an unrelated `Cargo.lock` + /// is just a file. #[test] - fn locks_on_different_paths_are_independent() { + fn a_legacy_lock_file_is_one_only_beside_a_record_lock_file() { let dir = tempfile::tempdir().unwrap(); - let _a = try_lock(&dir.path().join("a.lock")).unwrap().unwrap(); - assert!(try_lock(&dir.path().join("b.lock")).unwrap().is_some()); + let session = dir.path().join("s1"); + std::fs::create_dir_all(&session).unwrap(); + let file = dir.path().join("mcp-manifest.json"); + let _a = try_lock(Target::Dir(&session)).unwrap().unwrap(); + let _b = try_lock(Target::File(&file)).unwrap().unwrap(); + assert!(is_lock_file(&session.join("lock"))); + assert!(is_lock_file(&dir.path().join("mcp-manifest.json.lock"))); + std::fs::write(dir.path().join("Cargo.lock"), "").unwrap(); + assert!(!is_lock_file(&dir.path().join("Cargo.lock"))); + assert!(!is_lock_file(&dir.path().join("lock"))); } + /// An old binary held only the legacy `flock`. While one runs, it must still exclude this binary. #[cfg(unix)] #[test] - fn same_file_catches_a_path_whose_inode_moved() { + fn an_old_binary_holding_only_the_legacy_flock_excludes_this_one() { let dir = tempfile::tempdir().unwrap(); - let path = dir.path().join("lock"); - std::fs::write(&path, b"").unwrap(); - let held = File::open(&path).unwrap(); - assert!(same_file(&held, &path).unwrap()); - - // Another file takes the path — what a session directory renamed into `.trash/` looks like - // from the lock's point of view. - let other = dir.path().join("other"); - std::fs::write(&other, b"").unwrap(); - std::fs::rename(&other, &path).unwrap(); + let session = dir.path().join("s1"); + std::fs::create_dir_all(&session).unwrap(); + let target = Target::Dir(&session); + let Ok(mut old) = Command::new("python3") + .args([ + "-c", + "import fcntl,sys\nf=open(sys.argv[1],'a+')\nfcntl.flock(f,fcntl.LOCK_EX)\nprint('locked',flush=True)\nsys.stdin.read()", + ]) + .arg(target.legacy_path().unwrap()) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .spawn() + else { + eprintln!("no python3: skipping the mixed-rollout check"); + return; + }; + let mut line = String::new(); + BufReader::new(old.stdout.take().unwrap()) + .read_line(&mut line) + .unwrap(); + assert_eq!(line.trim(), "locked"); + let started = Instant::now(); assert!( - !same_file(&held, &path).unwrap(), - "the descriptor no longer names the file at that path" + try_lock(target).unwrap().is_none(), + "an old binary's flock still means held" + ); + assert!( + started.elapsed() < Duration::from_secs(2), + "and it says so within the bounded retry" + ); + drop(old.stdin.take()); + old.wait().unwrap(); + assert!( + try_lock(target).unwrap().is_some(), + "free once the old binary lets go" ); - - // And a path that is simply gone is not the same file either. - std::fs::remove_file(&path).unwrap(); - assert!(!same_file(&held, &path).unwrap()); } } diff --git a/crates/agent/src/serve_ws.rs b/crates/agent/src/serve_ws.rs index 048f247d..5e2c920d 100644 --- a/crates/agent/src/serve_ws.rs +++ b/crates/agent/src/serve_ws.rs @@ -590,9 +590,9 @@ where /// a replica that died mid-start) is left for an out-of-band sweep; see ARCHITECTURE.md. /// /// `remove_dir` *is* the emptiness check, and an atomic one: `000001.jsonl` is created with `O_EXCL` -/// before anything else and never deleted, so a directory holding nothing but `lock` has never been +/// before anything else and never deleted, so a directory holding nothing but its lock files has never been /// a session — and if a racing replica got one written in between, `ENOTEMPTY` leaves everything -/// alone. The `lock` file is unlinked first because `remove_dir` would otherwise always fail; the +/// alone. The lock files are unlinked first because `remove_dir` would otherwise always fail; the /// window that opens between that unlink and this task's own release is harmless, since the lock is /// liveness-only (correctness is the epoch fence) and this replica is already done with the session. async fn release_session_lock(lock: Option, path: Option) { @@ -601,16 +601,15 @@ async fn release_session_lock(lock: Option, path: Option 1, + "so must the lock files it is locked through" ); // The lock itself is released, so the next owner can take it. assert!(acquire_session_lock(&path).unwrap().is_some()); diff --git a/crates/agent/src/session_store.rs b/crates/agent/src/session_store.rs index 6a8bec69..e5205e2d 100644 --- a/crates/agent/src/session_store.rs +++ b/crates/agent/src/session_store.rs @@ -862,7 +862,7 @@ fn make_journal_key( } let deadline = std::time::Instant::now() + JOURNAL_KEY_LOCK_WAIT; let lock = loop { - match crate::file_lock::try_lock(&path) { + match crate::file_lock::try_lock(crate::file_lock::Target::Itself(&path)) { Ok(Some(lock)) => break lock, Ok(None) if std::time::Instant::now() < deadline => { std::thread::sleep(std::time::Duration::from_millis(2)); @@ -6402,24 +6402,18 @@ pub type SessionLock = crate::file_lock::FileLock; /// more. Correctness comes from the epoch fence, which needs no lock at all — so a lock lost to a /// crash, a network partition or a stuck NFS client costs a retry, never history. /// -/// Three details that matter on EFS: -/// - the descriptor is opened **read+write**, because NFS emulates `flock` with POSIX record locks and -/// those need a writable descriptor; -/// - after locking, `fstat` on the held descriptor is compared with `stat` of the path, so a directory -/// renamed into `.trash/` between the open and the lock is caught rather than silently "locked"; -/// - a session is opened **once per process**, because closing *any* descriptor to a POSIX-locked file -/// drops the lock — a second open in the same process would quietly release the first one's hold. -/// -/// All three are [`crate::file_lock::try_lock`]'s; this only names the session's lock file. +/// It is an open file description lock ([`crate::file_lock`]): no other descriptor of the lock file +/// releases it, a second acquire in this process is refused, and it is released with an explicit +/// unlock — a forked child inherits the description, but cannot keep the lock once its holder lets +/// go. After locking, `fstat` on the held descriptor is compared with `stat` of the path, so a +/// directory renamed into `.trash/` between the open and the lock is caught rather than silently +/// "locked". This only says which lock file is the session's. pub fn acquire_session_lock(session_path: &Path) -> std::io::Result> { - let lock_path = if session_path.is_dir() { - session_path.join("lock") + crate::file_lock::try_lock(if session_path.is_dir() { + crate::file_lock::Target::Dir(session_path) } else { - let mut name = session_path.as_os_str().to_os_string(); - name.push(".lock"); - PathBuf::from(name) - }; - crate::file_lock::try_lock(&lock_path) + crate::file_lock::Target::File(session_path) + }) } #[cfg(test)] @@ -6427,6 +6421,175 @@ mod tests { use super::*; use serde_json::json; + /// The root cause of `serve_ws::tests::a_session_directory_with_a_segment_is_left_alone`'s + /// flakes under load: a child forked by *another* thread (a tool, an MCP server) shares the + /// lock's open file description until it execs, so a lock released by closing its descriptor was + /// still held about 6% of the time while processes were being spawned (180/3000 measured). The + /// OFD lock is inherited the same way, but it is released with an explicit unlock, which drops it + /// whatever copies a child holds: zero. + #[test] + fn a_released_session_lock_is_free_at_once_while_other_threads_spawn_processes() { + let stop = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)); + let spawners: Vec<_> = (0..8) + .map(|_| { + let stop = stop.clone(); + std::thread::spawn(move || { + while !stop.load(std::sync::atomic::Ordering::Relaxed) { + let _ = std::process::Command::new("true").status(); + } + }) + }) + .collect(); + let root = tempfile::tempdir().unwrap(); + let mut held = 0; + for i in 0..3000 { + let p = root.path().join(format!("s{i}")); + std::fs::create_dir_all(&p).unwrap(); + let lock = acquire_session_lock(&p).unwrap().unwrap(); + drop(lock); + if acquire_session_lock(&p).unwrap().is_none() { + held += 1; + } + } + stop.store(true, std::sync::atomic::Ordering::Relaxed); + for s in spawners { + s.join().unwrap(); + } + assert_eq!( + held, 0, + "a released lock was still held by a forked child {held}/3000 times" + ); + } + + /// The lock reached through another name — a symlink (`innocent.txt -> /`, + /// which the model can make with `bash`) or a hard link — by every in-process reader that does + /// not know it is a lock file: the file tools' read, write, edit and search, the memory tool's + /// view (under `/memory`), and the `@file` read (`std::fs::read_to_string`). Each opens and + /// closes the lock file on its own descriptor, and none of them releases the lock, which belongs + /// to the open file description that took it. + #[cfg(unix)] + #[tokio::test] + async fn every_in_process_reader_reaching_the_lock_file_leaves_it_held() { + use agent_core::Tool as _; + let dir = tmpdir(); + let session_dir = dir.path().join("s1"); + std::fs::create_dir_all(session_dir.join("memory")).unwrap(); + let lock = acquire_session_lock(&session_dir).unwrap().unwrap(); + let target = crate::file_lock::Target::Dir(&session_dir); + let Some(true) = crate::file_lock::tests::held_for_another_process(target) else { + eprintln!("no python3: skipping"); + return; + }; + let record = crate::file_lock::tests::record_path(target); + let work = dir.path().join("work"); + std::fs::create_dir_all(&work).unwrap(); + let symlink = work.join("innocent.txt"); + std::os::unix::fs::symlink(&record, &symlink).unwrap(); + let hardlink = work.join("also-innocent.txt"); + std::fs::hard_link(&record, &hardlink).unwrap(); + std::os::unix::fs::symlink(&record, session_dir.join("memory").join("notes.md")).unwrap(); + + for link in [&symlink, &hardlink] { + let p = link.to_str().unwrap(); + let _ = crate::tools::read::Read::new(&work) + .run(serde_json::json!({ "path": p })) + .await; + let _ = crate::tools::write::Write::new(&work) + .run(serde_json::json!({ "path": p, "content": "scribble" })) + .await; + let _ = crate::tools::edit::Edit::new(&work) + .run(serde_json::json!({ "path": p, "old_string": "", "new_string": "x" })) + .await; + // What `@file` does with the path it is given. + let _ = std::fs::read_to_string(link); + } + let _ = crate::tools::grep::Grep::new(&work) + .run(serde_json::json!({ "pattern": "x", "path": work.to_str().unwrap() })) + .await; + let memory = crate::tools::memory::Memory::new(std::sync::Arc::new( + crate::memory::file::FileBackend::session_at(session_dir.join("memory")), + )); + let _ = memory + .run(serde_json::json!({ "command": "view", "path": "/session/notes.md" })) + .await; + + assert_eq!( + crate::file_lock::tests::held_for_another_process(target), + Some(true), + "no reader of the lock file released it" + ); + assert!( + acquire_session_lock(&session_dir).unwrap().is_none(), + "nor can this process take it a second time" + ); + assert_eq!( + std::fs::metadata(&record).unwrap().len(), + 0, + "and nothing replaced or wrote into the lock file" + ); + drop(lock); + assert_eq!( + crate::file_lock::tests::held_for_another_process(target), + Some(false) + ); + } + + /// Every public path that reads a locked session directory — listing, opening, forking, the + /// file tools' read and search of every file in it, the lock file included — leaves the lock + /// held for another process: none of their descriptors is the one that holds it. + #[tokio::test] + async fn reading_a_locked_session_through_any_public_api_leaves_it_locked() { + use crate::tools::fs::FsBackend as _; + let dir = tmpdir(); + let repo = SessionRepo::open_with( + dir.path(), + RepoOptions { + layout: Layout::Segmented { codec: None }, + ..RepoOptions::default() + }, + ) + .unwrap(); + let mut store = repo.create(SessionMeta::new("/w", "m")).unwrap(); + let mut session = Session::new(); + session.user("hello lock"); + store.append_new(&session.messages).unwrap(); + let id = store.meta().id.clone(); + let session_dir = store.path().to_path_buf(); + drop(store); + let lock = acquire_session_lock(&session_dir).unwrap().unwrap(); + let target = crate::file_lock::Target::Dir(&session_dir); + let Some(true) = crate::file_lock::tests::held_for_another_process(target) else { + eprintln!("no python3: skipping"); + return; + }; + + let _ = repo.list().unwrap(); + let _ = repo.open_id_read_only(&id).unwrap(); + let _ = repo.fork(&id, usize::MAX).unwrap(); + let fs = crate::tools::fs::local::LocalFs::new(); + for entry in std::fs::read_dir(&session_dir).unwrap().flatten() { + // Read every file there by the tool path; a lock file is refused, never opened. + let _ = fs.read_bytes(&entry.path(), 0, 1024).await; + } + let _ = fs + .search(&crate::tools::fs::local::query( + "lock", + dir.path().to_path_buf(), + 100, + )) + .await; + assert_eq!( + crate::file_lock::tests::held_for_another_process(target), + Some(true), + "no public read of the session released its lock" + ); + drop(lock); + assert_eq!( + crate::file_lock::tests::held_for_another_process(target), + Some(false) + ); + } + fn tmpdir() -> tempfile::TempDir { tempfile::tempdir().unwrap() } diff --git a/crates/agent/src/tools/mcp.rs b/crates/agent/src/tools/mcp.rs index a68de292..c0ea5779 100644 --- a/crates/agent/src/tools/mcp.rs +++ b/crates/agent/src/tools/mcp.rs @@ -148,6 +148,9 @@ struct HttpDial { /// The daemon's `mcp_http` client: ALPN, the `web` tool's SSRF resolver, and deliberately not /// the gateway's h2c pool — a tenant's connector is not the gateway. client: reqwest::Client, + /// The egress policy that client enforces, for the fresh client that ends this connection's + /// session on the way out of the process (`mcp_http_exit`). + policy: Arc, headers: Vec, } @@ -453,6 +456,9 @@ impl ClientHandler for McpHandler { } self.skills_changed .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + // rmcp runs this on a task of its own, after (or before) the response that followed the + // notification on the wire; this line is how an operator — or a test — knows it has run. + tracing::debug!(server = %self.server_name, "resources/list_changed handled: view cache cleared"); } async fn on_progress( @@ -2112,6 +2118,7 @@ fn granted_jobs( host: host.clone(), http: Some(HttpDial { client: egress.client.clone(), + policy: egress.policy.clone(), headers, }), apps, @@ -2511,13 +2518,21 @@ async fn connect_http( // Wrapped so extension results survive rmcp's result decoding — see `mcp_wire` — so an MCP // App view's read is refused over its cap without being read whole — see `mcp_view_http` — and // so a rejected OAuth token is refreshed and the request retried once — see `mcp_oauth`. + // The session this connection opens is ended on the way out of the process if nothing ended + // it before — see `mcp_http_exit`. + let exit = crate::tools::mcp_http_exit::HttpSession::new( + url, + auth.clone(), + dial.http.as_ref().map(|h| h.policy.clone()), + ); let transport = StreamableHttpClientTransport::with_client( crate::tools::mcp_oauth::OAuthHttp::new( crate::tools::mcp_view_http::ViewCappedHttp::new(crate::tools::mcp_wire::HttpClient { client, }), auth, - ), + ) + .ending_session_on_exit(exit), transport_config, ); @@ -3346,6 +3361,7 @@ mod tests { agent_core::ensure_provider(); let dial = HttpDial { client: reqwest::Client::new(), + policy: Arc::default(), headers: vec![header], }; assert!(!format!("{dial:?}").contains("tenant-token")); diff --git a/crates/agent/src/tools/mcp_http_exit.rs b/crates/agent/src/tools/mcp_http_exit.rs new file mode 100644 index 00000000..3303c9a3 --- /dev/null +++ b/crates/agent/src/tools/mcp_http_exit.rs @@ -0,0 +1,422 @@ +//! Ending streamable-HTTP MCP sessions on the way out of the process. +//! +//! A server that issued an `Mcp-Session-Id` keeps that session until it is told it is over — the +//! spec's `DELETE` — or until it times out on its own. rmcp sends that `DELETE` when a connection is +//! closed while the runtime is running (the idle reaper, a registry rebuild); a process that simply +//! exits (`agent run` finishing, a `serve` shutdown) never got that far, and every session it held +//! was left for the server to expire. +//! +//! So each connection records the session it opened in an [`HttpSession`] (kept by +//! [`OAuthHttp`](crate::tools::mcp_oauth::OAuthHttp), which sees every response's session id and +//! every `DELETE` rmcp sends), and [`begin_close_all`] — run beside the stdio sweep in +//! `mcp_stdio::sweep_before_exit` — sends a `DELETE` for every session still open: in parallel, with +//! the server's bearer token and the connection's own headers, under one bounded deadline. +//! +//! It runs on a thread of its own with a fresh runtime and a fresh, pool-less client: the process +//! runtime may already be gone (`main`'s last guard) or blocked in the very call that is exiting, +//! and a pooled connection belonging to a blocked runtime would never answer. A grant connector's +//! client gets the same SSRF resolver its connection used. A server that does not answer costs at +//! most the deadline; the session then expires on the server's side, as before. + +use std::collections::HashMap; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; + +use http::{HeaderName, HeaderValue}; + +/// The most the exit waits for every `DELETE` together. +pub const CLOSE_DEADLINE: Duration = Duration::from_millis(1500); + +/// Headers a `DELETE` never takes from the connection's (rmcp's reserved set, plus the two it sets). +const NOT_COPIED: [&str; 5] = [ + "accept", + "mcp-session-id", + "last-event-id", + "authorization", + "content-type", +]; + +/// One HTTP connection's open session, if it has one. See the module doc. +/// +/// Held by the connection's transport (`OAuthHttp`) and only **weakly** by the process: a connection +/// that goes away takes its entry with it, so a long-lived daemon never accumulates them. If it goes +/// away with its session still open (its transport died before rmcp's own `DELETE`), dropping it +/// sends that `DELETE` right away on the runtime it was dropped on, and keeps it as an [`Ending`] +/// for the exit only until that succeeds — or, during the runtime's own teardown, where nothing +/// spawned runs any more, for the exit to send. +pub(crate) struct HttpSession { + registry: Arc, + uri: String, + auth: Option>, + /// A grant connector's egress policy: its exit client resolves through the same SSRF resolver. + egress: Option>, + open: Mutex>, +} + +struct Open { + id: String, + headers: HashMap, +} + +/// One session's `DELETE`, everything needed to send it. +#[derive(Clone)] +struct Ending { + uri: String, + id: String, + headers: HashMap, + auth: Option>, + egress: Option>, +} + +/// Where the exit finds what it has to end. One per process ([`Registry::process`]); a test makes +/// its own, so tests never see each other's sessions. +pub(crate) struct Registry { + /// Connections with a session open, weakly. + open: Mutex>>, + /// Sessions whose connection went away still open, by a sequence number, while their `DELETE` + /// is in flight — or, if their runtime was already tearing down, until the exit sends it. + orphans: Mutex>, + /// How long an orphan's `DELETE` may take before it is given up. + orphan_timeout: Duration, +} + +impl Default for Registry { + fn default() -> Self { + Self { + open: Mutex::default(), + orphans: Mutex::default(), + orphan_timeout: CLOSE_DEADLINE * 4, + } + } +} + +impl Registry { + pub(crate) fn process() -> Arc { + static PROCESS: std::sync::LazyLock> = std::sync::LazyLock::new(Arc::default); + PROCESS.clone() + } +} + +fn lock(m: &Mutex) -> std::sync::MutexGuard<'_, T> { + m.lock().unwrap_or_else(std::sync::PoisonError::into_inner) +} + +impl HttpSession { + pub(crate) fn new( + uri: &str, + auth: Option>, + egress: Option>, + ) -> Arc { + Self::new_in(Registry::process(), uri, auth, egress) + } + + pub(crate) fn new_in( + registry: Arc, + uri: &str, + auth: Option>, + egress: Option>, + ) -> Arc { + Arc::new(Self { + registry, + uri: uri.to_owned(), + auth, + egress, + open: Mutex::new(None), + }) + } + + /// A response named the session `id` (the handshake's, or a re-initialized one's); `headers` + /// are the request's own, which a `DELETE` for it carries too. + pub(crate) fn opened(self: &Arc, id: &str, headers: &HashMap) { + let mut open = lock(&self.open); + if open.as_ref().is_some_and(|o| o.id == id) { + return; + } + *open = Some(Open { + id: id.to_owned(), + headers: headers.clone(), + }); + drop(open); + let mut all = lock(&self.registry.open); + all.retain(|s| s.strong_count() > 0); + if !all + .iter() + .any(|s| std::ptr::eq(s.as_ptr(), Arc::as_ptr(self))) + { + all.push(Arc::downgrade(self)); + } + } + + /// rmcp is ending the session `id` itself (its own `DELETE`): nothing left for the exit to do. + pub(crate) fn closed(self: &Arc, id: &str) { + let mut open = lock(&self.open); + if open.as_ref().is_some_and(|o| o.id == id) { + *open = None; + drop(open); + lock(&self.registry.open) + .retain(|s| s.strong_count() > 0 && !std::ptr::eq(s.as_ptr(), Arc::as_ptr(self))); + } + } + + fn ending(&self) -> Option { + let open = lock(&self.open); + let open = open.as_ref()?; + Some(Ending { + uri: self.uri.clone(), + id: open.id.clone(), + headers: open.headers.clone(), + auth: self.auth.clone(), + egress: self.egress.clone(), + }) + } +} + +impl Drop for HttpSession { + fn drop(&mut self) { + let Some(ending) = self.ending() else { + return; + }; + static SEQ: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); + let n = SEQ.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + lock(&self.registry.orphans).push((n, ending.clone())); + // Send it now if a runtime is here to run it. Once that attempt is over — answered, refused + // or timed out — the entry goes: an unanswered `DELETE` is not kept for the process's life. + // Only a drop during the runtime's own teardown, where nothing spawned runs any more, leaves + // it for the exit to send. + if let Ok(handle) = tokio::runtime::Handle::try_current() { + let registry = self.registry.clone(); + handle.spawn(async move { + let _ = tokio::time::timeout(registry.orphan_timeout, ending.send()).await; + lock(®istry.orphans).retain(|(m, _)| *m != n); + }); + } + } +} + +impl Ending { + /// Send this session's `DELETE`: the status, or why it was not sent. + async fn send(&self) -> Result { + let client = self.client()?; + let mut request = client + .delete(&self.uri) + .header("Mcp-Session-Id", self.id.as_str()); + for (name, value) in &self.headers { + if !NOT_COPIED.contains(&name.as_str()) { + request = request.header(name, value); + } + } + if let Some(auth) = &self.auth + && let Some(token) = auth.token().await + { + request = request.bearer_auth(token); + } + request + .send() + .await + .map(|r| r.status().as_u16()) + .map_err(|e| e.to_string()) + } + + fn client(&self) -> Result { + agent_core::ensure_provider(); + let mut builder = reqwest::Client::builder() + .pool_max_idle_per_host(0) + .redirect(reqwest::redirect::Policy::none()) + .connect_timeout(CLOSE_DEADLINE); + if let Some(policy) = &self.egress { + builder = builder.dns_resolver(crate::tools::web::ssrf::SsrfResolver::new( + (**policy).clone(), + )); + } + builder.build().map_err(|e| e.to_string()) + } +} + +/// Every session the exit still has to end: those of live connections, and those left by +/// connections that went away before their `DELETE` was answered. Taken: each is sent once. +fn take_endings(registry: &Registry) -> Vec { + let live: Vec> = std::mem::take(&mut *lock(®istry.open)) + .iter() + .filter_map(std::sync::Weak::upgrade) + .collect(); + let mut endings: Vec = live.iter().filter_map(|s| s.ending()).collect(); + // Ended here; a later drop of these connections must not send them again. + for session in &live { + *lock(&session.open) = None; + } + endings.extend( + std::mem::take(&mut *lock(®istry.orphans)) + .into_iter() + .map(|(_, e)| e), + ); + endings +} + +/// Start ending every session still open (see the module doc); the returned waiter blocks until +/// they are all answered or `deadline` has passed since this call, whichever is first. Start it, +/// do the rest of the exit, then wait — so the HTTP round trips overlap the stdio sweep. +pub fn begin_close_all(deadline: Duration) -> impl FnOnce() { + let endings = take_endings(&Registry::process()); + let until = Instant::now() + deadline; + let (done, finished) = std::sync::mpsc::channel::<()>(); + if !endings.is_empty() { + let spawned = std::thread::Builder::new() + .name("mcp-http-exit".into()) + .spawn(move || { + let Ok(runtime) = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + else { + return; + }; + runtime.block_on(async { + let sends = endings.iter().map(Ending::send); + let left = until.saturating_duration_since(Instant::now()); + match tokio::time::timeout(left, futures::future::join_all(sends)).await { + Ok(results) => { + for (ending, result) in endings.iter().zip(results) { + if let Err(e) = result { + tracing::debug!(uri = %ending.uri, error = %e, "ending an MCP session on exit failed"); + } + } + } + Err(_) => tracing::debug!( + "ending MCP sessions on exit ran out of time; the servers will expire them" + ), + } + }); + runtime.shutdown_timeout(Duration::from_millis(50)); + let _ = done.send(()); + }); + if spawned.is_err() { + return Box::new(|| {}) as Box; + } + } else { + let _ = done.send(()); + } + Box::new(move || { + let _ = finished.recv_timeout(until.saturating_duration_since(Instant::now())); + }) as Box +} + +#[cfg(test)] +mod tests { + use super::*; + + /// A loopback server that answers every request 200 and records each request's first line and + /// `Mcp-Session-Id`. + async fn server() -> (String, Arc>>) { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("http://{}/mcp", listener.local_addr().unwrap()); + let seen = Arc::new(Mutex::new(Vec::new())); + let record = seen.clone(); + tokio::spawn(async move { + while let Ok((mut stream, _)) = listener.accept().await { + let record = record.clone(); + tokio::spawn(async move { + use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _}; + let mut buf = vec![0u8; 8192]; + let n = stream.read(&mut buf).await.unwrap_or(0); + let head = String::from_utf8_lossy(&buf[..n]).to_ascii_lowercase(); + let line = head.lines().next().unwrap_or_default().to_owned(); + let session = head + .lines() + .find_map(|l| l.strip_prefix("mcp-session-id: ")) + .unwrap_or_default(); + record.lock().unwrap().push(format!("{line} {session}")); + let _ = stream + .write_all( + b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\nConnection: close\r\n\r\n", + ) + .await; + }); + } + }); + (url, seen) + } + + /// What `registry` holds: live connections with a session open, and orphaned `DELETE`s. + fn held(registry: &Registry) -> (usize, usize) { + let live = lock(®istry.open) + .iter() + .filter(|s| s.strong_count() > 0) + .count(); + (live, lock(®istry.orphans).len()) + } + + /// A connection that goes away with its session still open — its transport died before rmcp's + /// own `DELETE` — is not kept by the process until it exits: the `DELETE` goes out at once, and + /// once it is answered nothing of the connection is left. + #[tokio::test] + async fn a_connection_gone_with_its_session_open_ends_it_and_is_not_kept() { + agent_core::ensure_provider(); + let registry = Arc::new(Registry::default()); + let (url, seen) = server().await; + let session = HttpSession::new_in(registry.clone(), &url, None, None); + session.opened("sess-gone", &HashMap::new()); + assert_eq!(held(®istry), (1, 0)); + drop(session); + + let deadline = Instant::now() + Duration::from_secs(10); + while !seen.lock().unwrap().iter().any(|l| l.starts_with("delete")) { + assert!(Instant::now() < deadline, "the DELETE never went out"); + tokio::time::sleep(Duration::from_millis(10)).await; + } + assert_eq!( + seen.lock().unwrap().clone(), + ["delete /mcp http/1.1 sess-gone"] + ); + while held(®istry) != (0, 0) { + assert!(Instant::now() < deadline, "{:?}", held(®istry)); + tokio::time::sleep(Duration::from_millis(10)).await; + } + } + + /// A `DELETE` nobody answers is not kept either: once its attempt has timed out the entry goes, + /// rather than living as long as the process. + #[tokio::test] + async fn an_unanswered_orphan_delete_is_dropped_after_its_timeout() { + agent_core::ensure_provider(); + let registry = Arc::new(Registry { + orphan_timeout: Duration::from_millis(300), + ..Registry::default() + }); + // A listener that accepts and never answers. + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("http://{}/mcp", listener.local_addr().unwrap()); + let _accepting = tokio::spawn(async move { + let mut held = Vec::new(); + while let Ok((stream, _)) = listener.accept().await { + held.push(stream); + } + }); + let session = HttpSession::new_in(registry.clone(), &url, None, None); + session.opened("sess-silent", &HashMap::new()); + drop(session); + assert_eq!(held(®istry), (0, 1), "in flight"); + let deadline = Instant::now() + Duration::from_secs(10); + while held(®istry) != (0, 0) { + assert!( + Instant::now() < deadline, + "still kept: {:?}", + held(®istry) + ); + tokio::time::sleep(Duration::from_millis(20)).await; + } + } + + /// One rmcp ended itself is gone at once, and is not sent again on drop. + #[tokio::test] + async fn a_session_rmcp_ended_is_neither_kept_nor_sent_again() { + let registry = Arc::new(Registry::default()); + let (url, seen) = server().await; + let session = HttpSession::new_in(registry.clone(), &url, None, None); + session.opened("sess-closed", &HashMap::new()); + session.closed("sess-closed"); + assert_eq!(held(®istry), (0, 0)); + drop(session); + tokio::time::sleep(Duration::from_millis(200)).await; + assert!(seen.lock().unwrap().is_empty()); + assert_eq!(held(®istry), (0, 0)); + } +} diff --git a/crates/agent/src/tools/mcp_manifest.rs b/crates/agent/src/tools/mcp_manifest.rs index c5674ef7..30e55549 100644 --- a/crates/agent/src/tools/mcp_manifest.rs +++ b/crates/agent/src/tools/mcp_manifest.rs @@ -191,12 +191,9 @@ fn lock_store(dir: &ManifestDir) -> Option { if let Some(parent) = path.parent() { std::fs::create_dir_all(parent).ok()?; } - let mut lock_path = path.into_os_string(); - lock_path.push(".lock"); - let lock_path = PathBuf::from(lock_path); let deadline = Instant::now() + LOCK_TIMEOUT; loop { - match crate::file_lock::try_lock(&lock_path) { + match crate::file_lock::try_lock(crate::file_lock::Target::File(&path)) { Ok(Some(lock)) => return Some(lock), Ok(None) if Instant::now() < deadline => std::thread::sleep(LOCK_RETRY_INTERVAL), Ok(None) | Err(_) => return None, @@ -393,7 +390,11 @@ mod tests { let leftovers: Vec<_> = std::fs::read_dir(dir.path()) .unwrap() .map(|e| e.unwrap().file_name().to_string_lossy().into_owned()) - .filter(|n| n != FILE && *n != format!("{FILE}.lock")) + .filter(|n| { + n != FILE + && *n != format!("{FILE}.lock") + && !crate::file_lock::is_record_lock_file(std::path::Path::new(n)) + }) .collect(); assert!( leftovers.is_empty(), @@ -443,10 +444,10 @@ mod tests { let manifest = ManifestDir::at(dir.path()); let held = lock_store(&manifest).unwrap(); assert!( - crate::file_lock::try_lock(&dir.path().join(format!("{FILE}.lock"))) + crate::file_lock::try_lock(crate::file_lock::Target::File(&dir.path().join(FILE))) .unwrap() .is_none(), - "it is crate::file_lock on .lock" + "it is crate::file_lock on the manifest file" ); let session = dir.path().join(FILE); assert!( diff --git a/crates/agent/src/tools/mcp_oauth.rs b/crates/agent/src/tools/mcp_oauth.rs index b88cb138..34115e59 100644 --- a/crates/agent/src/tools/mcp_oauth.rs +++ b/crates/agent/src/tools/mcp_oauth.rs @@ -483,11 +483,42 @@ fn is_unauthorized(e: &StreamableHttpError) -> bool { pub(crate) struct OAuthHttp { inner: C, auth: Option>, + /// Where the session this connection opens is recorded, so the exit can end it + /// (`mcp_http_exit`). This layer sees every response's session id and every `DELETE`. + exit: Option>, } impl OAuthHttp { pub(crate) fn new(inner: C, auth: Option>) -> Self { - Self { inner, auth } + Self { + inner, + auth, + exit: None, + } + } + + /// Record the session this connection opens in `exit`, for the process exit to end. + pub(crate) fn ending_session_on_exit( + mut self, + exit: Arc, + ) -> Self { + self.exit = Some(exit); + self + } + + fn saw( + &self, + response: &StreamableHttpPostResponse, + headers: &HashMap, + ) { + let session = match response { + StreamableHttpPostResponse::Json(_, Some(id)) + | StreamableHttpPostResponse::Sse(_, Some(id)) => id, + _ => return, + }; + if let Some(exit) = &self.exit { + exit.opened(session, headers); + } } } @@ -553,16 +584,19 @@ where auth_header: Option, custom_headers: HashMap, ) -> Result> { - self.authed(auth_header, |token| { - self.inner.post_message( - uri.clone(), - message.clone(), - session_id.clone(), - token, - custom_headers.clone(), - ) - }) - .await + let response = self + .authed(auth_header, |token| { + self.inner.post_message( + uri.clone(), + message.clone(), + session_id.clone(), + token, + custom_headers.clone(), + ) + }) + .await?; + self.saw(&response, &custom_headers); + Ok(response) } async fn post_message_with_max_sse_event_size( @@ -574,17 +608,20 @@ where custom_headers: HashMap, max_sse_event_size: usize, ) -> Result> { - self.authed(auth_header, |token| { - self.inner.post_message_with_max_sse_event_size( - uri.clone(), - message.clone(), - session_id.clone(), - token, - custom_headers.clone(), - max_sse_event_size, - ) - }) - .await + let response = self + .authed(auth_header, |token| { + self.inner.post_message_with_max_sse_event_size( + uri.clone(), + message.clone(), + session_id.clone(), + token, + custom_headers.clone(), + max_sse_event_size, + ) + }) + .await?; + self.saw(&response, &custom_headers); + Ok(response) } async fn delete_session( @@ -594,6 +631,10 @@ where auth_header: Option, custom_headers: HashMap, ) -> Result<(), StreamableHttpError> { + // rmcp is ending the session itself: the exit has nothing left to end. + if let Some(exit) = &self.exit { + exit.closed(&session_id); + } self.authed(auth_header, |token| { self.inner.delete_session( uri.clone(), diff --git a/crates/agent/src/tools/mcp_stdio.rs b/crates/agent/src/tools/mcp_stdio.rs index d8c445a6..c53432b6 100644 --- a/crates/agent/src/tools/mcp_stdio.rs +++ b/crates/agent/src/tools/mcp_stdio.rs @@ -111,8 +111,15 @@ pub(crate) fn retire(server: &std::sync::Arc) { /// to reap. Bounded by a server's grace plus a margin for the sweep — so exiting takes a little /// longer whenever a stdio server ran (it is given the chance to exit cleanly), and up to that /// bound when one ignores its stdin closing. +/// +/// Every streamable-HTTP MCP session still open is ended too (`mcp_http_exit`): its `DELETE`s go out +/// first, in parallel and bounded by [`CLOSE_DEADLINE`](crate::tools::mcp_http_exit::CLOSE_DEADLINE), +/// overlapping the stdio grace. pub fn sweep_before_exit() { + let http_closed = + crate::tools::mcp_http_exit::begin_close_all(crate::tools::mcp_http_exit::CLOSE_DEADLINE); retire_all(); + http_closed(); #[cfg(unix)] crate::tools::exec::wait_for_pending_group_kills( SHUTDOWN_GRACE + std::time::Duration::from_secs(2), diff --git a/crates/agent/src/tools/mod.rs b/crates/agent/src/tools/mod.rs index 310e5ddb..bfe1d9b9 100644 --- a/crates/agent/src/tools/mod.rs +++ b/crates/agent/src/tools/mod.rs @@ -20,6 +20,7 @@ pub mod mcp; pub mod mcp_apps; pub mod mcp_events; pub mod mcp_host; +pub mod mcp_http_exit; pub mod mcp_manifest; pub(crate) mod mcp_oauth; pub mod mcp_skills; @@ -176,6 +177,15 @@ pub(crate) fn write_atomic(path: &str, content: &[u8]) -> std::io::Result<()> { let p = std::path::Path::new(path); let resolved = resolve_symlink_target(p); let p = resolved.as_deref().unwrap_or(p); + // Never a fresh file renamed over a lock file (reached by name or through a symlink — resolved + // just above): the lock would stay on the old inode, which the path no longer names, so the next + // owner would lock the new file. Opening one is harmless (see `file_lock`); replacing it is not. + if crate::file_lock::is_lock_file(p) { + return Err(std::io::Error::new( + std::io::ErrorKind::PermissionDenied, + format!("{} is a lock file; it is not written", p.display()), + )); + } let name = match p.file_name() { Some(name) => name, None => return Err(std::io::Error::other(format!("invalid path: {path}"))), @@ -886,6 +896,36 @@ mod tests { assert!(target.is_dir()); } + /// Never a fresh file renamed over a lock file — the record one or, during mixed rollout, the + /// legacy one an older binary locks — directly or through a symlink: the lock would be left on + /// an inode the path no longer names. An unrelated `*.lock` (a `Cargo.lock`) is written as ever. + #[test] + #[cfg(unix)] + fn write_atomic_never_replaces_a_lock_file() { + let dir = tempfile::tempdir().unwrap(); + let session = dir.path().join("s1"); + std::fs::create_dir_all(&session).unwrap(); + let _held = crate::file_lock::try_lock(crate::file_lock::Target::Dir(&session)) + .unwrap() + .unwrap(); + let legacy = session.join("lock"); + let link = dir.path().join("innocent.txt"); + std::os::unix::fs::symlink(&legacy, &link).unwrap(); + let before = std::fs::metadata(&legacy).unwrap(); + for target in [&legacy, &link] { + assert!( + write_atomic(target.to_str().unwrap(), b"x").is_err(), + "{} must not be replaced", + target.display() + ); + } + use std::os::unix::fs::MetadataExt as _; + assert_eq!(std::fs::metadata(&legacy).unwrap().ino(), before.ino()); + let cargo_lock = dir.path().join("Cargo.lock"); + write_atomic(cargo_lock.to_str().unwrap(), b"[[package]]\n").unwrap(); + assert_eq!(std::fs::read(&cargo_lock).unwrap(), b"[[package]]\n"); + } + #[test] #[cfg(unix)] fn write_atomic_does_not_follow_a_symlink_planted_at_the_old_deterministic_temp_name() { diff --git a/crates/agent/src/tools/output.rs b/crates/agent/src/tools/output.rs index a098285b..471fc535 100644 --- a/crates/agent/src/tools/output.rs +++ b/crates/agent/src/tools/output.rs @@ -1053,7 +1053,13 @@ mod tests { // `ensure_temp_file` now sweeps stale siblings (older than `STALE_TEMP_FILE_AGE`) before // creating a new one. A fresh sibling — e.g. one from a command that just finished, whose // `full_output_path` a caller may still legitimately be about to read — must survive. - let prefix = "test-stale-sweep"; + // Unique per run: the files live in the shared system temp dir, and a fixed name let two + // concurrent runs of this suite (another checkout, a stress loop) sweep each other's files. + let prefix = &format!( + "test-stale-sweep-{}-{}", + std::process::id(), + crate::tools::temp_suffix() + ); let dir = std::env::temp_dir(); let stale_path = dir.join(format!("{prefix}-aaaaaaaaaaaaaaaa.log")); let fresh_path = dir.join(format!("{prefix}-bbbbbbbbbbbbbbbb.log")); diff --git a/crates/agent/src/worktree.rs b/crates/agent/src/worktree.rs index 201a396e..8c459126 100644 --- a/crates/agent/src/worktree.rs +++ b/crates/agent/src/worktree.rs @@ -425,7 +425,9 @@ impl Git { tokio::task::spawn_blocking(move || { for rel in &files { let from = src.join(rel); - if !from.is_file() { + // A lock file is left behind: a copy of one is not a lock (see + // `file_lock`). + if !from.is_file() || crate::file_lock::is_lock_file(&from) { continue; } let to = dst.join(rel); @@ -1229,6 +1231,53 @@ mod tests { ); } + /// A lock file sitting untracked in the repo is left behind by the seed (a copy of a lock is no + /// lock), and copying a hard link to it — an in-process open and close of the lock file — does + /// not release it (see `file_lock`). + #[cfg(unix)] + #[tokio::test] + async fn seeding_never_copies_a_held_lock_file_and_leaves_it_held() { + let repo = repo().await; + let state = repo.path().join("state"); + let lock = crate::file_lock::try_lock(crate::file_lock::Target::File(&state)) + .unwrap() + .unwrap(); + let record = std::fs::read_dir(repo.path()) + .unwrap() + .flatten() + .map(|e| e.path()) + .find(|p| crate::file_lock::is_record_lock_file(p)) + .unwrap(); + std::fs::hard_link(&record, repo.path().join("innocent.txt")).unwrap(); + std::fs::write(repo.path().join("brand-new.rs"), "fn main() {}\n").unwrap(); + + let root = preflight(&Git::Local, repo.path()).await.unwrap(); + let wt = Worktree::create(&Git::Local, &root, "scout-0") + .await + .unwrap(); + assert!(wt.path().join("brand-new.rs").exists(), "seeding ran"); + let copied: Vec<_> = std::fs::read_dir(wt.path()) + .unwrap() + .flatten() + .map(|e| e.path()) + .collect(); + assert!( + !copied + .iter() + .any(|p| crate::file_lock::is_record_lock_file(p)), + "the record lock file was copied: {copied:?}" + ); + // A hard link to it is just a file to copy — and copying it (an open and a close of the lock + // file) leaves the lock held. + assert!(wt.path().join("innocent.txt").exists()); + if let Some(held) = crate::file_lock::tests::held_for_another_process( + crate::file_lock::Target::File(&state), + ) { + assert!(held, "seeding the worktree left the lock held"); + } + drop(lock); + } + #[tokio::test] async fn child_delta_contains_only_the_childs_changes_not_the_parents_wip() { // The correctness property the baseline commit exists for. Without it, the parent's own diff --git a/crates/agent/tests/common/mcp_oauth_fixture.rs b/crates/agent/tests/common/mcp_oauth_fixture.rs index 0c88c8fe..ca500726 100644 --- a/crates/agent/tests/common/mcp_oauth_fixture.rs +++ b/crates/agent/tests/common/mcp_oauth_fixture.rs @@ -146,6 +146,9 @@ pub fn write_json(stream: &mut TcpStream, status: &str, value: &Value) { write_response(stream, status, "Content-Type: application/json\r\n", &body); } +/// One request's `(Authorization, Mcp-Session-Id)`. +pub type HeaderPair = (Option, Option); + /// A running OAuth-protected MCP server and its controls. See the module doc. #[derive(Clone)] pub struct OAuthFixture { @@ -184,6 +187,8 @@ pub struct OAuthFixture { pub discover_reply: Arc>>, /// Answer notifications with an empty `200 OK` (no content type) instead of `202 Accepted`. pub notifications_empty_200: Arc, + /// Never answer a `DELETE` (hold the connection 30 s). + pub hang_deletes: Arc, /// What arrived: `initialize`s, `GET` streams, `DELETE`s, POSTs refused for a missing or /// expired session, and the `Mcp-Session-Id` each `tools/call` carried. pub initializes: Arc, @@ -191,6 +196,10 @@ pub struct OAuthFixture { pub deletes: Arc, pub session_refusals: Arc, pub call_sessions: Arc>>>, + /// Each `DELETE`'s `(Authorization, Mcp-Session-Id)`. + pub delete_requests: Arc>>, + /// The client's answers to server→client requests (JSON-RPC responses POSTed back), by id. + pub client_answers: Arc>>, issued: Arc>>, live_sessions: Arc>>, } @@ -230,11 +239,14 @@ impl OAuthFixture { sse_responses: Arc::default(), discover_reply: Arc::default(), notifications_empty_200: Arc::default(), + hang_deletes: Arc::default(), initializes: Arc::default(), get_streams: Arc::default(), deletes: Arc::default(), session_refusals: Arc::default(), call_sessions: Arc::default(), + delete_requests: Arc::default(), + client_answers: Arc::default(), issued: Arc::default(), live_sessions: Arc::default(), }; @@ -426,6 +438,17 @@ impl Shared { if !f.sessions.load(Ordering::SeqCst) { return write_response(stream, "405 Method Not Allowed", "", b""); } + if req.method == "DELETE" { + f.delete_requests.lock().unwrap().push(( + req.headers.get("authorization").cloned(), + req.headers.get("mcp-session-id").cloned(), + )); + if f.hang_deletes.load(Ordering::SeqCst) { + // A server that never answers the session's end. + thread::sleep(Duration::from_secs(30)); + return; + } + } if !self.authorized(req) { return self.unauthorized(stream, &Value::Null); } @@ -507,6 +530,14 @@ impl Shared { Some(_) => {} } } + // The client answering a server→client request (no `method`, a `result` or `error`). + if method.is_empty() && (request.get("result").is_some() || request.get("error").is_some()) + { + f.client_answers.lock().unwrap().push(request.clone()); + let (_, cv) = &self.gate; + cv.notify_all(); + return write_response(stream, "202 Accepted", "", b""); + } if request.get("id").is_none() { return if f.notifications_empty_200.load(Ordering::SeqCst) { write_response(stream, "200 OK", "", b"") @@ -547,7 +578,15 @@ impl Shared { "properties": { "text": { "type": "string" } }, "required": ["text"], }, + }, { + "name": "ask_user", + "description": "Asks the user for their name, then greets them.", + "inputSchema": { "type": "object", "properties": {} }, }] }), + _ if request.pointer("/params/name").and_then(Value::as_str) == Some("ask_user") => { + f.call_sessions.lock().unwrap().push(session.clone()); + return self.ask_user(stream, &request, &id, &extra); + } _ => { f.call_sessions.lock().unwrap().push(session.clone()); let text = request @@ -585,6 +624,61 @@ impl Shared { } } } + + /// `ask_user`, answered the way a real streamable-HTTP server does it: an SSE stream on the + /// call's own POST, **flushed event by event** — a progress notification, then an + /// `elicitation/create` request to the client — and the response only once the client's answer + /// to that request has arrived (POSTed back on another connection). A client that buffered the + /// stream to its end would never see the question, and the call would hang to the 10 s cap. + fn ask_user(&self, stream: &mut TcpStream, request: &Value, id: &Value, extra: &str) { + let f = &self.fixture; + let token = request.pointer("/params/_meta/progressToken").cloned(); + let ask_id = format!("srv-elicit-{}", self.calls.fetch_add(1, Ordering::SeqCst)); + let _ = stream.write_all( + format!("HTTP/1.1 200 OK\r\n{extra}Content-Type: text/event-stream\r\nCache-Control: no-cache\r\nConnection: close\r\n\r\n") + .as_bytes(), + ); + let mut event = |message: Value| { + let _ = stream.write_all(format!("data: {message}\n\n").as_bytes()); + let _ = stream.flush(); + }; + if let Some(token) = token { + event( + json!({ "jsonrpc": "2.0", "method": "notifications/progress", + "params": { "progressToken": token, "progress": 1, "total": 2, + "message": "asking the user" } }), + ); + } + thread::sleep(Duration::from_millis(50)); + event( + json!({ "jsonrpc": "2.0", "id": ask_id, "method": "elicitation/create", + "params": { "message": "What is your name?", + "requestedSchema": { "type": "object", + "properties": { "name": { "type": "string" } }, + "required": ["name"] } } }), + ); + // Wait for the answer, on another connection, before ending this stream. + let answered = + |answers: &Vec| answers.iter().find(|a| a["id"] == json!(ask_id)).cloned(); + let (count, cv) = &self.gate; + let deadline = std::time::Instant::now() + Duration::from_secs(10); + let mut answer = answered(&f.client_answers.lock().unwrap()); + while answer.is_none() && std::time::Instant::now() < deadline { + let guard = count.lock().unwrap(); + let _ = cv.wait_timeout(guard, Duration::from_millis(100)).unwrap(); + answer = answered(&f.client_answers.lock().unwrap()); + } + let text = match answer { + Some(a) if a["result"]["action"] == "accept" => format!( + "hello-{}", + a["result"]["content"]["name"].as_str().unwrap_or("?") + ), + Some(a) => format!("not answered: {}", a["result"]["action"]), + None => "no answer before the stream had to end".to_owned(), + }; + event(json!({ "jsonrpc": "2.0", "id": id, + "result": { "content": [{ "type": "text", "text": text }], "isError": false } })); + } } /// The original two-value form: the MCP URL and the accepted `Authorization` headers. diff --git a/crates/agent/tests/mcp_apps_session.rs b/crates/agent/tests/mcp_apps_session.rs index 13587409..d46ba66c 100644 --- a/crates/agent/tests/mcp_apps_session.rs +++ b/crates/agent/tests/mcp_apps_session.rs @@ -12,8 +12,8 @@ mod common; use std::time::Duration; use common::mcp_apps::{ - Ws, app_request, command, daemon, daemon_with, declare_apps, home_with_fixture, of_type, - prompt, tool_end_text, view_id, views_model, + Ws, app_request, command, daemon, daemon_opts, daemon_with, declare_apps, home_with_fixture, + of_type, prompt, tool_end_text, view_id, views_model, }; use common::{ body_json, spawn_model_server_routed, turn_text, turn_tool_use, ws_next_frame, ws_send, @@ -371,7 +371,20 @@ async fn view_html_is_cached_per_connection_until_the_server_says_resources_chan call("toolu_1"), ); let sessions = tempfile::tempdir().unwrap(); - let daemon = daemon(&home.home, &base, sessions.path(), "0").await; + // The daemon's debug log says when `list_changed` has been handled (rmcp runs the handler on a + // task of its own, so the response to `change_view` can arrive first): the wait below is for + // that line, not for a guessed interval. + let log = sessions.path().join("serve.stderr"); + let daemon = daemon_opts( + &home.home, + &base, + sessions.path(), + "0", + &[], + &[("RUST_LOG", "beyond_ai_agent::tools::mcp=debug")], + Some(&log), + ) + .await; let mut ws = daemon.connect("cache").await; declare_apps(&mut ws).await; @@ -402,8 +415,22 @@ async fn view_html_is_cached_per_connection_until_the_server_says_resources_chan ) .await; assert_eq!(r["data"]["result"]["content"][0]["text"], "view-v2", "{r}"); - // `list_changed` is handled on the client's own task; give it a moment to land. - tokio::time::sleep(Duration::from_millis(200)).await; + // Wait until the client has handled `list_changed` on its own task. + let handled = |log: &std::path::Path| { + std::fs::read_to_string(log) + .unwrap_or_default() + .matches("resources/list_changed handled") + .count() + }; + let deadline = std::time::Instant::now() + Duration::from_secs(20); + while handled(&log) == 0 { + assert!( + std::time::Instant::now() < deadline, + "list_changed was never handled: {}", + std::fs::read_to_string(&log).unwrap_or_default() + ); + tokio::time::sleep(Duration::from_millis(10)).await; + } let three = prompt(&mut ws, "zq3x").await; assert!(html(&three).contains("weather v2"), "{}", html(&three)); assert_eq!(home.reads_of(VIEW_URI), 2); diff --git a/crates/agent/tests/mcp_oauth_routing.rs b/crates/agent/tests/mcp_oauth_routing.rs index 66b067ba..7e9d6ee6 100644 --- a/crates/agent/tests/mcp_oauth_routing.rs +++ b/crates/agent/tests/mcp_oauth_routing.rs @@ -8,10 +8,13 @@ mod common; +use std::io::{BufRead, Write}; use std::sync::atomic::Ordering; +use std::time::{Duration, Instant}; use common::mcp_oauth_fixture::{echo, logged_in, run, tool_results}; -use common::turn_text; +use common::{SpawnGuarded, read_until_response, turn_text}; +use serde_json::{Value, json}; /// `toolu_`'s result in the model request that followed it: `(text, is_error)`. fn result_of(bodies: &[String], request: usize, id: &str) -> (String, bool) { @@ -118,6 +121,174 @@ fn sessions_sse_responses_and_the_standalone_stream_work_for_an_oauth_server() { ); } +#[test] +fn closing_an_oauth_servers_connection_ends_its_session_with_an_authenticated_delete() { + // The idle reaper closes the connection a second after the call: rmcp ends the session with a + // DELETE, which must carry the bearer (through the OAuth layer) and the session id. + let (home, fixture) = logged_in(); + fixture.sessions.store(true, Ordering::SeqCst); + let (base, _bodies) = common::spawn_model_server(vec![echo("toolu_1", "x"), turn_text("done")]); + let dir = tempfile::tempdir().unwrap(); + let session_file = dir.path().join("s.jsonl").to_string_lossy().into_owned(); + let mut child = common::serve_cmd(common::BIN, &base, &session_file) + .env("HOME", home.path()) + .env("BEYOND_AI_AGENT_MCP_IDLE_SECS", "1") + .spawn_guarded(); + let mut stdin = child.stdin.take().unwrap(); + let mut stdout = common::child_frames(&mut child); + writeln!(stdin, "{}", json!({ "type": "prompt", "message": "go" })).unwrap(); + stdin.flush().unwrap(); + read_until_response(&mut stdout, "prompt"); + + let deadline = Instant::now() + Duration::from_secs(15); + while fixture.deletes.load(Ordering::SeqCst) == 0 && Instant::now() < deadline { + std::thread::sleep(Duration::from_millis(100)); + } + drop(stdin); + child.wait().unwrap(); + assert_eq!( + fixture.deletes.load(Ordering::SeqCst), + 1, + "the reaped connection ends its session with a DELETE" + ); + assert_eq!( + *fixture.delete_requests.lock().unwrap(), + [( + Some("Bearer access-token-0".to_owned()), + Some("session-0".to_owned()) + )], + "the DELETE carries the bearer and the session id" + ); +} + +#[test] +fn an_agent_run_ends_its_oauth_session_with_one_authenticated_delete_before_it_exits() { + // A process that simply exits closes nothing on its own; the exit sweep ends every HTTP + // session still open (`mcp_http_exit`), so by the time `run` has exited the DELETE is in. + let (home, fixture) = logged_in(); + fixture.sessions.store(true, Ordering::SeqCst); + let bodies = run(home.path(), vec![echo("toolu_1", "x"), turn_text("done")]); + let (text, is_error) = result_of(&bodies, 1, "toolu_1"); + assert!(!is_error && text.contains('x'), "{text}"); + assert_eq!( + *fixture.delete_requests.lock().unwrap(), + [( + Some("Bearer access-token-0".to_owned()), + Some("session-0".to_owned()) + )], + "exactly one DELETE, carrying the bearer and the session id, before the process exited" + ); + assert_eq!( + fixture.deletes.load(Ordering::SeqCst), + 1, + "and it was accepted" + ); +} + +#[test] +fn an_exit_whose_delete_is_never_answered_stays_bounded() { + let (home, fixture) = logged_in(); + fixture.sessions.store(true, Ordering::SeqCst); + let started = Instant::now(); + run(home.path(), vec![echo("toolu_1", "x"), turn_text("done")]); + let baseline = started.elapsed(); + + let (home, fixture2) = logged_in(); + fixture2.sessions.store(true, Ordering::SeqCst); + fixture2.hang_deletes.store(true, Ordering::SeqCst); + let started = Instant::now(); + run(home.path(), vec![echo("toolu_1", "x"), turn_text("done")]); + let hung = started.elapsed(); + assert_eq!( + fixture2.delete_requests.lock().unwrap().len(), + 1, + "the DELETE was sent" + ); + // The server holds the DELETE for 30 s; the exit waits at most its deadline (1.5 s). + assert!( + hung < baseline + Duration::from_secs(4), + "an unanswered DELETE must not hold the exit: {hung:?} (baseline {baseline:?})" + ); + drop(fixture); +} + +#[test] +fn an_oauth_servers_in_post_question_is_routed_to_the_client_and_answered_mid_stream() { + // The fixture answers `ask_user` on the POST's own SSE stream, flushed event by event: a + // progress notification, then an `elicitation/create` to the client — and holds the response + // until the client's answer has come back on another POST. So the call only completes with + // the answer if each event was handled as it arrived, before the stream ended. + let (home, fixture) = logged_in(); + fixture.sessions.store(true, Ordering::SeqCst); + fixture.sse_responses.store(true, Ordering::SeqCst); + let (base, _bodies) = common::spawn_model_server(vec![ + common::turn_tool_use("toolu_e", "mcp__protected__ask_user", "{}"), + turn_text("done"), + ]); + let dir = tempfile::tempdir().unwrap(); + let session_file = dir.path().join("s.jsonl").to_string_lossy().into_owned(); + let mut child = common::serve_cmd(common::BIN, &base, &session_file) + .env("HOME", home.path()) + .env("BEYOND_AI_AGENT_MCP_IDLE_SECS", "0") + .spawn_guarded(); + let mut stdin = child.stdin.take().unwrap(); + let mut stdout = common::child_frames(&mut child); + writeln!(stdin, "{}", json!({ "type": "prompt", "message": "ask" })).unwrap(); + stdin.flush().unwrap(); + + // Frames up to the question: the progress event must already be among them. + let mut before_question = Vec::new(); + let deadline = Instant::now() + Duration::from_secs(20); + let mut line = String::new(); + let request_id = loop { + assert!( + Instant::now() < deadline, + "no elicitation_request: {before_question:?}" + ); + line.clear(); + assert_ne!(stdout.read_line(&mut line).unwrap(), 0, "serve exited"); + let Ok(frame) = serde_json::from_str::(line.trim()) else { + continue; + }; + if frame["type"] == "elicitation_request" { + assert_eq!(frame["server"], "protected"); + break frame["request_id"].as_str().unwrap().to_owned(); + } + before_question.push(frame); + }; + assert!( + before_question.iter().any(|f| f["type"] == "event" + && f["event"]["kind"] == "tool_progress" + && f.to_string().contains("asking the user")), + "the stream's first event reached the client before its last was sent: {before_question:?}" + ); + + writeln!( + stdin, + "{}", + json!({ "type": "elicit", "request_id": request_id, "action": "accept", + "content": { "name": "Ferris" } }) + ) + .unwrap(); + stdin.flush().unwrap(); + let frames = read_until_response(&mut stdout, "prompt"); + drop(stdin); + child.wait().unwrap(); + + let tool_text: String = frames + .iter() + .filter(|f| f["type"] == "event" && f["event"]["kind"] == "tool_end") + .filter_map(|f| f["event"]["result"].as_str()) + .collect(); + assert!( + tool_text.contains("hello-Ferris"), + "the answer reached the server mid-stream and the call completed with it: {tool_text:?}" + ); + let answers = fixture.client_answers.lock().unwrap().clone(); + assert_eq!(answers.len(), 1, "{answers:?}"); + assert_eq!(answers[0]["result"]["action"], "accept"); +} + #[test] fn an_expired_session_is_reinitialized_and_the_call_retried() { // The 404 for an expired session is rmcp's `SessionExpired`: it re-initializes and retries.