diff --git a/CHANGELOG.md b/CHANGELOG.md index db70a371a..b41024890 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added - `teams.json` agents accept a per-agent `model` field when `up --spawn` starts them; an explicit `--model` or `-m` inside `cli` still wins. +- The cloud dashboard can link a recorded Claude Code or Codex session to the fleet worker running it, so you can message that worker from the session page. After each spawn and supervised restart, the broker records the worker's provider session as `session_id` and `session_kind` in its Relaycast agent metadata. ### Changed diff --git a/crates/broker/src/relaycast/mod.rs b/crates/broker/src/relaycast/mod.rs index 38b6c2a16..387abed57 100644 --- a/crates/broker/src/relaycast/mod.rs +++ b/crates/broker/src/relaycast/mod.rs @@ -27,6 +27,6 @@ pub(crate) use workspace::{ pub(crate) use ws::{ format_worker_preregistration_error, register_new_spawn_identity, registration_retry_after_secs, retry_agent_registration, retry_agent_registration_create_only, - RegRetryOutcome, RelaycastHttpClient, RelaycastRegistrationError, WsControl, - MAX_AGENT_REGISTRATION_ELAPSED, MAX_AGENT_REGISTRATION_OUTER_TIMEOUT, + RegRetryOutcome, RelaycastHttpClient, RelaycastRegistrationError, SessionMetadataPublish, + WsControl, MAX_AGENT_REGISTRATION_ELAPSED, MAX_AGENT_REGISTRATION_OUTER_TIMEOUT, }; diff --git a/crates/broker/src/relaycast/ws.rs b/crates/broker/src/relaycast/ws.rs index 4ec900f7a..136e5d746 100644 --- a/crates/broker/src/relaycast/ws.rs +++ b/crates/broker/src/relaycast/ws.rs @@ -1,6 +1,6 @@ use std::{ collections::{BTreeSet, HashMap}, - sync::atomic::{AtomicU32, Ordering}, + sync::atomic::{AtomicU32, AtomicU64, Ordering}, sync::{Arc, Mutex as StdMutex}, time::{Duration, Instant}, }; @@ -64,10 +64,61 @@ pub struct RelaycastHttpClient { /// cache-miss registrations from both taking the name over — the second /// response would invalidate the token handed to the first caller. takeover_locks: Arc>>>>, + /// One fence per agent name, ordering this broker's metadata PATCHes for + /// that name. See [`MetadataPublishFence`]. + metadata_fences: Arc>>>, pub agent_name: String, pub default_cli: String, } +/// Orders the metadata PATCHes this broker sends for one agent name. +/// +/// The engine merges `PATCH /v1/agents/:name` by reading the stored metadata +/// and writing the combined object back, with no per-agent lock or version +/// check. Two of our PATCHes in flight at once for the same name (declared +/// keys and the session id, both detached from the spawn) can therefore each +/// read the pre-update record and the later write drops the other's keys. +/// `lock` is held across every metadata PATCH for the name, so they land one +/// after another. +/// +/// `latest_session` fences session-id publishes to the most recent spawn of +/// the name: every spawn claims a new sequence number before its publish task +/// starts, and a publish whose number is no longer the latest is dropped, so a +/// released worker's late publish cannot label a replacement that reuses its +/// name with the old session. +#[derive(Default)] +struct MetadataPublishFence { + lock: tokio::sync::Mutex<()>, + latest_session: AtomicU64, +} + +/// A spawn's claim on publishing its session id for an agent name, taken +/// synchronously at spawn time by [`RelaycastHttpClient::claim_session_metadata`]. +/// A later claim for the same name supersedes it. +pub struct SessionMetadataClaim { + agent_name: String, + fence: Arc, + sequence: u64, +} + +impl SessionMetadataClaim { + pub fn agent_name(&self) -> &str { + &self.agent_name + } +} + +/// What [`RelaycastHttpClient::publish_session_metadata`] did. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum SessionMetadataPublish { + /// The PATCH was sent and accepted. + Published, + /// The session id was blank; nothing was sent. + NothingToPublish, + /// A later spawn of the same name claimed the session metadata before this + /// publish ran; nothing was sent. + Superseded, +} + pub type RelaycastRegistrationError = AgentRegistrationError; pub type RegRetryOutcome = AgentRegistrationRetryOutcome; pub(crate) use relaycast::registration_retry_after_secs; @@ -203,6 +254,7 @@ impl RelaycastHttpClient { relay: Arc::new(None), registration: Arc::new(None), takeover_locks: Arc::new(StdMutex::new(HashMap::new())), + metadata_fences: Arc::new(StdMutex::new(HashMap::new())), agent_name: agent_name.into(), default_cli: String::new(), } @@ -229,6 +281,7 @@ impl RelaycastHttpClient { relay, registration, takeover_locks: Arc::new(StdMutex::new(HashMap::new())), + metadata_fences: Arc::new(StdMutex::new(HashMap::new())), agent_name: agent_name.into(), default_cli, } @@ -817,13 +870,103 @@ impl RelaycastHttpClient { if declared_metadata.is_empty() { return Ok(()); } + self.merge_agent_metadata(name, declared_metadata).await + } + + /// Claim the right to publish a spawn's session id for `agent_name`. + /// + /// Call synchronously when the spawn succeeds, before detaching the + /// publish, so claims follow spawn order. The claim supersedes every + /// earlier one for the name, including a still-pending publish from a + /// released worker whose name this spawn reuses. + pub fn claim_session_metadata( + &self, + agent_name: &str, + ) -> std::result::Result { + let name = agent_name.trim(); + if name.is_empty() { + return Err(RelaycastRegistrationError::InvalidAgentName); + } + let fence = self.metadata_fence(name); + let sequence = fence.latest_session.fetch_add(1, Ordering::SeqCst) + 1; + Ok(SessionMetadataClaim { + agent_name: name.to_string(), + fence, + sequence, + }) + } + + /// Publish the worker's provider session identity (`session_id`, and + /// `session_kind` when known) onto its already-registered agent, merged + /// over the metadata the engine already holds. + /// + /// `session_id` is the provider's own session id — the Claude Code session + /// UUID or the Codex thread id — which is also the id a recorded session + /// carries in relay history. Publishing it lets a dashboard showing that + /// session find the fleet worker's `@name` and message it. The keys match + /// the ones desktop session agents already publish. + /// + /// Nothing is sent when a later spawn of the same name has claimed the + /// session metadata since `claim` was taken; the check runs under the + /// name's metadata lock, so a superseded publish can never land after the + /// newer one. + /// + /// Best-effort like [`Self::publish_declared_metadata`]: callers log a + /// failure rather than failing the spawn. + pub async fn publish_session_metadata( + &self, + claim: &SessionMetadataClaim, + session_id: &str, + session_kind: Option<&str>, + ) -> std::result::Result { + let metadata = session_metadata_map(session_id, session_kind); + if metadata.is_empty() { + return Ok(SessionMetadataPublish::NothingToPublish); + } + let _guard = claim.fence.lock.lock().await; + if claim.fence.latest_session.load(Ordering::SeqCst) != claim.sequence { + return Ok(SessionMetadataPublish::Superseded); + } + self.patch_agent_metadata(&claim.agent_name, metadata) + .await?; + Ok(SessionMetadataPublish::Published) + } + + fn metadata_fence(&self, name: &str) -> Arc { + let mut fences = self + .metadata_fences + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + Arc::clone(fences.entry(name.to_string()).or_default()) + } + + /// [`Self::patch_agent_metadata`] under the name's metadata lock, so it + /// never overlaps another metadata PATCH this broker sends for the name. + async fn merge_agent_metadata( + &self, + name: &str, + metadata: serde_json::Map, + ) -> std::result::Result<(), RelaycastRegistrationError> { + let fence = self.metadata_fence(name); + let _guard = fence.lock.lock().await; + self.patch_agent_metadata(name, metadata).await + } + + /// `PATCH /v1/agents/:name` with only `metadata`, which the engine merges + /// over the record's existing metadata. Callers hold the name's metadata + /// lock (see [`MetadataPublishFence`]). + async fn patch_agent_metadata( + &self, + name: &str, + metadata: serde_json::Map, + ) -> std::result::Result<(), RelaycastRegistrationError> { let relay = self .relay_client() .ok_or_else(|| RelaycastRegistrationError::Transport { agent_name: name.to_string(), detail: "SDK relay client not initialized".to_string(), })?; - // Send ONLY the declared keys. `PATCH /v1/agents/:name` merges them over + // Send ONLY the given keys. `PATCH /v1/agents/:name` merges them over // the record's existing metadata server-side — verified in the engine at // both the ref fleet-e2e pins (v7.0.0, eb7563ff) and relaycast `main` // (`packages/engine/src/routes/agent.ts`: @@ -839,7 +982,7 @@ impl RelaycastHttpClient { .update_agent( name, UpdateAgentRequest { - metadata: Some(declared_metadata), + metadata: Some(metadata), ..Default::default() }, ) @@ -2567,6 +2710,29 @@ fn declared_metadata_map(declared: &AgentRegistrationMetadata) -> serde_json::Ma metadata } +/// `session_id` (and `session_kind` when known), trimmed, with blanks omitted. +/// +/// A blank session id yields an empty map: there is nothing to link, and an +/// empty value would overwrite a session id the engine already holds. +fn session_metadata_map( + session_id: &str, + session_kind: Option<&str>, +) -> serde_json::Map { + let mut metadata = serde_json::Map::new(); + let session_id = session_id.trim(); + if session_id.is_empty() { + return metadata; + } + metadata.insert( + "session_id".to_string(), + Value::String(session_id.to_string()), + ); + if let Some(kind) = session_kind.map(str::trim).filter(|kind| !kind.is_empty()) { + metadata.insert("session_kind".to_string(), Value::String(kind.to_string())); + } + metadata +} + /// Convert a terminal SDK error into the typed registration error, keeping /// the two facts the broker's retry loops need from the SDK layer: the /// server's `Retry-After` (so the broker paces on the same cadence the SDK @@ -2658,7 +2824,7 @@ mod tests { retry_workspace_busy_reconcile, with_registration_attempts, workspace_busy_reconcile_delay, workspace_busy_retry_allowed, ImpersonationAwareRegistrationError, MessageInjectionMode, RecipientReachability, RegRetryOutcome, RegisterIntent, RelaycastHttpClient, - RelaycastRegistrationError, MAX_AGENT_REGISTRATION_ELAPSED, + RelaycastRegistrationError, SessionMetadataPublish, MAX_AGENT_REGISTRATION_ELAPSED, MAX_AGENT_REGISTRATION_OUTER_TIMEOUT, MAX_AGENT_REGISTRATION_RETRY_DELAY, WORKSPACE_BUSY_ACTION_SAFETY_CAP, WORKSPACE_BUSY_CREATE_ONLY_SAFETY_CAP, WORKSPACE_BUSY_RECONCILE_BUDGET, WORKSPACE_BUSY_RECONCILE_SAFETY_CAP, @@ -3708,6 +3874,236 @@ mod tests { any_write.assert_hits(0); } + /// The provider session id rides a merge-only PATCH carrying exactly + /// `session_id` and `session_kind`, so the engine-owned `fleet` placement + /// record and any declared keys survive untouched. + #[tokio::test] + async fn publish_session_metadata_sends_only_session_keys() { + let server = MockServer::start(); + let read = server.mock(|when, then| { + when.method(GET).path("/v1/agents/worker-a"); + then.status(500); + }); + let update = server.mock(|when, then| { + when.method(PATCH) + .path("/v1/agents/worker-a") + .json_body(json!({ + "metadata": { + "session_id": "0f5c8d3e-1b2a-4c5d-9e8f-7a6b5c4d3e2f", + "session_kind": "claude-terminal" + } + })); + then.status(200).json_body(json!({ + "ok": true, + "data": { + "id": "agent_worker_a", + "name": "worker-a", + "type": "agent", + "status": "online", + "persona": null, + "metadata": {} + } + })); + }); + + let client = seeded_http_client(&server.base_url()); + let claim = client.claim_session_metadata("worker-a").unwrap(); + let outcome = client + .publish_session_metadata( + &claim, + " 0f5c8d3e-1b2a-4c5d-9e8f-7a6b5c4d3e2f ", + Some("claude-terminal"), + ) + .await + .expect("publishing session metadata should succeed"); + + assert_eq!(outcome, SessionMetadataPublish::Published); + read.assert_hits(0); + update.assert_hits(1); + } + + /// Without a known kind only `session_id` is sent; a blank kind must not + /// overwrite one the engine already holds with an empty string. + #[tokio::test] + async fn publish_session_metadata_omits_blank_kind() { + let server = MockServer::start(); + let update = server.mock(|when, then| { + when.method(PATCH) + .path("/v1/agents/worker-a") + .json_body(json!({ "metadata": { "session_id": "thread-123" } })); + then.status(200).json_body(json!({ + "ok": true, + "data": { + "id": "agent_worker_a", + "name": "worker-a", + "type": "agent", + "status": "online", + "persona": null, + "metadata": {} + } + })); + }); + + let client = seeded_http_client(&server.base_url()); + let claim = client.claim_session_metadata("worker-a").unwrap(); + client + .publish_session_metadata(&claim, "thread-123", Some(" ")) + .await + .expect("publishing session metadata should succeed"); + + update.assert_hits(1); + } + + /// Must-not-fire: a blank session id links nothing, so no request is made. + #[tokio::test] + async fn publish_session_metadata_makes_no_request_for_blank_session_id() { + let server = MockServer::start(); + let any_write = server.mock(|when, then| { + when.method(PATCH).path("/v1/agents/worker-a"); + then.status(500); + }); + + let client = seeded_http_client(&server.base_url()); + let claim = client.claim_session_metadata("worker-a").unwrap(); + let outcome = client + .publish_session_metadata(&claim, " ", Some("codex")) + .await + .expect("a blank session id is a no-op, not an error"); + + assert_eq!(outcome, SessionMetadataPublish::NothingToPublish); + any_write.assert_hits(0); + } + + /// Must-not-fire: a released worker's publish that has not run by the time + /// a replacement reuses its name must not label the replacement with the + /// old session; only the replacement's session id is sent. + #[tokio::test] + async fn superseded_session_metadata_publish_sends_nothing() { + let server = MockServer::start(); + let stale = server.mock(|when, then| { + when.method(PATCH) + .path("/v1/agents/worker-a") + .json_body(json!({ "metadata": { "session_id": "session-old" } })); + then.status(500); + }); + let current = server.mock(|when, then| { + when.method(PATCH) + .path("/v1/agents/worker-a") + .json_body(json!({ "metadata": { "session_id": "session-new" } })); + then.status(200).json_body(json!({ + "ok": true, + "data": { + "id": "agent_worker_a", + "name": "worker-a", + "type": "agent", + "status": "online", + "persona": null, + "metadata": {} + } + })); + }); + + let client = seeded_http_client(&server.base_url()); + let old_claim = client.claim_session_metadata("worker-a").unwrap(); + let new_claim = client.claim_session_metadata("worker-a").unwrap(); + + let new_outcome = client + .publish_session_metadata(&new_claim, "session-new", None) + .await + .expect("the latest claim publishes"); + let old_outcome = client + .publish_session_metadata(&old_claim, "session-old", None) + .await + .expect("a superseded publish is a no-op, not an error"); + + assert_eq!(new_outcome, SessionMetadataPublish::Published); + assert_eq!(old_outcome, SessionMetadataPublish::Superseded); + current.assert_hits(1); + stale.assert_hits(0); + } + + /// The engine merges metadata PATCHes with an unlocked read-modify-write, + /// so two of ours in flight for one name could each drop the other's keys. + /// A session publish issued while a declared-metadata PATCH is in flight + /// must wait for it to complete before sending. + #[tokio::test] + async fn metadata_publishes_for_one_name_do_not_overlap() { + let server = MockServer::start(); + let agent_body = json!({ + "ok": true, + "data": { + "id": "agent_worker_a", + "name": "worker-a", + "type": "agent", + "status": "online", + "persona": null, + "metadata": {} + } + }); + let declared_body = agent_body.clone(); + let declared = server.mock(move |when, then| { + when.method(PATCH) + .path("/v1/agents/worker-a") + .json_body(json!({ "metadata": { "role": "reviewer" } })); + then.status(200) + .delay(Duration::from_millis(600)) + .json_body(declared_body); + }); + let session = server.mock(move |when, then| { + when.method(PATCH) + .path("/v1/agents/worker-a") + .json_body(json!({ "metadata": { "session_id": "thread-123" } })); + then.status(200).json_body(agent_body); + }); + + let client = seeded_http_client(&server.base_url()); + let declared_task = { + let client = client.clone(); + tokio::spawn(async move { + let metadata = AgentRegistrationMetadata { + role: Some("reviewer".to_string()), + ..Default::default() + }; + client + .publish_declared_metadata("worker-a", &metadata) + .await + }) + }; + while declared.hits() == 0 { + tokio::time::sleep(Duration::from_millis(5)).await; + } + let claim = client.claim_session_metadata("worker-a").unwrap(); + let session_task = { + let client = client.clone(); + tokio::spawn(async move { + client + .publish_session_metadata(&claim, "thread-123", None) + .await + }) + }; + + tokio::time::sleep(Duration::from_millis(200)).await; + assert_eq!( + session.hits(), + 0, + "session PATCH was sent while the declared PATCH was still in flight" + ); + + declared_task + .await + .unwrap() + .expect("declared publish succeeds"); + assert_eq!( + session_task + .await + .unwrap() + .expect("session publish succeeds"), + SessionMetadataPublish::Published + ); + declared.assert_hits(1); + session.assert_hits(1); + } + /// A presence update used to call POST /v1/agents/release with no reason. /// That invalidated the credential of a participant that could still be /// running, and left an unattributable `release.reason = null` record. The diff --git a/crates/broker/src/runtime/api.rs b/crates/broker/src/runtime/api.rs index e1e82a123..0f043664b 100644 --- a/crates/broker/src/runtime/api.rs +++ b/crates/broker/src/runtime/api.rs @@ -804,6 +804,11 @@ impl BrokerRuntime { name.as_str(), registration_metadata, ); + super::fleet::spawn_session_metadata_publish( + relaycast_http, + name.as_str(), + &effective_spec, + ); } if owns_identity { if let Some(worker) = workers.workers.get(&name) { diff --git a/crates/broker/src/runtime/fleet.rs b/crates/broker/src/runtime/fleet.rs index 388115db7..61d8935b7 100644 --- a/crates/broker/src/runtime/fleet.rs +++ b/crates/broker/src/runtime/fleet.rs @@ -8,6 +8,7 @@ use crate::{ listen_api::{DeliveryRouteError, ListenApiRequest, SetInboundDeliveryModeOk}, node_control::{delivery_ack, handler_unavailable_result, DeliveryDecision, ReceiptAckability}, node_delivery_probe::DeliverDisposition, + relaycast::SessionMetadataPublish, terminal_control::{ request_terminal_reconnect, TerminalControlCommand, TerminalControlEvent, TerminalDeliveryDiagnostics, TerminalFromCloud, TerminalMode, TerminalToCloud, @@ -2294,6 +2295,98 @@ pub(super) fn spawn_declared_metadata_publish( }); } +/// The provider session a spawned worker runs, as published on its Relaycast +/// agent: `(session_id, session_kind)`. +/// +/// `session_id` is the id the worker's CLI records its session under — the +/// Claude Code session UUID or the Codex thread id — resolved by +/// `WorkerRegistry::spawn` before the process starts (and reported as +/// `sessionId` by `GET /api/spawned`). `None` when the spawn has no session id, +/// e.g. a CLI the broker cannot pre-assign one for. +/// +/// `session_kind` names the CLI family, using the same values desktop session +/// agents publish where they overlap: `claude-terminal` for an interactive +/// (PTY) Claude Code session, `codex` for Codex. Other CLIs report their +/// normalized CLI name (headless Claude reports `claude`). +pub(super) fn worker_session_metadata(spec: &AgentSpec) -> Option<(String, Option)> { + let session_id = spec + .session_id + .as_deref() + .map(str::trim) + .filter(|id| !id.is_empty())? + .to_string(); + let cli = spec + .cli + .as_deref() + // Parse the executable the way worker startup does, so a quoted path + // with spaces (`"/opt/AI Tools/codex" --flag`) names `codex`. + .and_then(|cli| crate::cli::command_parse::parse_cli_command(cli).ok()) + .map(|(command, _)| crate::cli::command_parse::normalize_cli_name(&command).to_lowercase()) + .or_else(|| { + spec.provider + .as_ref() + .map(|provider| super::headless::headless_provider_cli_name(provider).to_string()) + }); + let kind = cli.map(|cli| { + let family = cli.split(':').next().unwrap_or(&cli).to_string(); + match family.as_str() { + "claude" if spec.runtime == AgentRuntime::Pty => "claude-terminal".to_string(), + _ => family, + } + }); + Some((session_id, kind)) +} + +/// Publish a freshly spawned (or respawned) worker's provider session id onto +/// its Relaycast agent, on its own task. +/// +/// Called after every successful spawn of a worker that holds a hosted +/// identity, so a respawn that resumes or starts a different session updates +/// the published id. The session claim is taken here, synchronously and even +/// when the spawn has no session id, so any publish still pending from an +/// earlier worker of the same name is superseded and cannot label this one +/// with the old session. Detached and best-effort for the same reasons as +/// [`spawn_declared_metadata_publish`]: a failure is logged, never fatal. +pub(super) fn spawn_session_metadata_publish( + relaycast_http: &RelaycastHttpClient, + name: &str, + spec: &AgentSpec, +) { + let Ok(claim) = relaycast_http.claim_session_metadata(name) else { + return; + }; + let Some((session_id, session_kind)) = worker_session_metadata(spec) else { + return; + }; + let http = relaycast_http.clone(); + tokio::spawn(async move { + let agent = claim.agent_name(); + match http + .publish_session_metadata(&claim, &session_id, session_kind.as_deref()) + .await + { + Ok(SessionMetadataPublish::Published) => tracing::debug!( + worker = %agent, + session_id = %session_id, + "published provider session id for spawned agent" + ), + Ok(SessionMetadataPublish::Superseded) => tracing::debug!( + worker = %agent, + session_id = %session_id, + "skipped provider session id publish; a later spawn of this name owns it" + ), + Ok(SessionMetadataPublish::NothingToPublish) => {} + Err(error) => tracing::error!( + worker = %agent, + session_id = %session_id, + error = %error, + "failed to publish provider session id; the agent is registered and running \ + but its session cannot be linked to it" + ), + } + }); +} + /// Bind an agent to this node by sending node-control `agent.register` and /// awaiting the engine reply with the minted agent token. This is the single /// "register agent via node" step both the `/api/spawn` path and the node @@ -4652,6 +4745,61 @@ mod tests { assert_eq!(delivery_book.active_agent_id("worker-a"), None); } + #[test] + fn worker_session_metadata_reports_provider_session_and_kind() { + let mut spec = test_agent_spec(Some(" thread-123 "), None); + assert_eq!( + worker_session_metadata(&spec), + Some(("thread-123".to_string(), Some("codex".to_string()))) + ); + + spec.cli = Some("/usr/local/bin/claude --model opus".to_string()); + assert_eq!( + worker_session_metadata(&spec), + Some(( + "thread-123".to_string(), + Some("claude-terminal".to_string()) + )) + ); + + // A quoted executable path with spaces names the binary, not the + // first whitespace-separated fragment of its directory. + spec.cli = Some("\"/opt/AI Tools/codex\" --full-auto".to_string()); + assert_eq!( + worker_session_metadata(&spec), + Some(("thread-123".to_string(), Some("codex".to_string()))) + ); + + spec.cli = Some("claude:opus".to_string()); + spec.runtime = AgentRuntime::Headless; + assert_eq!( + worker_session_metadata(&spec), + Some(("thread-123".to_string(), Some("claude".to_string()))) + ); + + spec.cli = None; + spec.provider = Some(ProtocolHeadlessProvider::Opencode); + assert_eq!( + worker_session_metadata(&spec), + Some(("thread-123".to_string(), Some("opencode".to_string()))) + ); + + spec.provider = None; + assert_eq!( + worker_session_metadata(&spec), + Some(("thread-123".to_string(), None)) + ); + } + + #[test] + fn worker_session_metadata_is_none_without_a_session_id() { + assert_eq!(worker_session_metadata(&test_agent_spec(None, None)), None); + assert_eq!( + worker_session_metadata(&test_agent_spec(Some(" "), None)), + None + ); + } + #[tokio::test] async fn refresh_fleet_inventory_session_ref_publishes_immediate_sync() { let (tx, mut rx) = mpsc::channel(4); diff --git a/crates/broker/src/runtime/maintenance.rs b/crates/broker/src/runtime/maintenance.rs index aac7a0c22..7e03fbcd6 100644 --- a/crates/broker/src/runtime/maintenance.rs +++ b/crates/broker/src/runtime/maintenance.rs @@ -657,6 +657,7 @@ impl BrokerRuntime { } } + let has_hosted_identity = worker_relay_key.is_some(); match workers .spawn( rst.payload.spec.clone(), @@ -673,6 +674,15 @@ impl BrokerRuntime { { Ok(effective_spec) => { fleet_load_changed = true; + // A respawn may resume the same session or start a new + // one; either way the published id must follow it. + if has_hosted_identity { + super::fleet::spawn_session_metadata_publish( + relaycast_http, + name.as_str(), + &effective_spec, + ); + } // A supervised tokenless worker owns the identity // across respawns, but each process is a new exact // generation. Refresh custody before any later diff --git a/crates/broker/src/runtime/relaycast_events.rs b/crates/broker/src/runtime/relaycast_events.rs index 64c005097..89ea410c6 100644 --- a/crates/broker/src/runtime/relaycast_events.rs +++ b/crates/broker/src/runtime/relaycast_events.rs @@ -873,6 +873,16 @@ pub(super) async fn spawn_worker_from_request( .await { Ok(effective_spec) => { + if worker_relay_key.is_some() { + // The provider session id is only final once the spawn has + // resolved it, so it is published here rather than alongside + // the declared metadata at registration. + super::fleet::spawn_session_metadata_publish( + workspace_http, + name.as_str(), + &effective_spec, + ); + } if owns_identity { if let Some(worker) = workers.workers.get(&name) { workers diff --git a/crates/broker/src/runtime/tests.rs b/crates/broker/src/runtime/tests.rs index c5d6994bd..69e8ec626 100644 --- a/crates/broker/src/runtime/tests.rs +++ b/crates/broker/src/runtime/tests.rs @@ -7946,6 +7946,18 @@ async fn assert_http_spawn_metadata_publication(supplied_token: bool, valid_cwd: then.status(200) .json_body(json!({"ok":true,"data":identity})); }); + // The provider session id rides its own merge-only PATCH once the spawn + // has resolved it, so the dashboard can link the recorded session to this + // worker without disturbing the declared keys or the `fleet` record. + let session_metadata = server.mock(|when, then| { + when.method(PATCH) + .path("/v1/agents/metadata-worker") + .json_body(json!({"metadata":{ + "session_id":"metadata-session", "session_kind":"cat" + }})); + then.status(200) + .json_body(json!({"ok":true,"data":identity})); + }); let unexpected_cleanup = server.mock(|when, then| { when.method(POST).path("/v1/agents/release"); then.status(500); @@ -8015,20 +8027,22 @@ async fn assert_http_spawn_metadata_publication(supplied_token: bool, valid_cwd: if valid_cwd { assert_eq!(response.unwrap()["success"], true); let published = tokio::time::timeout(Duration::from_secs(2), async { - while metadata.hits() == 0 { + while metadata.hits() == 0 || session_metadata.hits() == 0 { tokio::task::yield_now().await; } }) .await; assert!( published.is_ok(), - "successful spawn did not publish declared metadata" + "successful spawn did not publish declared and session metadata" ); metadata.assert_hits(1); + session_metadata.assert_hits(1); } else { assert!(response.unwrap_err().contains("cwd")); tokio::time::sleep(Duration::from_millis(100)).await; metadata.assert_hits(0); + session_metadata.assert_hits(0); } create.assert_hits(usize::from(!supplied_token)); bind.assert_hits(usize::from(!supplied_token));