diff --git a/CHANGELOG.md b/CHANGELOG.md index 08926e5e46..6760f829cc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed +- Agent names that had already been registered can be reclaimed again. Relaycast made registration create-only, so the broker's register-or-rotate paths failed with an opaque `401 Agent token required`; supervisor restart, offline-agent attach and the broker's own reconnect now reclaim the name through an audited takeover instead, and crash recovery uses the explicit recover route. - `agent-relay fleet spawn --node ` now uses the active workspace to mint and clean up a short-lived launcher identity when no agent token is present. ### Added diff --git a/Cargo.lock b/Cargo.lock index b22943d7c4..fe50482be1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2018,9 +2018,9 @@ dependencies = [ [[package]] name = "relaycast" -version = "6.0.0" +version = "7.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b78f728d34cfb36d07d78b1d28d70784810bf4d1dcffc896906d92eac27b1669" +checksum = "b8a323473bd0deb2fba53f540250038a1d012d495d494066e3bb8b2128e166b7" dependencies = [ "futures-util", "reqwest", diff --git a/crates/broker/Cargo.toml b/crates/broker/Cargo.toml index c74a07d479..f6c67634d9 100644 --- a/crates/broker/Cargo.toml +++ b/crates/broker/Cargo.toml @@ -28,7 +28,7 @@ serde_json = "1.0" sha2 = "0.10" shlex = "1.3" thiserror = "2.0" -relaycast = "=6.0.0" +relaycast = "=7.0.0" tokio = { version = "1.44", features = ["full"] } tracing = "0.1" tracing-appender = "0.2" diff --git a/crates/broker/src/relaycast/auth.rs b/crates/broker/src/relaycast/auth.rs index 9a1dd574ee..c37d825b80 100644 --- a/crates/broker/src/relaycast/auth.rs +++ b/crates/broker/src/relaycast/auth.rs @@ -1,6 +1,9 @@ use anyhow::{Context, Result}; use chrono::{DateTime, Utc}; -use relaycast::{CreateAgentRequest, RelayCast, RelayCastOptions, RelayError}; +use relaycast::{ + CreateAgentRequest, RecoverAgentRequest, RelayCast, RelayCastOptions, RelayError, + WorkspaceProvenance, +}; use reqwest::StatusCode; use serde::{Deserialize, Serialize}; use serde_json::Value; @@ -458,9 +461,16 @@ impl AuthClient { let api_key = normalize_workspace_key(&cached.api_key) .context("cached api_key is not a valid workspace key")?; + // Rotation is authenticated as the agent itself (relaycast 7.0.0): it + // needs this agent's current token, not the workspace key. Without a + // cached token there is nothing to roll over, and silently falling back + // to the workspace key is what used to fail as an opaque 401. + let agent_token = cached.agent_token.as_deref().context( + "cannot rotate token without the agent's current token; re-register or recover the identity", + )?; let relay = build_relay_client(&api_key, self.base_url.as_deref())?; let result = relay - .rotate_agent_token(agent_name) + .rotate_agent_token(agent_name, agent_token) .await .map_err(relay_error_to_anyhow)?; let token = result.token; @@ -696,7 +706,13 @@ impl AuthClient { } async fn create_workspace(&self, name: &str) -> Result<(String, String)> { - match RelayCast::create_workspace(name, self.base_url.as_deref()).await { + match RelayCast::create_workspace( + name, + self.base_url.as_deref(), + WorkspaceProvenance::sdk(), + ) + .await + { Ok(result) => Ok((result.workspace_id, result.api_key)), Err(error) if is_workspace_name_conflict(&error) => { let suffix = Uuid::new_v4().simple().to_string(); @@ -706,9 +722,13 @@ impl AuthClient { fallback_name = %fallback_name, "workspace already exists; retrying with a fresh fallback name" ); - let result = RelayCast::create_workspace(&fallback_name, self.base_url.as_deref()) - .await - .map_err(relay_error_to_anyhow)?; + let result = RelayCast::create_workspace( + &fallback_name, + self.base_url.as_deref(), + WorkspaceProvenance::sdk(), + ) + .await + .map_err(relay_error_to_anyhow)?; Ok((result.workspace_id, result.api_key)) } Err(error) => Err(relay_error_to_anyhow(error)), @@ -735,7 +755,7 @@ impl AuthClient { .map(ToOwned::to_owned) .unwrap_or_else(|| format!("agent-{}", Uuid::new_v4().simple())); - admit_agent_registration(&relay, &name, agent_type, identity_key).await + admit_agent_registration(&relay, workspace_key, &name, agent_type, identity_key).await } pub async fn workspace_key_is_live(&self, workspace_key: &str) -> Result { @@ -984,6 +1004,9 @@ pub(crate) fn identity_key_fingerprint(raw: &str) -> String { /// collision is rejected. async fn admit_agent_registration( relay: &RelayCast, + // Needed only for the pre-8.2.0 rotate fallback below, where the workspace + // key is still an accepted credential for reclaiming an agent's token. + workspace_key: &str, name: &str, agent_type: Option<&str>, identity_key: Option<&str>, @@ -1031,14 +1054,58 @@ async fn admit_agent_registration( })); } + // Reclaiming a crashed work unit's identity. We have just proven + // ownership via the work-unit identity key but do not hold the + // agent's token — which is precisely the case `recover` exists for + // (relaycast 7.0.0). Rotation is self-rollover only and cannot serve + // this path: attempting it with the workspace key is what produced + // an opaque `401 Agent token required` here. let token_response = relay - .rotate_agent_token(&existing.name) - .await - .map_err(relay_error_to_anyhow)?; + .recover_agent( + &existing.name, + RecoverAgentRequest { + expected_agent_id: existing.id.clone(), + recovery_proof: None, + reason: Some( + "work-unit identity key proved ownership after a crash".to_string(), + ), + // Hashed, never raw: the identity key is a replayable + // credential and the surrounding code hashes it before + // anything workspace-readable. An audit record is + // workspace-readable, so it gets the same treatment — + // still correlatable, not replayable. + session_ref: identity_key.map(hash_identity_key), + node_id: None, + }, + ) + .await; + // Engines before 8.2.0 have no `/recover` route — the identity + // recovery surface arrived with it — and answer "Route not found". + // Those engines still let the workspace key rotate an agent's + // token, which is what this path did before, so fall back rather + // than failing every node restart against an older engine. The + // agent-specific 404 (`agent_not_found`) is a real failure and is + // deliberately excluded. + let token_response = match token_response { + Ok(response) => response.token, + Err(RelayError::Api { + status: 404, + ref code, + .. + }) if code != "agent_not_found" => relay + .rotate_agent_token(&existing.name, workspace_key) + .await + .map_err(relay_error_to_anyhow) + .context( + "recover unavailable on this engine and the legacy rotate fallback failed", + )? + .token, + Err(error) => return Err(relay_error_to_anyhow(error)), + }; Ok(( existing.id, existing.name, - token_response.token, + token_response, existing.workspace_id, )) } @@ -1620,13 +1687,16 @@ mod tests { .header("content-type", "application/json") .body(r#"{"ok":true,"data":{"id":"a_existing","name":"lead","type":"agent","status":"offline","persona":null,"metadata":{},"last_seen":"2025-01-01T00:00:00Z","channels":[]}}"#); }); + // Identity reclaim uses `recover`, not rotation (relaycast 7.0.0). let rotate = server.mock(|when, then| { when.method(POST) - .path("/v1/agents/lead/rotate-token") + .path("/v1/agents/lead/recover") .header("authorization", "Bearer rk_live_shared"); then.status(200) .header("content-type", "application/json") - .body(r#"{"ok":true,"data":{"name":"lead","token":"at_live_rotated"}}"#); + .body( + r#"{"ok":true,"data":{"agent_id":"a_existing","name":"lead","token":"at_live_rotated","audit_id":"aud_1"}}"#, + ); }); let client = AuthClient::new(Some(server.base_url())); @@ -1652,6 +1722,64 @@ mod tests { } } + /// Engines before 8.2.0 have no `/recover`. A node restarting against one + /// must still reclaim its identity via the legacy workspace-key rotate — + /// this is the exact path the two-node fleet e2e drives, against a pinned + /// v7.0.0 engine. + #[tokio::test] + async fn identity_reclaim_falls_back_to_legacy_rotate_on_older_engines() { + let _env_guard = clear_relay_env(); + let identity = "work-unit-42"; + let identity_hash = hash_identity_key(identity); + let server = MockServer::start(); + unsafe { + std::env::set_var("RELAY_API_KEY", "rk_live_shared"); + std::env::set_var("RELAY_AGENT_IDENTITY_KEY", identity); + } + server.mock(|when, then| { + when.method(POST).path("/v1/agents"); + then.status(409) + .header("content-type", "application/json") + .body(r#"{"ok":false,"error":{"code":"agent_already_exists","message":"name_taken"}}"#); + }); + server.mock(|when, then| { + when.method(GET).path("/v1/agents/lead"); + then.status(200) + .header("content-type", "application/json") + .body(format!( + r#"{{"ok":true,"data":{{"id":"a_existing","name":"lead","type":"agent","status":"offline","persona":null,"metadata":{{"identity_key":"{identity_hash}"}},"last_seen":"2025-01-01T00:00:00Z","channels":[]}}}}"# + )); + }); + // Older engine: the recovery surface simply is not mounted. + let recover = server.mock(|when, then| { + when.method(POST).path("/v1/agents/lead/recover"); + then.status(404) + .header("content-type", "application/json") + .body(r#"{"ok":false,"error":{"code":"not_found","message":"Route not found"}}"#); + }); + let legacy_rotate = server.mock(|when, then| { + when.method(POST) + .path("/v1/agents/lead/rotate-token") + .header("authorization", "Bearer rk_live_shared"); + then.status(200) + .header("content-type", "application/json") + .body(r#"{"ok":true,"data":{"name":"lead","token":"at_live_legacy_reclaim"}}"#); + }); + + let client = AuthClient::new(Some(server.base_url())); + let session = client + .startup_session_set_with_identity(Some("lead"), true, None, Some(identity)) + .await + .expect("an older engine must still let a restarting node reclaim its name") + .default_session() + .cloned() + .expect("a session was registered"); + + assert_eq!(session.token, "at_live_legacy_reclaim"); + recover.assert_hits(1); + legacy_rotate.assert_hits(1); + } + #[tokio::test] async fn strict_name_conflict_with_matching_identity_reclaims_existing_agent() { // Crash-recovery resume: the SAME work unit re-registers under the @@ -1691,13 +1819,16 @@ mod tests { r#"{{"ok":true,"data":{{"id":"a_existing","name":"lead","type":"agent","status":"offline","persona":null,"metadata":{{"identity_key":"{identity_hash}"}},"last_seen":"2025-01-01T00:00:00Z","channels":[]}}}}"# )); }); + // Identity reclaim uses `recover`, not rotation (relaycast 7.0.0). let rotate = server.mock(|when, then| { when.method(POST) - .path("/v1/agents/lead/rotate-token") + .path("/v1/agents/lead/recover") .header("authorization", "Bearer rk_live_shared"); then.status(200) .header("content-type", "application/json") - .body(r#"{"ok":true,"data":{"name":"lead","token":"at_live_rotated"}}"#); + .body( + r#"{"ok":true,"data":{"agent_id":"a_existing","name":"lead","token":"at_live_rotated","audit_id":"aud_1"}}"#, + ); }); let client = AuthClient::new(Some(server.base_url())); @@ -1780,13 +1911,18 @@ mod tests { r#"{{"ok":true,"data":{{"id":"a_existing","name":"node-a","type":"agent","status":"offline","persona":null,"metadata":{{"identity_key":"{identity_hash}"}},"last_seen":"2025-01-01T00:00:00Z","channels":[]}}}}"# )); }); + // Reclaiming a crashed identity goes through `recover`, not rotation: + // rotation is self-rollover only and this path holds the work-unit + // identity proof rather than the agent's token (relaycast 7.0.0). let rotate = server.mock(|when, then| { when.method(POST) - .path("/v1/agents/node-a/rotate-token") + .path("/v1/agents/node-a/recover") .header("authorization", "Bearer rk_live_shared"); then.status(200) .header("content-type", "application/json") - .body(r#"{"ok":true,"data":{"name":"node-a","token":"at_live_rotated"}}"#); + .body( + r#"{"ok":true,"data":{"agent_id":"a_existing","name":"node-a","token":"at_live_rotated","audit_id":"aud_1"}}"#, + ); }); let client = AuthClient::new(Some(server.base_url())); @@ -2415,7 +2551,11 @@ mod tests { let first_conflict = server.mock(|when, then| { when.method(POST) .path("/v1/workspaces") - .json_body(json!({ "name": workspace_name })); + // `body_contains` rather than an exact `json_body`: the request + // now also carries workspace provenance (relaycast 7.0.0), and + // this mock only cares that the first attempt uses the + // deterministic name. + .body_contains(format!("\"name\":\"{workspace_name}\"")); then.status(409) .header("content-type", "application/json") .body( @@ -2457,7 +2597,8 @@ mod tests { let rotate = server.mock(|when, then| { when.method(POST) .path("/v1/agents/lead/rotate-token") - .header("authorization", "Bearer rk_live_cached"); + // Self-rollover authenticates as the agent, not the workspace. + .header("authorization", "Bearer at_live_current"); then.status(200) .header("content-type", "application/json") .body(r#"{"ok":true,"data":{"token":"at_live_rotated","name":"lead"}}"#); @@ -2471,7 +2612,7 @@ mod tests { agent_id: "a_old".into(), api_key: "rk_live_cached".into(), agent_name: Some("lead".into()), - agent_token: None, + agent_token: Some("at_live_current".into()), updated_at: chrono::Utc::now(), }; @@ -2490,7 +2631,8 @@ mod tests { let rotate_404 = server.mock(|when, then| { when.method(POST) .path("/v1/agents/lead/rotate-token") - .header("authorization", "Bearer rk_live_cached"); + // Self-rollover authenticates as the agent, not the workspace. + .header("authorization", "Bearer at_live_current"); then.status(404) .header("content-type", "application/json") .body(r#"{"ok":false,"error":{"code":"not_found","message":"not found"}}"#); @@ -2512,7 +2654,7 @@ mod tests { agent_id: "a_old".into(), api_key: "rk_live_cached".into(), agent_name: Some("lead".into()), - agent_token: None, + agent_token: Some("at_live_current".into()), updated_at: chrono::Utc::now(), }; diff --git a/crates/broker/src/relaycast/ws.rs b/crates/broker/src/relaycast/ws.rs index c04711b9c7..86d92a6ba4 100644 --- a/crates/broker/src/relaycast/ws.rs +++ b/crates/broker/src/relaycast/ws.rs @@ -1,13 +1,17 @@ -use std::{collections::BTreeSet, sync::Arc, time::Duration}; +use std::{ + collections::{BTreeSet, HashMap}, + sync::{Arc, Mutex as StdMutex}, + time::Duration, +}; use anyhow::{Context, Result}; use relaycast::{ agent::DmOptions, format_registration_error, retry_agent_registration as sdk_retry_agent_registration, ActionDefinition, ActionInvocation, - AgentClient, AgentRegistrationClient, AgentRegistrationError, AgentRegistrationRetryOutcome, - CompleteInvocationRequest, CreateObserverTokenRequest, EmitSessionEventRequest, - MessageListQuery, ObserverToken, RegisterActionRequest, RelayCast, RelayCastOptions, - RelayError, ReleaseAgentRequest, UpdateAgentRequest, + AgentClient, AgentIdentityRecoveryResponse, AgentRegistrationClient, AgentRegistrationError, + AgentRegistrationRetryOutcome, CompleteInvocationRequest, CreateObserverTokenRequest, + EmitSessionEventRequest, MessageListQuery, ObserverToken, RegisterActionRequest, RelayCast, + RelayCastOptions, RelayError, ReleaseAgentRequest, TakeOverAgentRequest, UpdateAgentRequest, }; use serde_json::Value; @@ -34,6 +38,11 @@ pub struct RelaycastHttpClient { pub api_key: String, relay: Arc>, registration: Arc>, + /// One lock per agent name, guarding collision recovery. Takeover keeps the + /// same agent id, so `expected_agent_id` cannot stop two concurrent + /// cache-miss registrations from both taking the name over — the second + /// response would invalidate the token handed to the first caller. + takeover_locks: Arc>>>>, pub agent_name: String, pub default_cli: String, } @@ -146,6 +155,7 @@ impl RelaycastHttpClient { api_key, relay, registration, + takeover_locks: Arc::new(StdMutex::new(HashMap::new())), agent_name: agent_name.into(), default_cli, } @@ -227,9 +237,154 @@ impl RelaycastHttpClient { detail: "SDK relay client not initialized".to_string(), } })?; - registration + match registration .register_agent_token(trimmed_name, cli_hint) .await + { + Ok(token) => Ok(token), + // Registration is create-only as of relaycast 8.2.0 / SDK 7.0.0, so + // a name the broker already owns can no longer be re-registered and + // rotation is self-rollover the broker cannot perform. The broker + // holds the workspace key, which makes reclaiming one of its own + // agents a `takeover` — the explicit, audited operation #349 built + // for exactly this — rather than the silent identity replacement it + // removed. Callers that must not seize a live agent go through + // `register_agent_token_with_intent`, whose presence probe runs + // first. + Err(AgentRegistrationError::AlreadyExists { .. }) => { + self.take_over_agent_identity(trimmed_name, None).await + } + Err(other) => Err(other), + } + } + + /// Reclaim an agent name this workspace already owns, leaving an audit + /// record. Used when create-only registration reports the name is taken. + async fn take_over_agent_identity( + &self, + agent_name: &str, + known_agent_id: Option<&str>, + ) -> std::result::Result { + // Singleflight per agent name. Two concurrent cache misses would + // otherwise both take the name over, and the second response would + // invalidate the token already returned to the first caller — + // `expected_agent_id` cannot catch that, because takeover preserves the + // agent id. + let lock = { + let mut locks = self + .takeover_locks + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + Arc::clone( + locks + .entry(agent_name.to_string()) + .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))), + ) + }; + let _guard = lock.lock().await; + + // Re-check after acquiring: a concurrent caller may have completed the + // takeover while we waited, in which case its token is the live one and + // taking over again would invalidate it. + if let Some(registration) = self.registration.as_ref().as_ref() { + if let Some(cached) = registration.cached_agent_token(agent_name) { + return Ok(cached); + } + } + + let relay = + (*self.relay) + .as_ref() + .ok_or_else(|| RelaycastRegistrationError::Transport { + agent_name: agent_name.to_string(), + detail: "SDK relay client not initialized".to_string(), + })?; + + // Callers that already resolved the incumbent (the impersonation + // presence probe does) pass its id through rather than paying a second + // lookup on every registration. + let existing_id = match known_agent_id { + Some(id) => id.to_string(), + None => { + relay + .get_agent(agent_name) + .await + .map_err(|error| RelaycastRegistrationError::Transport { + agent_name: agent_name.to_string(), + detail: format!("failed to resolve existing agent for takeover: {error}"), + })? + .id + } + }; + + let response = match relay + .take_over_agent( + agent_name, + TakeOverAgentRequest { + expected_agent_id: existing_id.clone(), + actor: self.agent_name.clone(), + reason: "broker reclaimed an agent name it owns after create-only registration reported a collision".to_string(), + session_ref: existing_id.clone(), + node_id: self.agent_name.clone(), + }, + ) + .await + { + Ok(response) => response, + // Engines before 8.2.0 have no `/takeover` route — the whole + // identity-recovery surface arrived with it. Those engines still + // allow the workspace key to rotate an agent's token, which is what + // this path did before, so fall back rather than stranding every + // self-hosted deployment that has not upgraded yet. + // + // The code, not just the status, decides. 8.2.0 also answers 404 + // with `agent_not_found` when the agent itself is gone (it looks the + // target up before taking over), and a vanished agent is a real + // failure that must surface as one — falling back there would send a + // workspace key at a `requireAgentToken` route and report the 401 as + // "takeover unavailable on this engine", which is exactly the kind + // of misdirecting error this whole change exists to remove. + Err(RelayError::Api { status: 404, ref code, .. }) + if code != "agent_not_found" => + { + let rotated = relay + .rotate_agent_token(agent_name, self.api_key.clone()) + .await + .map_err(|error| RelaycastRegistrationError::Transport { + agent_name: agent_name.to_string(), + detail: format!( + "takeover unavailable on this engine and the legacy rotate fallback failed: {error}" + ), + })?; + AgentIdentityRecoveryResponse { + agent_id: existing_id, + name: agent_name.to_string(), + token: rotated.token, + audit_id: String::new(), + } + } + Err(error) => { + return Err(RelaycastRegistrationError::Transport { + agent_name: agent_name.to_string(), + detail: format!("takeover failed: {error}"), + }) + } + }; + + if response.token.trim().is_empty() { + return Err(RelaycastRegistrationError::MissingToken { + agent_name: agent_name.to_string(), + }); + } + + // Seed the credential the SDK cache would have held had registration + // succeeded. Without this, the next call re-registers, collides, takes + // over again and invalidates the token just handed out — and eventually + // meets the agent this broker itself brought online and refuses to + // impersonate it. + self.seed_agent_token(agent_name, &response.token); + + Ok(response.token) } /// Intent-aware token acquisition. `SpawnNew` preserves the existing @@ -293,7 +448,10 @@ impl RelaycastHttpClient { } match intent { - RegisterIntent::SpawnNew => registration + // `self.register_agent_token` rather than the SDK directly: it adds + // the audited takeover fallback that create-only registration now + // requires when the broker is reclaiming a name it owns. + RegisterIntent::SpawnNew => self .register_agent_token(trimmed_name, cli_hint) .await .map_err(ImpersonationAwareRegistrationError::Sdk), @@ -334,10 +492,21 @@ impl RelaycastHttpClient { // uses for a worker with no live credential — no // live token to invalidate, so the SDK's // register-or-rotate path is safe to run. - registration + // The probe has established there is no live credential + // to strand, which is exactly the condition that makes + // an audited takeover safe here. Reuse the agent the + // probe already fetched instead of looking it up again. + match registration .register_agent_token(trimmed_name, cli_hint) .await - .map_err(ImpersonationAwareRegistrationError::Sdk) + { + Ok(token) => Ok(token), + Err(AgentRegistrationError::AlreadyExists { .. }) => self + .take_over_agent_identity(trimmed_name, Some(&agent.id)) + .await + .map_err(ImpersonationAwareRegistrationError::Sdk), + Err(other) => Err(ImpersonationAwareRegistrationError::Sdk(other)), + } } Ok(Ok(agent)) => { // Neither a known-live nor a known-offline status — @@ -446,15 +615,16 @@ impl RelaycastHttpClient { } async fn registered_agent_client(&self) -> Result { - let registration = self - .registration - .as_ref() - .as_ref() - .context("SDK relay client not initialized")?; - registration - .registered_agent_client(&self.agent_name, Some(&self.default_cli)) + // Build the client from `register_agent_token` rather than the SDK's + // `registered_agent_client`, so this path inherits the audited takeover + // fallback create-only registration now requires. The SDK helper is + // exactly `register_agent_token` + `as_agent`, so this is the same + // composition with the collision case handled. + let token = self + .register_agent_token(&self.agent_name.clone(), Some(&self.default_cli.clone())) .await - .map_err(|error| anyhow::anyhow!("{error}")) + .map_err(|error| anyhow::anyhow!("{error}"))?; + AgentClient::new(token, self.base_url.clone()).map_err(|error| anyhow::anyhow!("{error}")) } /// Authenticate as `agent_name` rather than this broker's own identity. @@ -1676,7 +1846,7 @@ mod tests { /// broker restart. /// /// **Must-fire**: without the fix, this test hits - /// `POST /v1/agents/worker-a/rotate-token` (mounted below with a + /// `POST /v1/agents/worker-a/takeover` (mounted below with a /// deliberate 500 so a rotation would obviously fail the test) and /// invalidates the running worker. With the fix, the rotate mock is /// never called. @@ -1706,7 +1876,7 @@ mod tests { // well, instead of quietly returning a token that the test would // then have to explicitly assert against. let rotate = server.mock(|when, then| { - when.method(POST).path("/v1/agents/worker-a/rotate-token"); + when.method(POST).path("/v1/agents/worker-a/takeover"); then.status(500).json_body(json!({ "ok": false, "error": { "code": "must_not_rotate", "message": "must not rotate a live agent" } @@ -1757,7 +1927,7 @@ mod tests { then.status(500); }); let rotate = server.mock(|when, then| { - when.method(POST).path("/v1/agents/worker-a/rotate-token"); + when.method(POST).path("/v1/agents/worker-a/takeover"); then.status(500); }); @@ -1809,10 +1979,10 @@ mod tests { })); }); let rotate = server.mock(|when, then| { - when.method(POST).path("/v1/agents/worker-a/rotate-token"); + when.method(POST).path("/v1/agents/worker-a/takeover"); then.status(200).json_body(json!({ "ok": true, - "data": { "name": "worker-a", "token": "at_live_rotated_ok" } + "data": { "agent_id": "a_worker-a", "name": "worker-a", "token": "at_live_rotated_ok", "audit_id": "aud_1" } })); }); @@ -1864,7 +2034,7 @@ mod tests { }); // Any rotate is a regression. 500 makes it obvious. let rotate = server.mock(|when, then| { - when.method(POST).path("/v1/agents/worker-a/rotate-token"); + when.method(POST).path("/v1/agents/worker-a/takeover"); then.status(500).json_body(json!({ "ok": false, "error": { "code": "must_not_rotate", "message": "must not rotate a live worker" } @@ -1900,6 +2070,251 @@ mod tests { mark_read.assert_hits(0); } + /// Engines before 8.2.0 have no `/takeover`. The broker must fall back to + /// the legacy workspace-key rotate rather than stranding the agent — this + /// is the path the fleet e2e exercises, and every self-hosted deployment + /// that has not upgraded. + #[tokio::test] + async fn takeover_falls_back_to_legacy_rotate_on_older_engines() { + let server = MockServer::start(); + server.mock(|when, then| { + when.method(POST).path("/v1/agents"); + then.status(409).json_body(json!({ + "ok": false, + "error": { "code": "agent_already_exists", "message": "exists" } + })); + }); + server.mock(|when, then| { + when.method(GET).path("/v1/agents/worker-a"); + then.status(200).json_body(json!({ + "ok": true, + "data": { + "id": "a_worker-a", + "name": "worker-a", + "type": "agent", + "status": "offline", + "persona": null, + "metadata": {}, + "last_seen": "2026-08-16T20:00:00.000Z" + } + })); + }); + // Older engine: the route simply is not there. + let takeover = server.mock(|when, then| { + when.method(POST).path("/v1/agents/worker-a/takeover"); + then.status(404).json_body(json!({ + "ok": false, + "error": { "code": "not_found", "message": "no such route" } + })); + }); + // ...and rotate still accepts the workspace key there. + let legacy_rotate = server.mock(|when, then| { + when.method(POST) + .path("/v1/agents/worker-a/rotate-token") + .header("authorization", "Bearer rk_live_test"); + then.status(200).json_body(json!({ + "ok": true, + "data": { "name": "worker-a", "token": "at_live_legacy" } + })); + }); + + let client = + RelaycastHttpClient::new(Some(server.base_url()), "rk_live_test", "broker", "codex"); + + let token = client + .register_agent_token("worker-a", Some("codex")) + .await + .expect("an older engine must still be able to reclaim the name"); + assert_eq!(token, "at_live_legacy"); + takeover.assert_hits(1); + legacy_rotate.assert_hits(1); + } + + /// A 404 that means "this agent is gone" must NOT be treated as "this + /// engine is old". Falling back there would send a workspace key at a + /// `requireAgentToken` route and report the resulting 401 as an engine + /// capability problem — the misdirecting error this change exists to remove. + #[tokio::test] + async fn agent_not_found_does_not_trigger_the_legacy_fallback() { + let server = MockServer::start(); + server.mock(|when, then| { + when.method(POST).path("/v1/agents"); + then.status(409).json_body(json!({ + "ok": false, + "error": { "code": "agent_already_exists", "message": "exists" } + })); + }); + server.mock(|when, then| { + when.method(GET).path("/v1/agents/worker-a"); + then.status(200).json_body(json!({ + "ok": true, + "data": { + "id": "a_worker-a", + "name": "worker-a", + "type": "agent", + "status": "offline", + "persona": null, + "metadata": {}, + "last_seen": "2026-08-16T20:00:00.000Z" + } + })); + }); + // Modern engine, but the agent vanished between the lookup and the + // takeover — 404 with the agent-specific code. + server.mock(|when, then| { + when.method(POST).path("/v1/agents/worker-a/takeover"); + then.status(404).json_body(json!({ + "ok": false, + "error": { "code": "agent_not_found", "message": "Agent \"worker-a\" not found" } + })); + }); + // Must never be reached. + let legacy_rotate = server.mock(|when, then| { + when.method(POST).path("/v1/agents/worker-a/rotate-token"); + then.status(200).json_body(json!({ + "ok": true, + "data": { "name": "worker-a", "token": "at_live_should_not_be_used" } + })); + }); + + let client = + RelaycastHttpClient::new(Some(server.base_url()), "rk_live_test", "broker", "codex"); + + let error = client + .register_agent_token("worker-a", Some("codex")) + .await + .expect_err("a vanished agent is a real failure, not an old engine"); + let rendered = error.to_string(); + assert!( + rendered.contains("agent_not_found") || rendered.contains("not found"), + "the error must name the real cause; got: {rendered}" + ); + assert!( + !rendered.contains("takeover unavailable on this engine"), + "must not be misreported as an engine capability problem; got: {rendered}" + ); + legacy_rotate.assert_hits(0); + } + + /// Two concurrent cache-miss registrations for the same name must produce + /// exactly ONE takeover. Without singleflight both fire, and the second + /// response invalidates the token already returned to the first caller — + /// `expected_agent_id` cannot catch it, because takeover preserves the id. + #[tokio::test] + async fn concurrent_collisions_take_over_once() { + let server = MockServer::start(); + server.mock(|when, then| { + when.method(POST).path("/v1/agents"); + then.status(409).json_body(json!({ + "ok": false, + "error": { "code": "agent_already_exists", "message": "exists" } + })); + }); + server.mock(|when, then| { + when.method(GET).path("/v1/agents/worker-a"); + then.status(200).json_body(json!({ + "ok": true, + "data": { + "id": "a_worker-a", + "name": "worker-a", + "type": "agent", + "status": "offline", + "persona": null, + "metadata": {}, + "last_seen": "2026-08-16T20:00:00.000Z" + } + })); + }); + let takeover = server.mock(|when, then| { + when.method(POST).path("/v1/agents/worker-a/takeover"); + then.status(200) + // A slow response widens the window both callers would race in. + .delay(std::time::Duration::from_millis(150)) + .json_body(json!({ + "ok": true, + "data": { "agent_id": "a_worker-a", "name": "worker-a", "token": "at_live_taken", "audit_id": "aud_1" } + })); + }); + + let client = + RelaycastHttpClient::new(Some(server.base_url()), "rk_live_test", "broker", "codex"); + + let a = client.clone(); + let b = client.clone(); + let (first, second) = tokio::join!( + async move { a.register_agent_token("worker-a", Some("codex")).await }, + async move { b.register_agent_token("worker-a", Some("codex")).await }, + ); + + let first = first.expect("first concurrent caller succeeds"); + let second = second.expect("second concurrent caller succeeds"); + assert_eq!(first, "at_live_taken"); + assert_eq!( + second, first, + "both callers must hold the same live token, not one invalidated by the other" + ); + takeover.assert_hits(1); + } + + /// A takeover token must land in the registration cache. Without it the + /// next call re-registers, collides, takes over again, and invalidates the + /// token just handed out. + #[tokio::test] + async fn takeover_seeds_the_registration_cache() { + let server = MockServer::start(); + let register = server.mock(|when, then| { + when.method(POST).path("/v1/agents"); + then.status(409).json_body(json!({ + "ok": false, + "error": { "code": "agent_already_exists", "message": "exists" } + })); + }); + let lookup = server.mock(|when, then| { + when.method(GET).path("/v1/agents/worker-a"); + then.status(200).json_body(json!({ + "ok": true, + "data": { + "id": "a_worker-a", + "name": "worker-a", + "type": "agent", + "status": "offline", + "persona": null, + "metadata": {}, + "last_seen": "2026-08-16T20:00:00.000Z" + } + })); + }); + let takeover = server.mock(|when, then| { + when.method(POST).path("/v1/agents/worker-a/takeover"); + then.status(200).json_body(json!({ + "ok": true, + "data": { "agent_id": "a_worker-a", "name": "worker-a", "token": "at_live_taken", "audit_id": "aud_1" } + })); + }); + + let client = + RelaycastHttpClient::new(Some(server.base_url()), "rk_live_test", "broker", "codex"); + + let first = client + .register_agent_token("worker-a", Some("codex")) + .await + .expect("first call takes the name over"); + assert_eq!(first, "at_live_taken"); + + let second = client + .register_agent_token("worker-a", Some("codex")) + .await + .expect("second call is served from cache"); + assert_eq!( + second, "at_live_taken", + "the token must not be rotated out from under the caller" + ); + + register.assert_hits(1); + lookup.assert_hits(1); + takeover.assert_hits(1); + } + /// Must-not-fire: the existing `register_agent_token` API keeps its /// spawn-time behaviour (register or rotate on collision) so /// supervisor-driven worker restart at maintenance.rs:569 continues to @@ -1914,11 +2329,27 @@ mod tests { "error": { "code": "agent_already_exists", "message": "exists" } })); }); + // Takeover resolves the incumbent first so it can pin expected_agent_id. + let lookup = server.mock(|when, then| { + when.method(GET).path("/v1/agents/worker-a"); + then.status(200).json_body(json!({ + "ok": true, + "data": { + "id": "a_worker-a", + "name": "worker-a", + "type": "agent", + "status": "offline", + "persona": null, + "metadata": {}, + "last_seen": "2026-08-16T20:00:00.000Z" + } + })); + }); let rotate = server.mock(|when, then| { - when.method(POST).path("/v1/agents/worker-a/rotate-token"); + when.method(POST).path("/v1/agents/worker-a/takeover"); then.status(200).json_body(json!({ "ok": true, - "data": { "name": "worker-a", "token": "at_live_spawned" } + "data": { "agent_id": "a_worker-a", "name": "worker-a", "token": "at_live_spawned", "audit_id": "aud_1" } })); }); @@ -1930,6 +2361,7 @@ mod tests { .expect("spawn intent must rotate on collision"); assert_eq!(token, "at_live_spawned"); register.assert_hits(1); + lookup.assert_hits(1); rotate.assert_hits(1); } @@ -1962,7 +2394,7 @@ mod tests { })); }); let rotate = server.mock(|when, then| { - when.method(POST).path("/v1/agents/worker-a/rotate-token"); + when.method(POST).path("/v1/agents/worker-a/takeover"); then.status(500); }); @@ -2022,7 +2454,7 @@ mod tests { })); }); let rotate = server.mock(|when, then| { - when.method(POST).path("/v1/agents/worker-a/rotate-token"); + when.method(POST).path("/v1/agents/worker-a/takeover"); then.status(500); }); let register = server.mock(|when, then| { @@ -2062,11 +2494,27 @@ mod tests { #[tokio::test] async fn registered_agent_client_as_bypasses_impersonation_for_broker_own_identity() { let server = MockServer::start(); - // If the self-identity bypass regresses, this call would fire - // (broker probes itself) — 500 makes that loud. + // This endpoint now serves two different purposes: the presence probe + // (which the self-identity bypass must skip) and pinning + // `expected_agent_id` for an audited takeover (which is legitimate). + // A hit count alone can no longer tell them apart, so the guarantee is + // asserted by the call succeeding below: were the bypass to regress, + // the broker would refuse its own identity as a live-agent + // impersonation and the `expect` would fail. let presence = server.mock(|when, then| { when.method(GET).path("/v1/agents/broker"); - then.status(500); + then.status(200).json_body(json!({ + "ok": true, + "data": { + "id": "a_broker", + "name": "broker", + "type": "agent", + "status": "offline", + "persona": null, + "metadata": {}, + "last_seen": "2026-08-16T20:00:00.000Z" + } + })); }); let register = server.mock(|when, then| { when.method(POST).path("/v1/agents"); @@ -2076,10 +2524,10 @@ mod tests { })); }); let rotate = server.mock(|when, then| { - when.method(POST).path("/v1/agents/broker/rotate-token"); + when.method(POST).path("/v1/agents/broker/takeover"); then.status(200).json_body(json!({ "ok": true, - "data": { "name": "broker", "token": "at_live_broker_self" } + "data": { "agent_id": "a_broker", "name": "broker", "token": "at_live_broker_self", "audit_id": "aud_1" } })); }); @@ -2093,7 +2541,9 @@ mod tests { .registered_agent_client_as("broker", None) .await .expect("broker's own identity must not be refused as a live-agent impersonation"); - presence.assert_hits(0); + // One lookup, for the takeover's expected_agent_id — not a presence + // probe loop. + presence.assert_hits(1); register.assert_hits(1); rotate.assert_hits(1); }