From c5370fc34726e44a1d43e89c147a04cca480854c Mon Sep 17 00:00:00 2001 From: Proactive Runtime Bot Date: Tue, 8 Sep 2026 12:22:56 +0200 Subject: [PATCH] fix(attach): isolate cloud registration and back off on database overload --- .../compact_msichwe8jqzl_2026-09-08.json | 37 ++ .../compact_msichwe8jqzl_2026-09-08.md | 18 + .../2026-09/traj_pg9aglu4aqlw/compaction.json | 5 + .../2026-09/traj_pg9aglu4aqlw/summary.md | 31 + .../2026-09/traj_pg9aglu4aqlw/trajectory.json | 53 ++ CHANGELOG.md | 8 +- crates/broker/Cargo.toml | 1 + crates/broker/src/listen_api.rs | 52 ++ crates/broker/src/runtime/api.rs | 561 +++++++++++------- crates/broker/src/runtime/event_loop.rs | 10 + crates/broker/src/runtime/init.rs | 2 + crates/broker/src/runtime/tests.rs | 301 +++++++++- .../cli/src/cli/lib/attach-fleet-node.test.ts | 92 +++ packages/cli/src/cli/lib/attach-fleet-node.ts | 28 +- .../cli/src/cli/lib/attach-native.test.ts | 18 +- packages/cli/src/cli/lib/attach-native.ts | 15 +- 16 files changed, 1024 insertions(+), 208 deletions(-) create mode 100644 .agentworkforce/trajectories/compacted/compact_msichwe8jqzl_2026-09-08.json create mode 100644 .agentworkforce/trajectories/compacted/compact_msichwe8jqzl_2026-09-08.md create mode 100644 .agentworkforce/trajectories/completed/2026-09/traj_pg9aglu4aqlw/compaction.json create mode 100644 .agentworkforce/trajectories/completed/2026-09/traj_pg9aglu4aqlw/summary.md create mode 100644 .agentworkforce/trajectories/completed/2026-09/traj_pg9aglu4aqlw/trajectory.json diff --git a/.agentworkforce/trajectories/compacted/compact_msichwe8jqzl_2026-09-08.json b/.agentworkforce/trajectories/compacted/compact_msichwe8jqzl_2026-09-08.json new file mode 100644 index 0000000000..898f8ca261 --- /dev/null +++ b/.agentworkforce/trajectories/compacted/compact_msichwe8jqzl_2026-09-08.json @@ -0,0 +1,37 @@ +{ + "id": "compact_msichwe8jqzl", + "version": 1, + "type": "compacted", + "compactedAt": "2026-09-08T10:22:43.328Z", + "sourceTrajectories": [ + "traj_pg9aglu4aqlw" + ], + "dateRange": { + "start": "2026-09-08T10:04:45.586Z", + "end": "2026-09-08T10:21:50.230Z" + }, + "summary": { + "totalDecisions": 1, + "totalEvents": 1, + "uniqueAgents": [ + "default" + ] + }, + "decisionGroups": [ + { + "category": "api", + "decisions": [ + { + "question": "Move HTTP-spawn cloud preparation out of the broker actor and bound it", + "chosen": "Move HTTP-spawn cloud preparation out of the broker actor and bound it", + "reasoning": "Wrangler traced a 35.7s duplicate registration blocking local attach. Preserve node tokens, merge only completed identity/cursor state, reserve local names, reject pending releases, and skip launch when the caller is gone. Retry only explicit pre-allocation D1 overload for fleet attach.", + "fromTrajectory": "traj_pg9aglu4aqlw" + } + ] + } + ], + "keyLearnings": [], + "keyFindings": [], + "filesAffected": [], + "commits": [] +} \ No newline at end of file diff --git a/.agentworkforce/trajectories/compacted/compact_msichwe8jqzl_2026-09-08.md b/.agentworkforce/trajectories/compacted/compact_msichwe8jqzl_2026-09-08.md new file mode 100644 index 0000000000..428e189abe --- /dev/null +++ b/.agentworkforce/trajectories/compacted/compact_msichwe8jqzl_2026-09-08.md @@ -0,0 +1,18 @@ +# Trajectory Compaction: Sep 8, 2026 - Sep 8, 2026 + +## Summary +- Sessions: 1 +- Decisions: 1 +- Events: 1 +- Agents: default +- Files: 0 +- Commits: 0 + +## Api +- Move HTTP-spawn cloud preparation out of the broker actor and bound it -> Move HTTP-spawn cloud preparation out of the broker actor and bound it (traj_pg9aglu4aqlw) + +## Key Learnings +- None + +## Key Findings +- None \ No newline at end of file diff --git a/.agentworkforce/trajectories/completed/2026-09/traj_pg9aglu4aqlw/compaction.json b/.agentworkforce/trajectories/completed/2026-09/traj_pg9aglu4aqlw/compaction.json new file mode 100644 index 0000000000..1ab99ba63a --- /dev/null +++ b/.agentworkforce/trajectories/completed/2026-09/traj_pg9aglu4aqlw/compaction.json @@ -0,0 +1,5 @@ +{ + "trajectoryId": "traj_pg9aglu4aqlw", + "compactedInto": "compact_msichwe8jqzl", + "compactedAt": "2026-09-08T10:22:43.428Z" +} \ No newline at end of file diff --git a/.agentworkforce/trajectories/completed/2026-09/traj_pg9aglu4aqlw/summary.md b/.agentworkforce/trajectories/completed/2026-09/traj_pg9aglu4aqlw/summary.md new file mode 100644 index 0000000000..32917a905e --- /dev/null +++ b/.agentworkforce/trajectories/completed/2026-09/traj_pg9aglu4aqlw/summary.md @@ -0,0 +1,31 @@ +# Trajectory: Keep local attach responsive during slow cloud registration + +> **Status:** ✅ Completed +> **Confidence:** 90% +> **Started:** September 8, 2026 at 12:04 PM +> **Completed:** September 8, 2026 at 12:21 PM + +--- + +## Summary + +Isolated HTTP-spawn cloud registration from local attach; bounded registration and overload retries; accurate health/attach diagnostics. 1065 broker tests passed (4 ignored), 43 attach tests passed, CLI typecheck passed. Same real-pipe regression fails on baseline 11.10.4 and passes with fix. Production broker untouched; rollout pending. + +**Approach:** Standard approach + +--- + +## Key Decisions + +### Move HTTP-spawn cloud preparation out of the broker actor and bound it +- **Chose:** Move HTTP-spawn cloud preparation out of the broker actor and bound it +- **Reasoning:** Wrangler traced a 35.7s duplicate registration blocking local attach. Preserve node tokens, merge only completed identity/cursor state, reserve local names, reject pending releases, and skip launch when the caller is gone. Retry only explicit pre-allocation D1 overload for fleet attach. + +--- + +## Chapters + +### 1. Work +*Agent: default* + +- Move HTTP-spawn cloud preparation out of the broker actor and bound it: Move HTTP-spawn cloud preparation out of the broker actor and bound it diff --git a/.agentworkforce/trajectories/completed/2026-09/traj_pg9aglu4aqlw/trajectory.json b/.agentworkforce/trajectories/completed/2026-09/traj_pg9aglu4aqlw/trajectory.json new file mode 100644 index 0000000000..8dcf66afd5 --- /dev/null +++ b/.agentworkforce/trajectories/completed/2026-09/traj_pg9aglu4aqlw/trajectory.json @@ -0,0 +1,53 @@ +{ + "id": "traj_pg9aglu4aqlw", + "version": 1, + "task": { + "title": "Keep local attach responsive during slow cloud registration" + }, + "status": "completed", + "startedAt": "2026-09-08T10:04:45.586Z", + "completedAt": "2026-09-08T10:21:50.230Z", + "agents": [ + { + "name": "default", + "role": "lead", + "joinedAt": "2026-09-08T10:16:44.611Z" + } + ], + "chapters": [ + { + "id": "chap_k90zf0m041ko", + "title": "Work", + "agentName": "default", + "startedAt": "2026-09-08T10:16:44.611Z", + "endedAt": "2026-09-08T10:21:50.230Z", + "events": [ + { + "ts": 1788862604611, + "type": "decision", + "content": "Move HTTP-spawn cloud preparation out of the broker actor and bound it: Move HTTP-spawn cloud preparation out of the broker actor and bound it", + "raw": { + "question": "Move HTTP-spawn cloud preparation out of the broker actor and bound it", + "chosen": "Move HTTP-spawn cloud preparation out of the broker actor and bound it", + "alternatives": [], + "reasoning": "Wrangler traced a 35.7s duplicate registration blocking local attach. Preserve node tokens, merge only completed identity/cursor state, reserve local names, reject pending releases, and skip launch when the caller is gone. Retry only explicit pre-allocation D1 overload for fleet attach." + }, + "significance": "high" + } + ] + } + ], + "retrospective": { + "summary": "Isolated HTTP-spawn cloud registration from local attach; bounded registration and overload retries; accurate health/attach diagnostics. 1065 broker tests passed (4 ignored), 43 attach tests passed, CLI typecheck passed. Same real-pipe regression fails on baseline 11.10.4 and passes with fix. Production broker untouched; rollout pending.", + "approach": "Standard approach", + "confidence": 0.9 + }, + "commits": [], + "filesChanged": [], + "projectId": "AgentWorkforce/relay", + "tags": [], + "_trace": { + "startRef": "d754fff143c464c367ac743ccc3c8085bf5ef04d", + "endRef": "d754fff143c464c367ac743ccc3c8085bf5ef04d" + } +} \ No newline at end of file diff --git a/CHANGELOG.md b/CHANGELOG.md index 7cbd402e81..16f2d1da12 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,7 +5,13 @@ All notable changes to Agent Relay will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). -## [Unreleased] +## [Unreleased - Patch] + +### Fixed + +- Local agent attach, snapshots, and input remain responsive while HTTP spawns wait for cloud registration; registration now has bounded concurrency and an overall deadline. +- Fleet attach retries explicit database-overload rejections with server-directed backoff, and attach lookup failures explain that the agent may still be running. +- Broker health marks runtime status as unavailable instead of claiming zero agents when its status lookup times out. ## [11.10.4] - 2026-09-08 diff --git a/crates/broker/Cargo.toml b/crates/broker/Cargo.toml index 34572f2c58..68073d2d20 100644 --- a/crates/broker/Cargo.toml +++ b/crates/broker/Cargo.toml @@ -47,6 +47,7 @@ nix = { version = "0.30", features = ["signal", "process", "term", "fs"] } libc = "0.2" [dev-dependencies] +tokio = { version = "1.44", features = ["test-util"] } httpmock = "0.7" tempfile = "3.19" tower = { version = "0.5", features = ["util"] } diff --git a/crates/broker/src/listen_api.rs b/crates/broker/src/listen_api.rs index 6b1c7bea82..065dacbe7e 100644 --- a/crates/broker/src/listen_api.rs +++ b/crates/broker/src/listen_api.rs @@ -612,6 +612,29 @@ async fn listen_api_health( let mut payload = listen_api_health_payload(state.default_workspace_id, state.memberships); if let Some(status) = fetch_status_for_health(&state.tx).await { merge_status_into_health_payload(&mut payload, &status); + payload["runtimeResponsive"] = json!(true); + } else { + // Keep the existing public liveness HTTP contract: a busy actor must + // not induce a supervisor restart storm. But never present fabricated + // zero-agent/offline values as a successful runtime observation. + payload["runtimeResponsive"] = json!(false); + if payload["status"] == "ok" { + payload["status"] = json!("degraded"); + } + payload["runtimeStatusError"] = json!("runtime_status_unavailable"); + for key in [ + "agentCount", + "uptimeMs", + "pendingDeliveryCount", + "deadLetterCount", + "wsConnections", + "memoryMb", + "relaycastConnected", + "nodeConnected", + "nodeDelivery", + ] { + payload[key] = Value::Null; + } } axum::Json(payload) } @@ -3880,6 +3903,35 @@ mod auth_tests { .expect("request should succeed"); assert_eq!(response.status(), StatusCode::OK); + let payload = response_json(response).await; + assert_eq!(payload["runtimeResponsive"], false); + assert_eq!(payload["runtimeStatusError"], "runtime_status_unavailable"); + assert!(payload["agentCount"].is_null()); + assert!(payload["nodeConnected"].is_null()); + } + + #[tokio::test] + async fn health_route_reports_a_responsive_runtime_without_losing_counts() { + let (router, mut rx) = test_router(None); + let runtime = tokio::spawn(async move { + if let Some(super::ListenApiRequest::GetStatus { reply }) = rx.recv().await { + let _ = reply.send(Ok(json!({ "agent_count": 1 }))); + } + }); + let response = router + .oneshot( + Request::builder() + .uri("/health") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let payload = response_json(response).await; + assert_eq!(payload["runtimeResponsive"], true); + assert_eq!(payload["agentCount"], 1); + runtime.await.unwrap(); } #[test] diff --git a/crates/broker/src/runtime/api.rs b/crates/broker/src/runtime/api.rs index f09afb5fda..f94f1742ab 100644 --- a/crates/broker/src/runtime/api.rs +++ b/crates/broker/src/runtime/api.rs @@ -268,8 +268,353 @@ fn observer_token_filters_are_empty(filters: &ObserverTokenFilters) -> bool { && created_after.is_none() } +// Remote registration is polled alongside the runtime, never awaited inside +// its actor. Bound both concurrency and the complete cloud preparation phase. +const MAX_PENDING_API_SPAWNS: usize = 8; +const API_SPAWN_REGISTRATION_TIMEOUT: Duration = Duration::from_secs(20); + +pub(super) struct PreparedApiSpawn { + pub(super) request: ListenApiRequest, + registration: Result, +} + +struct ApiSpawnRegistration { + spec: AgentSpec, + worker_relay_key: Option, + fleet_registration: Option<( + crate::node_control::AgentRegistrationToken, + Option, + Option, + )>, + preregistration_warning: Option, +} + +async fn prepare_api_spawn_registration( + request: &ListenApiRequest, + relaycast_http: &RelaycastHttpClient, + fleet_control_tx: &mpsc::Sender, + fleet_node_name: &str, +) -> Result { + let ListenApiRequest::Spawn { + name, + cli, + transport, + model, + args, + registration_metadata, + channels, + cwd, + team, + shadow_of, + shadow_mode, + restart_policy, + harness_config, + agent_token, + .. + } = request + else { + unreachable!("only spawn requests are prepared") + }; + let ( + name, + cli, + transport, + model, + args, + registration_metadata, + channels, + cwd, + team, + shadow_of, + shadow_mode, + restart_policy, + harness_config, + agent_token, + ) = ( + name.clone(), + cli.clone(), + transport.clone(), + model.clone(), + args.clone(), + registration_metadata.clone(), + channels.clone(), + cwd.clone(), + team.clone(), + shadow_of.clone(), + shadow_mode.clone(), + restart_policy.clone(), + harness_config.clone(), + agent_token.clone(), + ); + // Merge only this result's identity/cursor, never a stale copy of the book. + let mut fleet_delivery_book = FleetDeliveryBook::default(); + let effective_channels = if channels.is_empty() { + default_spawn_channels() + } else { + channels.clone() + }; + let spec = match build_http_api_spawn_spec( + name.clone(), + cli.clone(), + transport, + model.clone(), + args, + effective_channels.clone(), + cwd, + team, + shadow_of, + shadow_mode, + *restart_policy, + harness_config, + ) { + Ok(spec) => spec, + Err(error) => { + return Err(error.to_string()); + } + }; + let mut preregistration_warning: Option = None; + // Caller-supplied agent_token is authoritative. In fleet mode it + // was minted by the node control connection, and the worker must + // receive that exact token before its harness starts. + // + // Otherwise bind the agent to this node via node-control + // `agent.register` — the same step the engine `action.invoke` + // spawn converges on — so the agent is born `via_node`-bound and + // delivery flows over /v1/node/ws. The minted token is injected + // as RELAY_AGENT_TOKEN (which also sets RELAY_SKIP_BOOTSTRAP) so + // the worker MCP never re-registers over HTTP. If node binding is + // unavailable, fall back to HTTP pre-registration so a tokenless + // node (e.g. mint failure) still spawns a working agent. + let mut fleet_registration = None; + let session_ref = super::fleet::fleet_initial_session_ref(&spec); + let worker_relay_key = if let Some(token) = agent_token { + seed_supplied_agent_token(relaycast_http, &name, &token); + match super::fleet::resolve_fleet_agent_token_identity( + relaycast_http, + &mut fleet_delivery_book, + &name, + &token, + ) + .await + { + Ok(registration) => { + fleet_registration = Some((registration, None, session_ref.clone())); + } + Err(error) => { + tracing::warn!( + worker = %name, + error = %error, + "could not resolve supplied agent token for reconnect inventory" + ); + } + } + Some(token) + } else { + // Derive the session ref from the resolved spec the same way + // the fleet/sidecar paths do, so an HTTP spawn carrying a + // `harnessConfig.session_id` registers as a resumable session + // rather than a fresh spawn. No invocation id exists on the + // HTTP path. + match super::fleet::register_node_agent_token( + fleet_control_tx, + &mut fleet_delivery_book, + name.as_str(), + None, + session_ref.clone(), + ) + .await + { + Ok(token) => { + tracing::info!( + worker = %name, + "bound agent to node via agent.register for HTTP spawn" + ); + super::fleet::spawn_declared_metadata_publish( + relaycast_http, + name.as_str(), + registration_metadata, + ); + let relay_key = token.token.clone(); + fleet_registration = Some((token, None, session_ref)); + Some(relay_key) + } + Err(node_error) => { + tracing::warn!( + worker = %name, + error = %node_error, + "node agent.register unavailable; falling back to HTTP pre-registration" + ); + // The ordinary cache-aware registration: it honours + // the SDK's cached token and rate-limit block, so a + // name already seeded by preflight or an earlier + // spawn is reused rather than re-created. Declared + // metadata is published separately, exactly as on + // the node path. + match retry_agent_registration(relaycast_http, &name, Some(&cli)).await { + Ok(token) => { + super::fleet::spawn_declared_metadata_publish( + relaycast_http, + name.as_str(), + registration_metadata, + ); + // HTTP registration alone leaves the agent + // without a node binding; the engine only + // delivers to `via_node` agents in node-only + // delivery. Bind it to this node so it is + // deliverable, surfacing a loud warning if the + // bind fails. + let bind_warning = + super::relaycast_events::bind_http_registered_agent_to_node( + relaycast_http, + fleet_node_name, + &name, + ) + .await; + if let Some(warning) = bind_warning { + preregistration_warning = Some(warning); + } else { + match super::fleet::resolve_fleet_agent_token_identity( + relaycast_http, + &mut fleet_delivery_book, + &name, + &token, + ) + .await + { + Ok(registration) => { + fleet_registration = + Some((registration, None, session_ref.clone())); + } + Err(error) => { + tracing::warn!( + worker = %name, + error = %error, + "could not resolve HTTP-registered agent for reconnect inventory" + ); + } + } + } + Some(token) + } + Err(RegRetryOutcome::RetryableExhausted(error)) => { + let message = format_worker_preregistration_error(&name, &error); + tracing::warn!( + worker = %name, + error = %error, + "continuing spawn without pre-registration after retries exhausted" + ); + preregistration_warning = Some(message); + None + } + Err(RegRetryOutcome::Fatal(error)) => { + return Err(format_worker_preregistration_error(&name, &error)); + } + } + } + } + }; + if let Some(token) = worker_relay_key.as_deref() { + // Node registration returns a token without populating the + // HTTP client's worker cache. Seed it before authenticating + // as the worker so channel reconciliation cannot rotate an + // already-live identity's token. + seed_supplied_agent_token(relaycast_http, &name, token); + if let Err(error) = relaycast_http + .ensure_agent_channels(&name, Some(&cli), &effective_channels) + .await + { + tracing::error!( + worker = %name, + channels = ?effective_channels, + error = %error, + "worker channel membership reconciliation failed" + ); + let membership_warning = + format!("worker channel membership was not fully reconciled: {error}"); + preregistration_warning = Some(match preregistration_warning.take() { + Some(existing) => format!("{existing}; {membership_warning}"), + None => membership_warning, + }); + } + } + + Ok(ApiSpawnRegistration { + spec, + worker_relay_key, + fleet_registration, + preregistration_warning, + }) +} + impl BrokerRuntime { pub(super) async fn handle_api_request(&mut self, req: ListenApiRequest) { + if let ListenApiRequest::Spawn { name, .. } = &req { + let error = + if self.workers.has_worker(name) || self.pending_api_spawn_names.contains(name) { + Some(format!( + "agent '{}' already exists or has a spawn in progress", + name + )) + } else if self.pending_api_spawns.len() >= MAX_PENDING_API_SPAWNS { + Some( + "spawn_admission_busy: too many registrations in progress; retry later" + .to_string(), + ) + } else { + None + }; + if let Some(error) = error { + if let ListenApiRequest::Spawn { reply, .. } = req { + let _ = reply.send(Err(error)); + } + return; + } + self.pending_api_spawn_names.insert(name.clone()); + let http = self.relaycast_http.clone(); + let control = self.fleet_control_tx.clone(); + let node = self.fleet_node_name.clone(); + self.pending_api_spawns.push(Box::pin(async move { + let registration = timeout(API_SPAWN_REGISTRATION_TIMEOUT, + prepare_api_spawn_registration(&req, &http, &control, &node)).await + .unwrap_or_else(|_| Err("spawn_registration_timeout: cloud registration exceeded 20s; no local worker was started; remote registration may have completed".to_string())); + PreparedApiSpawn { request: req, registration } + })); + return; + } + // Do not acknowledge release of a name that could still launch later. + if let ListenApiRequest::Release { name, .. } = &req { + if self.pending_api_spawn_names.contains(name) { + if let ListenApiRequest::Release { reply, .. } = req { + let _ = reply.send(Err("spawn_in_progress: wait for registration to finish before releasing this agent".to_string())); + } + return; + } + } + self.handle_ready_api_request(req, None).await; + } + + pub(super) async fn finish_api_spawn(&mut self, prepared: PreparedApiSpawn) { + if let ListenApiRequest::Spawn { name, reply, .. } = &prepared.request { + self.pending_api_spawn_names.remove(name); + if reply.is_closed() { + return; + } + if self.workers.has_worker(name) { + if let ListenApiRequest::Spawn { reply, name, .. } = prepared.request { + let _ = reply.send(Err(format!("agent '{}' already exists", name))); + } + return; + } + } + self.handle_ready_api_request(prepared.request, Some(prepared.registration)) + .await; + } + + async fn handle_ready_api_request( + &mut self, + req: ListenApiRequest, + prepared: Option>, + ) { let paths = &self.paths; let state = &mut self.state; let workspaces = &self.workspaces; @@ -319,219 +664,37 @@ impl BrokerRuntime { ListenApiRequest::Spawn { name, cli, - transport, - model, - args, task, - registration_metadata, - channels, - cwd, - team, - shadow_of, - shadow_mode, continue_from, idle_threshold_secs, exit_after_task, skip_relay_prompt, - restart_policy, - harness_config, - agent_token, agent_result_schema, replay_buffer, reply, + .. } => { - let effective_channels = if channels.is_empty() { - default_spawn_channels() - } else { - channels.clone() - }; - let spec = match build_http_api_spawn_spec( - name.clone(), - cli.clone(), - transport, - model.clone(), - args, - effective_channels.clone(), - cwd, - team, - shadow_of, - shadow_mode, - *restart_policy, - harness_config, - ) { - Ok(spec) => spec, + let ApiSpawnRegistration { + spec, + worker_relay_key, + mut fleet_registration, + preregistration_warning, + } = match prepared.expect("spawn must be prepared off the runtime actor") { + Ok(registration) => registration, Err(error) => { - let _ = reply.send(Err(error.to_string())); + let _ = reply.send(Err(error)); return; } }; - let mut preregistration_warning: Option = None; - // Caller-supplied agent_token is authoritative. In fleet mode it - // was minted by the node control connection, and the worker must - // receive that exact token before its harness starts. - // - // Otherwise bind the agent to this node via node-control - // `agent.register` — the same step the engine `action.invoke` - // spawn converges on — so the agent is born `via_node`-bound and - // delivery flows over /v1/node/ws. The minted token is injected - // as RELAY_AGENT_TOKEN (which also sets RELAY_SKIP_BOOTSTRAP) so - // the worker MCP never re-registers over HTTP. If node binding is - // unavailable, fall back to HTTP pre-registration so a tokenless - // node (e.g. mint failure) still spawns a working agent. - let mut fleet_registration = None; - let session_ref = super::fleet::fleet_initial_session_ref(&spec); - let worker_relay_key = if let Some(token) = agent_token { - seed_supplied_agent_token(relaycast_http, &name, &token); - match super::fleet::resolve_fleet_agent_token_identity( - relaycast_http, - fleet_delivery_book, - &name, - &token, - ) - .await - { - Ok(registration) => { - fleet_registration = Some((registration, None, session_ref.clone())); - } - Err(error) => { - tracing::warn!( - worker = %name, - error = %error, - "could not resolve supplied agent token for reconnect inventory" - ); - } - } - Some(token) - } else { - // Derive the session ref from the resolved spec the same way - // the fleet/sidecar paths do, so an HTTP spawn carrying a - // `harnessConfig.session_id` registers as a resumable session - // rather than a fresh spawn. No invocation id exists on the - // HTTP path. - match super::fleet::register_node_agent_token( - fleet_control_tx, - fleet_delivery_book, - name.as_str(), - None, - session_ref.clone(), - ) - .await - { - Ok(token) => { - tracing::info!( - worker = %name, - "bound agent to node via agent.register for HTTP spawn" - ); - super::fleet::spawn_declared_metadata_publish( - relaycast_http, - name.as_str(), - registration_metadata, - ); - let relay_key = token.token.clone(); - fleet_registration = Some((token, None, session_ref)); - Some(relay_key) - } - Err(node_error) => { - tracing::warn!( - worker = %name, - error = %node_error, - "node agent.register unavailable; falling back to HTTP pre-registration" - ); - // The ordinary cache-aware registration: it honours - // the SDK's cached token and rate-limit block, so a - // name already seeded by preflight or an earlier - // spawn is reused rather than re-created. Declared - // metadata is published separately, exactly as on - // the node path. - match retry_agent_registration(relaycast_http, &name, Some(&cli)).await - { - Ok(token) => { - super::fleet::spawn_declared_metadata_publish( - relaycast_http, - name.as_str(), - registration_metadata, - ); - // HTTP registration alone leaves the agent - // without a node binding; the engine only - // delivers to `via_node` agents in node-only - // delivery. Bind it to this node so it is - // deliverable, surfacing a loud warning if the - // bind fails. - let bind_warning = super::relaycast_events::bind_http_registered_agent_to_node( - relaycast_http, - fleet_node_name, - &name, - ) - .await; - if let Some(warning) = bind_warning { - preregistration_warning = Some(warning); - } else { - match super::fleet::resolve_fleet_agent_token_identity( - relaycast_http, - fleet_delivery_book, - &name, - &token, - ) - .await - { - Ok(registration) => { - fleet_registration = - Some((registration, None, session_ref.clone())); - } - Err(error) => { - tracing::warn!( - worker = %name, - error = %error, - "could not resolve HTTP-registered agent for reconnect inventory" - ); - } - } - } - Some(token) - } - Err(RegRetryOutcome::RetryableExhausted(error)) => { - let message = - format_worker_preregistration_error(&name, &error); - tracing::warn!( - worker = %name, - error = %error, - "continuing spawn without pre-registration after retries exhausted" - ); - preregistration_warning = Some(message); - None - } - Err(RegRetryOutcome::Fatal(error)) => { - let _ = reply.send(Err(format_worker_preregistration_error( - &name, &error, - ))); - return; - } - } - } - } - }; - if let Some(token) = worker_relay_key.as_deref() { - // Node registration returns a token without populating the - // HTTP client's worker cache. Seed it before authenticating - // as the worker so channel reconciliation cannot rotate an - // already-live identity's token. - seed_supplied_agent_token(relaycast_http, &name, token); - if let Err(error) = relaycast_http - .ensure_agent_channels(&name, Some(&cli), &effective_channels) - .await - { - tracing::error!( - worker = %name, - channels = ?effective_channels, - error = %error, - "worker channel membership reconciliation failed" + if let Some((token, _, _)) = &fleet_registration { + fleet_delivery_book + .bind_authoritative_identity(token.name.clone(), token.agent_id.clone()); + if let Some(up_to_seq) = token.delivery_ack_seq { + fleet_delivery_book.seed_cursor( + token.name.clone(), + token.agent_id.clone(), + up_to_seq, ); - let membership_warning = - format!("worker channel membership was not fully reconciled: {error}"); - preregistration_warning = Some(match preregistration_warning.take() { - Some(existing) => format!("{existing}; {membership_warning}"), - None => membership_warning, - }); } } diff --git a/crates/broker/src/runtime/event_loop.rs b/crates/broker/src/runtime/event_loop.rs index b6d9428a6f..334576c872 100644 --- a/crates/broker/src/runtime/event_loop.rs +++ b/crates/broker/src/runtime/event_loop.rs @@ -1,6 +1,7 @@ use super::*; use futures_util::future::{join, join_all}; +use futures_util::{future::BoxFuture, stream::FuturesUnordered, StreamExt}; /// Current PTY resize owner for a worker under the single-resizer policy. /// @@ -206,6 +207,9 @@ pub(crate) struct BrokerRuntime { pub(super) pty_observability: HashMap, pub(super) api_rx: mpsc::Receiver, pub(super) api_open: bool, + pub(super) pending_api_spawns: + FuturesUnordered>, + pub(super) pending_api_spawn_names: HashSet, pub(super) ws_inbound_rx: mpsc::Receiver, pub(super) relaycast_open: bool, pub(super) fleet_control_tx: mpsc::Sender, @@ -285,6 +289,7 @@ enum RuntimeEvent { Sigterm, Api(Box), ApiClosed, + ApiSpawnPrepared(Box), Stdin(std::io::Result>), Relaycast(Option), Fleet(Option), @@ -330,6 +335,8 @@ impl BrokerRuntime { Some(request) => RuntimeEvent::Api(Box::new(request)), None => RuntimeEvent::ApiClosed, }, + Some(prepared) = self.pending_api_spawns.next(), if !self.pending_api_spawns.is_empty() => + RuntimeEvent::ApiSpawnPrepared(Box::new(prepared)), result = self.sdk_lines.next_line(), if self.stdin_open => RuntimeEvent::Stdin(result), message = self.ws_inbound_rx.recv(), if self.relaycast_open => RuntimeEvent::Relaycast(message), event = self.fleet_event_rx.recv(), if self.fleet_control_open => RuntimeEvent::Fleet(event), @@ -355,6 +362,9 @@ impl BrokerRuntime { RuntimeEvent::ApiClosed => { self.api_open = false; } + RuntimeEvent::ApiSpawnPrepared(prepared) => { + self.finish_api_spawn(*prepared).await; + } RuntimeEvent::Stdin(result) => { if matches!(result, Ok(None) | Err(_)) { self.stdin_open = false; diff --git a/crates/broker/src/runtime/init.rs b/crates/broker/src/runtime/init.rs index 12058b3864..fb8240fb6b 100644 --- a/crates/broker/src/runtime/init.rs +++ b/crates/broker/src/runtime/init.rs @@ -676,6 +676,8 @@ pub(crate) async fn run_init(cmd: InitCommand, telemetry: TelemetryClient) -> Re hosted_agent_event_tx, pty_observability: HashMap::new(), api_rx, + pending_api_spawns: Default::default(), + pending_api_spawn_names: HashSet::new(), api_open: true, ws_inbound_rx, relaycast_open: true, diff --git a/crates/broker/src/runtime/tests.rs b/crates/broker/src/runtime/tests.rs index d667190bf9..cd2f1945f1 100644 --- a/crates/broker/src/runtime/tests.rs +++ b/crates/broker/src/runtime/tests.rs @@ -74,6 +74,10 @@ fn env_test_lock() -> &'static Mutex<()> { } async fn make_worker_registry_with_worker(name: &str) -> WorkerRegistry { + make_worker_registry_with_worker_stdout(name, false).await +} + +async fn make_worker_registry_with_worker_stdout(name: &str, piped: bool) -> WorkerRegistry { let (tx, _rx) = mpsc::channel::(16); let mut registry = WorkerRegistry::new( tx.clone(), @@ -83,7 +87,7 @@ async fn make_worker_registry_with_worker(name: &str) -> WorkerRegistry { ); let mut child = tokio::process::Command::new("cat") .stdin(Stdio::piped()) - .stdout(Stdio::null()) + .stdout(if piped { Stdio::piped() } else { Stdio::null() }) .stderr(Stdio::null()) .spawn() .expect("test worker process should spawn"); @@ -262,6 +266,8 @@ fn worker_event_runtime_fixture( hosted_agent_event_tx, pty_observability: HashMap::new(), api_rx, + pending_api_spawns: Default::default(), + pending_api_spawn_names: HashSet::new(), api_open: true, ws_inbound_rx, relaycast_open: true, @@ -339,6 +345,299 @@ fn delivery_lifecycle_worker_event( } } +fn attach_regression_spawn( + name: &str, +) -> ( + crate::listen_api::ListenApiRequest, + tokio::sync::oneshot::Receiver>, +) { + let (reply, rx) = tokio::sync::oneshot::channel(); + ( + crate::listen_api::ListenApiRequest::Spawn { + name: name.into(), + cli: "claude".to_string(), + transport: None, + model: None, + args: vec![], + task: None, + registration_metadata: Default::default(), + channels: vec![], + cwd: None, + team: None, + shadow_of: None, + shadow_mode: None, + continue_from: None, + idle_threshold_secs: None, + exit_after_task: false, + skip_relay_prompt: true, + restart_policy: Box::new(None), + harness_config: None, + agent_token: None, + agent_result_schema: None, + replay_buffer: crate::replay_buffer::ReplayBuffer::new(16), + reply, + }, + rx, + ) +} + +#[tokio::test(start_paused = true)] +async fn api_registration_has_an_overall_deadline_and_frees_admission() { + use futures_util::StreamExt; + let registry = make_worker_registry_with_worker("existing-agent").await; + let mut fixture = worker_event_runtime_fixture(registry, HashMap::new()); + let (request, reply) = attach_regression_spawn("timeout-agent"); + fixture.runtime.handle_api_request(request).await; + // Keep fleet_control_rx alive and its registration reply unanswered. + // Tokio's test clock advances to the phase deadline without a real sleep. + let started = tokio::time::Instant::now(); + let prepared = fixture.runtime.pending_api_spawns.next().await.unwrap(); + assert_eq!(started.elapsed(), Duration::from_secs(20)); + fixture.runtime.finish_api_spawn(prepared).await; + assert!(reply + .await + .unwrap() + .unwrap_err() + .contains("spawn_registration_timeout")); + assert!(!fixture.runtime.workers.has_worker("timeout-agent")); + assert!(fixture.runtime.pending_api_spawn_names.is_empty()); + cleanup_worker_registry(fixture.runtime.workers).await; +} + +#[tokio::test] +async fn api_registration_does_not_launch_after_the_caller_disconnects() { + use futures_util::StreamExt; + let server = httpmock::MockServer::start_async().await; + let _membership = server + .mock_async(|_when, then| { + then.status(401).json_body(json!({"ok": false, "error": { + "code": "unauthorized", "message": "test channel reconciliation unavailable" + }})); + }) + .await; + let registry = make_worker_registry_with_worker("existing-agent").await; + let mut fixture = worker_event_runtime_fixture(registry, HashMap::new()); + fixture.runtime.relaycast_http = + RelaycastHttpClient::new(Some(server.base_url()), "rk_live_test", "broker", "codex"); + let (request, reply) = attach_regression_spawn("abandoned-agent"); + fixture.runtime.handle_api_request(request).await; + drop(reply); + let register = async { + let Some(FleetControlCommand::RegisterAgent { reply, .. }) = + fixture.fleet_control_rx.recv().await + else { + panic!("expected registration") + }; + reply + .send(Ok(crate::node_control::AgentRegistrationToken { + name: "abandoned-agent".to_string(), + agent_id: "abandoned-id".to_string(), + token: "at_test".to_string(), + delivery_ack_seq: Some(42), + })) + .unwrap(); + }; + let (prepared, ()) = tokio::join!(fixture.runtime.pending_api_spawns.next(), register); + fixture.runtime.finish_api_spawn(prepared.unwrap()).await; + assert!(!fixture.runtime.workers.has_worker("abandoned-agent")); + assert!(fixture.runtime.pending_api_spawn_names.is_empty()); + assert!(fixture + .runtime + .fleet_delivery_book + .active_agent_id("abandoned-agent") + .is_none()); + assert!(fixture.runtime.fleet_inventory.is_empty()); + cleanup_worker_registry(fixture.runtime.workers).await; +} + +// Exercise the actual runtime select loop, not a stand-in task that just +// sleeps. Hold the node registration's reply indefinitely and prove the +// attach bootstrap APIs still answer before that registration completes. +#[tokio::test] +async fn slow_api_registration_does_not_block_attach_or_shutdown() { + use crate::listen_api::ListenApiRequest; + use tokio::io::AsyncBufReadExt; + let mut registry = make_worker_registry_with_worker_stdout("existing-agent", true).await; + let worker = registry.workers.get_mut("existing-agent").unwrap(); + let generation = worker.generation; + let stdout = worker.child.stdout.take().unwrap(); + let mut fixture = worker_event_runtime_fixture(registry, HashMap::new()); + let (worker_tx, worker_rx) = mpsc::channel(16); + fixture.runtime.worker_event_rx = worker_rx; + // A real child pipe echoes the broker's frames; turn them into worker + // responses so both snapshot and input traverse the actual request/ack lane. + let responder = tokio::spawn(async move { + let mut lines = tokio::io::BufReader::new(stdout).lines(); + while let Ok(Some(line)) = lines.next_line().await { + let frame: Value = serde_json::from_str(&line).unwrap(); + let kind = frame["type"].as_str().unwrap(); + if frame.get("request_id").is_none() { + continue; + } + let _ = worker_tx + .send(WorkerEvent::Message { + name: "existing-agent".into(), + generation, + value: json!({ + "type": format!("{kind}_response"), + "request_id": frame["request_id"], + "payload": { "screen": "still alive", "bytes_written": 4 } + }), + }) + .await; + } + }); + let (api_tx, api_rx) = mpsc::channel(16); + fixture.runtime.api_rx = api_rx; + fixture.runtime.reap_tick.reset(); + fixture.runtime.lease_check.reset(); + let (request, mut spawn_reply) = attach_regression_spawn("slow-new-agent"); + api_tx.send(request).await.unwrap(); + let exercise = async { + let registration = + tokio::time::timeout(Duration::from_secs(2), fixture.fleet_control_rx.recv()) + .await + .unwrap() + .unwrap(); + let FleetControlCommand::RegisterAgent { + reply: registration_reply, + .. + } = registration + else { + panic!("expected real node registration"); + }; + let (reply, rx) = tokio::sync::oneshot::channel(); + api_tx.send(ListenApiRequest::List { reply }).await.unwrap(); + let list = tokio::time::timeout(Duration::from_millis(500), rx) + .await + .expect("slow registration blocked attach lookup") + .unwrap() + .unwrap(); + assert_eq!(list["agents"][0]["name"], "existing-agent"); + + let (reply, rx) = tokio::sync::oneshot::channel(); + api_tx + .send(ListenApiRequest::WorkerRequest { + name: "existing-agent".into(), + kind: "snapshot_pty".to_string(), + payload: json!({}), + timeout: Duration::from_secs(5), + reply, + }) + .await + .unwrap(); + let snapshot = tokio::time::timeout(Duration::from_millis(500), rx) + .await + .expect("slow registration blocked snapshot") + .unwrap() + .unwrap(); + assert_eq!(snapshot["screen"], "still alive"); + let (reply, rx) = tokio::sync::oneshot::channel(); + api_tx + .send(ListenApiRequest::SendInput { + name: "existing-agent".into(), + data: "test".to_string(), + reply, + }) + .await + .unwrap(); + let input = tokio::time::timeout(Duration::from_millis(500), rx) + .await + .expect("slow registration blocked input acknowledgement") + .unwrap() + .unwrap(); + assert_eq!(input["bytes_written"], 4); + + let (reply, rx) = tokio::sync::oneshot::channel(); + api_tx + .send(ListenApiRequest::GetInboundDeliveryMode { + name: "existing-agent".into(), + reply, + }) + .await + .unwrap(); + assert!(tokio::time::timeout(Duration::from_millis(500), rx) + .await + .expect("slow registration blocked delivery-mode lookup") + .unwrap() + .is_ok()); + let (reply, rx) = tokio::sync::oneshot::channel(); + api_tx + .send(ListenApiRequest::GetStatus { reply }) + .await + .unwrap(); + assert!(tokio::time::timeout(Duration::from_millis(500), rx) + .await + .expect("slow registration blocked health status") + .unwrap() + .is_ok()); + assert!(matches!( + spawn_reply.try_recv(), + Err(tokio::sync::oneshot::error::TryRecvError::Empty) + )); + let (reply, rx) = tokio::sync::oneshot::channel(); + api_tx + .send(ListenApiRequest::Shutdown { reply }) + .await + .unwrap(); + tokio::time::timeout(Duration::from_millis(500), rx) + .await + .expect("slow registration blocked shutdown") + .unwrap() + .unwrap(); + // The reply stays live until after shutdown was acknowledged; dropping + // it earlier would let the old blocking path fall back and hide the bug. + drop(registration_reply); + }; + let (result, ()) = tokio::join!(fixture.runtime.run(), exercise); + result.unwrap(); + responder.await.unwrap(); +} + +#[tokio::test] +async fn api_registration_reserves_names_and_bounds_admission() { + use crate::listen_api::ListenApiRequest; + let registry = make_worker_registry_with_worker("existing-agent").await; + let mut fixture = worker_event_runtime_fixture(registry, HashMap::new()); + let (existing, existing_reply) = attach_regression_spawn("existing-agent"); + fixture.runtime.handle_api_request(existing).await; + assert!(existing_reply + .await + .unwrap() + .unwrap_err() + .contains("already exists")); + assert!(fixture.runtime.pending_api_spawns.is_empty()); + let mut replies = vec![]; + for i in 0..8 { + let (request, reply) = attach_regression_spawn(&format!("pending-{i}")); + fixture.runtime.handle_api_request(request).await; + replies.push(reply); + } + let (duplicate, reply) = attach_regression_spawn("pending-0"); + fixture.runtime.handle_api_request(duplicate).await; + assert!(reply.await.unwrap().unwrap_err().contains("already exists")); + let (extra, reply) = attach_regression_spawn("over-capacity"); + fixture.runtime.handle_api_request(extra).await; + assert!(reply + .await + .unwrap() + .unwrap_err() + .contains("spawn_admission_busy")); + let (reply, rx) = tokio::sync::oneshot::channel(); + fixture + .runtime + .handle_api_request(ListenApiRequest::Release { + name: "pending-0".into(), + reason: None, + reply, + }) + .await; + assert!(rx.await.unwrap().unwrap_err().contains("spawn_in_progress")); + assert_eq!(fixture.runtime.pending_api_spawns.len(), 8); + assert_eq!(fixture.runtime.pending_api_spawn_names.len(), 8); + cleanup_worker_registry(fixture.runtime.workers).await; +} + fn inbound_ctx<'a>(event_id: &'a str) -> InboundContext<'a> { InboundContext { from: "Alice", diff --git a/packages/cli/src/cli/lib/attach-fleet-node.test.ts b/packages/cli/src/cli/lib/attach-fleet-node.test.ts index 252894fc61..aef9d98673 100644 --- a/packages/cli/src/cli/lib/attach-fleet-node.test.ts +++ b/packages/cli/src/cli/lib/attach-fleet-node.test.ts @@ -141,6 +141,98 @@ describe('startFleetNodeAttachProxy terminal-session request retries', () => { } }); + it('recovers from a database overload and honors the server Retry-After', async () => { + const remote = await startFakeRemote(); + cleanup.push(remote.close); + const success = fakeTicketFetch(remote.url); + let calls = 0; + const retryDelays: number[] = []; + const proxy = await startFleetNodeAttachProxy({ + agent: 'running-agent', + node: 'finn-mini', + mode: 'drive', + baseUrl: 'https://fake.example', + workspaceKey: 'wk', + fetch: (async (input, init) => { + calls += 1; + if (calls === 1) { + return Response.json( + { ok: false, error: { code: 'database_overloaded', message: 'D1 overloaded' } }, + { status: 503, headers: { 'Retry-After': '8' } } + ); + } + return success(input, init); + }) as typeof globalThis.fetch, + sessionRequest: { + sleep: async (ms) => { + retryDelays.push(ms); + }, + }, + }); + cleanup.push(proxy.close); + expect(calls).toBe(2); + expect(retryDelays).toEqual([8_000]); + sendReady(await remote.nextConnection()); + }); + + it.each([ + ['database_overloaded', 503, 5], + ['internal_error', 503, 1], + ['database_overloaded', 401, 1], + ])('bounds retries for %s HTTP %s to %s attempts', async (code, status, expectedCalls) => { + let calls = 0; + const delays: number[] = []; + const failure = await startFleetNodeAttachProxy({ + agent: 'running-agent', + node: 'finn-mini', + mode: 'drive', + baseUrl: 'https://fake.example', + workspaceKey: 'wk', + fetch: (async () => { + calls += 1; + return Response.json( + { ok: false, error: { code, message: 'unavailable' } }, + { status, headers: { 'Retry-After': 'invalid' } } + ); + }) as typeof globalThis.fetch, + sessionRequest: { + sleep: async (ms) => { + delays.push(ms); + }, + }, + }).catch((error: Error) => error); + expect(failure).toBeInstanceOf(Error); + expect(calls).toBe(expectedCalls); + expect(delays.length).toBe(expectedCalls - 1); + if (expectedCalls === 5) { + expect(delays).toEqual([6_000, 7_200, 8_400, 9_600]); + expect((failure as Error).message).toContain('does not mean the agent has stopped'); + } + }); + + it('does not retry early when Retry-After exceeds the overall request budget', async () => { + vi.useFakeTimers(); + let calls = 0; + const attempt = startFleetNodeAttachProxy({ + agent: 'running-agent', + node: 'finn-mini', + mode: 'view', + baseUrl: 'https://fake.example', + workspaceKey: 'wk', + fetch: (async () => { + calls += 1; + return Response.json( + { ok: false, error: { code: 'database_overloaded', message: 'overloaded' } }, + { status: 503, headers: { 'Retry-After': '120' } } + ); + }) as typeof globalThis.fetch, + }).catch((error: Error) => error); + await vi.advanceTimersByTimeAsync(90_000); + const failure = await attempt; + expect(calls).toBe(1); + expect((failure as Error).message).toContain('overall budget was exhausted'); + }); + // MUST FIRE: before relay#1571 the first 503 escaped directly and this // never reached the successful second response. it('retries a transient node_unreachable response before opening the terminal', async () => { diff --git a/packages/cli/src/cli/lib/attach-fleet-node.ts b/packages/cli/src/cli/lib/attach-fleet-node.ts index a6375fd40a..dacb8815c5 100644 --- a/packages/cli/src/cli/lib/attach-fleet-node.ts +++ b/packages/cli/src/cli/lib/attach-fleet-node.ts @@ -262,6 +262,9 @@ function retryDelayBudgetMs(attempts: number, initialMs: number, maxMs: number): } function terminalSessionFailureSummary(error: TerminalSessionAttemptError): string { + if (error.code === 'database_overloaded') { + return 'The fleet control-plane database is overloaded; this does not mean the agent has stopped'; + } if (error.code === 'node_not_found') return 'Control-plane node lookup found no matching record'; if (error.code === 'node_unreachable') { return /no terminal transport/i.test(error.message) @@ -276,12 +279,26 @@ function terminalSessionFailureSummary(error: TerminalSessionAttemptError): stri return 'The terminal-session request was rejected'; } -function isRetryableTerminalSessionFailure(code: string | undefined): boolean { +function isRetryableTerminalSessionFailure(code: string | undefined, status: number): boolean { // These structured responses are emitted before a session is returned, so // retrying cannot duplicate a successful allocation. A fetch timeout, // network failure, or unclassified 5xx is different: this POST may already // have completed server-side, and retrying it could create a second session. - return code === 'node_unreachable' || code === 'terminal_session_unavailable'; + // The gateway performs D1 auth/node/agent reads BEFORE terminal/create. + // Its explicit D1 overload rejection is therefore safe to replay. Do not + // generalize this to other endpoints or unknown 5xx/transport failures. + return ( + code === 'node_unreachable' || + code === 'terminal_session_unavailable' || + (status === 503 && code === 'database_overloaded') + ); +} + +function retryAfterMs(value: string | null | undefined): number { + if (!value) return 0; + const seconds = /^\d+(?:\.\d+)?$/.test(value.trim()) ? Number(value) : NaN; + const delay = Number.isFinite(seconds) ? seconds * 1_000 : Date.parse(value) - Date.now(); + return Number.isFinite(delay) && delay > 0 ? delay : 0; } /** Start a broker-compatible loopback proxy for one remote terminal session. */ @@ -307,6 +324,7 @@ export async function startFleetNodeAttachProxy( let sessionRequestAttempts = 0; let sessionRequestBudgetExhaustedBetweenAttempts = false; let lastSessionError: TerminalSessionAttemptError | undefined; + let serverRetryAfterMs = 0; const sessionResult = await collectWithRetry( 'terminal session request', async () => { @@ -333,6 +351,7 @@ export async function startFleetNodeAttachProxy( throw exhausted; } sessionRequestAttempts += 1; + serverRetryAfterMs = 0; const controller = new AbortController(); let timedOut = false; const attemptTimeoutMs = Math.min(sessionRequestTimeoutMs, remainingRequestBudgetMs); @@ -373,13 +392,14 @@ export async function startFleetNodeAttachProxy( const resumeToken = ticketPayload.data?.resume_token; if (!ticketResponse.ok || !terminalUrl || !sessionId || !resumeToken) { const code = ticketPayload.error?.code; + serverRetryAfterMs = retryAfterMs(ticketResponse.headers?.get('retry-after')); const message = ticketPayload.error?.message ?? `terminal session request failed (HTTP ${ticketResponse.status})`; lastSessionError = new TerminalSessionAttemptError( message, code, ticketResponse.status, - isRetryableTerminalSessionFailure(code), + isRetryableTerminalSessionFailure(code, ticketResponse.status), false ); throw lastSessionError; @@ -392,7 +412,7 @@ export async function startFleetNodeAttachProxy( sleep: async (delayMs) => { const remainingRequestBudgetMs = sessionRequestDeadline - Date.now(); if (remainingRequestBudgetMs <= 0) return; - await sessionRequestSleep(Math.min(delayMs, remainingRequestBudgetMs)); + await sessionRequestSleep(Math.min(Math.max(delayMs, serverRetryAfterMs), remainingRequestBudgetMs)); }, shouldRetry: (error) => error instanceof TerminalSessionAttemptError && error.retryable, } diff --git a/packages/cli/src/cli/lib/attach-native.test.ts b/packages/cli/src/cli/lib/attach-native.test.ts index 2daf5fc2c9..9180691665 100644 --- a/packages/cli/src/cli/lib/attach-native.test.ts +++ b/packages/cli/src/cli/lib/attach-native.test.ts @@ -1,7 +1,12 @@ import { describe, expect, it } from 'vitest'; import { EventEmitter } from 'node:events'; -import { attachNative, renderNativeHarnessDiagnostic, renderAgentEvent } from './attach-native.js'; +import { + attachNative, + isNativeHarness, + renderNativeHarnessDiagnostic, + renderAgentEvent, +} from './attach-native.js'; const envelope = (kind: string, fields: Record = {}) => ({ @@ -13,6 +18,17 @@ const envelope = (kind: string, fields: Record = {}) => }) as never; describe('native harness attach rendering', () => { + it('explains a broker lookup timeout without reporting the running agent as dead', async () => { + await expect( + isNativeHarness( + 'session-thread-rollout', + { brokerUrl: 'http://127.0.0.1:1', apiKey: 'test-key' }, + (async () => { + throw new DOMException('The operation timed out.', 'TimeoutError'); + }) as typeof fetch + ) + ).rejects.toThrow(/broker could not complete the agent lookup.*agent may still be running/); + }); it('rejects passthrough without probing a terminal or broker', async () => { let error = ''; await expect( diff --git a/packages/cli/src/cli/lib/attach-native.ts b/packages/cli/src/cli/lib/attach-native.ts index 7a00d2e47b..7c0551209a 100644 --- a/packages/cli/src/cli/lib/attach-native.ts +++ b/packages/cli/src/cli/lib/attach-native.ts @@ -305,6 +305,17 @@ export async function isNativeHarness( }); if (!connection) return false; const client = createBrokerClient(connection, fetchFn); - const agent = (await client.listAgents()).find((candidate) => candidate.name === name); - return agent?.runtime_kind === 'native'; + try { + const agent = (await client.listAgents()).find((candidate) => candidate.name === name); + return agent?.runtime_kind === 'native'; + } catch (error) { + throw new Error( + `Cannot attach to ${JSON.stringify(name)}: the broker could not complete the agent lookup. ` + + 'The agent may still be running; this lookup failure does not establish that it needs a restart. ' + + `Check broker responsiveness and retry. Cause: ${describeError(error)}`, + { cause: error } + ); + } finally { + client.disconnect(); + } }