diff --git a/crates/broker/src/listen_api.rs b/crates/broker/src/listen_api.rs index 74f863b60..e45e42a99 100644 --- a/crates/broker/src/listen_api.rs +++ b/crates/broker/src/listen_api.rs @@ -818,27 +818,21 @@ async fn listen_api_spawn( let spawn_mode = body .get("spawn_mode") .or_else(|| body.get("spawnMode")) - .and_then(Value::as_str) - .map(|value| value.trim().to_ascii_lowercase()); - let spawn_mode_exit_after_task = match spawn_mode.as_deref() { - None | Some("") | Some("interactive") => false, - Some("task_exit" | "task-exit" | "single_shot" | "single-shot") => true, - Some(other) => { - return ( - axum::http::StatusCode::BAD_REQUEST, - axum::Json(json!({ - "success": false, - "error": format!("unsupported spawnMode '{other}' (expected 'interactive' or 'task_exit')") - })), - ); - } - }; - let exit_after_task = body + .and_then(Value::as_str); + let explicit_exit_after_task = body .get("exit_after_task") .or_else(|| body.get("exitAfterTask")) - .and_then(Value::as_bool) - .unwrap_or(false) - || spawn_mode_exit_after_task; + .and_then(Value::as_bool); + let exit_after_task = + match crate::runtime::resolve_exit_after_task(spawn_mode, explicit_exit_after_task) { + Ok(value) => value, + Err(error) => { + return ( + axum::http::StatusCode::BAD_REQUEST, + axum::Json(json!({ "success": false, "error": error })), + ); + } + }; let skip_relay_prompt = body .get("skip_relay_prompt") .or_else(|| body.get("skipRelayPrompt")) @@ -3729,6 +3723,37 @@ mod auth_tests { spawn_replier.await.expect("spawn replier should complete"); } + #[tokio::test] + async fn spawn_route_rejects_unsupported_spawn_mode() { + let (router, _rx) = test_router(Some("secret")); + let response = router + .oneshot( + Request::builder() + .uri("/api/spawn") + .method("POST") + .header("x-api-key", "secret") + .header("content-type", "application/json") + .body(Body::from( + json!({ + "name": "worker-a", + "cli": "codex", + "spawnMode": "detached" + }) + .to_string(), + )) + .expect("request should build"), + ) + .await + .expect("request should succeed"); + + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + let body = response_json(response).await; + assert!(body["error"] + .as_str() + .expect("error should be a string") + .contains("unsupported spawnMode 'detached'")); + } + #[tokio::test] async fn spawn_route_rejects_harness_id() { let (router, _rx) = test_router(Some("secret")); diff --git a/crates/broker/src/runtime/fleet.rs b/crates/broker/src/runtime/fleet.rs index 73c7939f3..2d6466e3d 100644 --- a/crates/broker/src/runtime/fleet.rs +++ b/crates/broker/src/runtime/fleet.rs @@ -374,6 +374,22 @@ impl BrokerRuntime { let channel = action_invoke_string(&invoke.input, &["channel"]); let model = action_invoke_string(&invoke.input, &["model"]); + // Honor the task-exit lifecycle exactly like the local HTTP spawn API: + // `spawn_mode: task_exit` / `exit_after_task: true` make the agent exit + // once its task is done instead of idling. Reject an unknown spawn_mode + // loudly rather than silently defaulting to interactive. + let spawn_mode = action_invoke_string(&invoke.input, &["spawn_mode", "spawnMode"]); + let explicit_exit_after_task = + action_invoke_bool(&invoke.input, &["exit_after_task", "exitAfterTask"]); + let exit_after_task = + match resolve_exit_after_task(spawn_mode.as_deref(), explicit_exit_after_task) { + Ok(value) => value, + Err(error) => { + self.reply_action_error(&invoke.invocation_id, &error).await; + return; + } + }; + // Reuse the action input as the `ws_value` the spawn fn reads // harnessConfig / supplied tokens from, mirroring the firehose payload // shape (top-level and nested-`agent` lookups both work). @@ -406,6 +422,7 @@ impl BrokerRuntime { task, channel, model, + exit_after_task, &ws_value, &workspace_id, None, @@ -940,6 +957,25 @@ fn action_invoke_string(input: &Value, keys: &[&str]) -> Option { None } +/// Read the first boolean at any of the given top-level keys of an +/// `action.invoke` input object (also checks under a nested `agent` object), +/// mirroring [`action_invoke_string`]'s lookup order for the flattened-vs-nested +/// spawn payload shape. +fn action_invoke_bool(input: &Value, keys: &[&str]) -> Option { + for key in keys { + if let Some(value) = input.get(key).and_then(Value::as_bool) { + return Some(value); + } + } + let agent = input.get("agent")?; + for key in keys { + if let Some(value) = agent.get(key).and_then(Value::as_bool) { + return Some(value); + } + } + None +} + /// Message fields extracted from a node `deliver` payload, ready to build a /// [`RelayDelivery`]. struct FleetDeliveryFields { @@ -1279,6 +1315,55 @@ mod tests { ); } + #[test] + fn action_invoke_bool_reads_top_level_and_nested_agent() { + // Top-level (flattened) and nested-`agent` shapes both resolve, matching + // the fleet TS layer that flattens `{...spawn.agent, task, ...}`. + assert_eq!( + action_invoke_bool(&json!({"exit_after_task": true}), &["exit_after_task"]), + Some(true) + ); + assert_eq!( + action_invoke_bool( + &json!({"agent": {"exitAfterTask": false}}), + &["exit_after_task", "exitAfterTask"] + ), + Some(false) + ); + // Absent on both levels yields None so the caller can default. + assert_eq!( + action_invoke_bool(&json!({"cli": "codex"}), &["exit_after_task"]), + None + ); + } + + #[test] + fn action_invoke_spawn_input_resolves_task_exit_lifecycle() { + // The engine-dispatched spawn reads spawn_mode/exit_after_task from the + // invoke input exactly like the local HTTP spawn, from either the + // flattened top level or the nested `agent` object. + let top_level = json!({"cli": "codex", "spawn_mode": "task_exit"}); + assert!(resolve_exit_after_task( + action_invoke_string(&top_level, &["spawn_mode", "spawnMode"]).as_deref(), + action_invoke_bool(&top_level, &["exit_after_task", "exitAfterTask"]), + ) + .expect("valid spawn_mode")); + + let nested = json!({"agent": {"cli": "codex", "spawnMode": "interactive"}}); + assert!(!resolve_exit_after_task( + action_invoke_string(&nested, &["spawn_mode", "spawnMode"]).as_deref(), + action_invoke_bool(&nested, &["exit_after_task", "exitAfterTask"]), + ) + .expect("valid spawn_mode")); + + let explicit = json!({"cli": "codex", "exit_after_task": true}); + assert!(resolve_exit_after_task( + action_invoke_string(&explicit, &["spawn_mode", "spawnMode"]).as_deref(), + action_invoke_bool(&explicit, &["exit_after_task", "exitAfterTask"]), + ) + .expect("valid explicit flag")); + } + #[test] fn fleet_initial_session_ref_prefers_explicit_spec_session() { let spec = test_agent_spec(Some("session-spec"), Some("session-harness")); diff --git a/crates/broker/src/runtime/relaycast_events.rs b/crates/broker/src/runtime/relaycast_events.rs index 24ff6919e..4a1bbffa3 100644 --- a/crates/broker/src/runtime/relaycast_events.rs +++ b/crates/broker/src/runtime/relaycast_events.rs @@ -265,7 +265,9 @@ pub(super) async fn release_worker_locally( /// directly via `action.invoke`. The spawn fields (`cli`, `task`, `channel`, /// `model`) previously came off the typed event payload and are now passed in; /// `ws_value` is retained for `harnessConfig`/token extraction exactly as -/// before. `control_dedup_key` carries the firehose control dedup key so the +/// before. `exit_after_task` carries the resolved task-exit lifecycle so an +/// engine-dispatched spawn exits after its task identically to a local HTTP +/// spawn. `control_dedup_key` carries the firehose control dedup key so the /// local spawn-echo dedup behaves identically. #[allow(clippy::too_many_arguments)] pub(super) async fn spawn_worker_from_request( @@ -274,6 +276,7 @@ pub(super) async fn spawn_worker_from_request( task: Option, channel: Option, model: Option, + exit_after_task: bool, ws_value: &Value, workspace_id: &WorkspaceId, control_dedup_key: Option<&str>, @@ -383,7 +386,14 @@ pub(super) async fn spawn_worker_from_request( channels: channels.clone(), restart_policy: None, }; - let mut effective_task = normalize_initial_task(task.clone()); + // Mirror the local HTTP spawn path (`runtime/api.rs`): a task-exit spawn + // appends the clean-exit contract to the initial task so the agent exits + // once it is done instead of idling forever. + let mut effective_task = if exit_after_task { + Some(apply_exit_after_task_instruction(task.clone())) + } else { + normalize_initial_task(task.clone()) + }; // Pre-register an agent token for every spawned worker. // The Agent Relay MCP server needs RELAY_AGENT_TOKEN + diff --git a/crates/broker/src/runtime/session.rs b/crates/broker/src/runtime/session.rs index 87b053b4f..366a730d0 100644 --- a/crates/broker/src/runtime/session.rs +++ b/crates/broker/src/runtime/session.rs @@ -106,6 +106,31 @@ pub(crate) fn apply_exit_after_task_instruction(task: Option) -> String } } +/// Resolve a spawn request's effective `exit_after_task` lifecycle flag from its +/// `spawn_mode` selector and any explicit `exit_after_task` boolean. +/// +/// Shared by the local HTTP spawn API and the engine-dispatched node spawn so +/// task-exit semantics are identical on both paths: a `spawn_mode` of +/// `task_exit`/`single_shot` — or an explicit `exit_after_task: true` — makes +/// the agent exit once its task is done; `interactive`/absent keeps it running. +/// An unrecognized `spawn_mode` is rejected with a caller-facing message. +pub(crate) fn resolve_exit_after_task( + spawn_mode: Option<&str>, + exit_after_task: Option, +) -> Result { + let normalized = spawn_mode.map(|value| value.trim().to_ascii_lowercase()); + let spawn_mode_exit_after_task = match normalized.as_deref() { + None | Some("") | Some("interactive") => false, + Some("task_exit" | "task-exit" | "single_shot" | "single-shot") => true, + Some(other) => { + return Err(format!( + "unsupported spawnMode '{other}' (expected 'interactive' or 'task_exit')" + )); + } + }; + Ok(exit_after_task.unwrap_or(false) || spawn_mode_exit_after_task) +} + pub(crate) struct RelaySessionOptions<'a> { pub(crate) paths: &'a RuntimePaths, pub(crate) requested_name: &'a str, diff --git a/crates/broker/src/runtime/tests.rs b/crates/broker/src/runtime/tests.rs index 592567500..5f6d09211 100644 --- a/crates/broker/src/runtime/tests.rs +++ b/crates/broker/src/runtime/tests.rs @@ -43,13 +43,13 @@ use super::{ normalize_channel, normalize_initial_task, normalize_sender, parse_sort_key_from_raw_timestamp, persist_dead_letters_on_shutdown, persist_pending_on_shutdown, queue_inbound_for_delivery_mode, relaycast_spawn_control_dedup_key, relaycast_ws_should_apply_local_spawn_echo_dedup, - relaycast_ws_spawn_token, requeue_dead_letter, resolve_workspace, retry_pending_delivery, - save_dead_letters, seed_supplied_agent_token, send_broker_event, sender_is_dashboard_label, - should_clear_pending_delivery_for_event, synthetic_delivery_read_ack_reason, AgentRuntime, - DeadLetterEntry, DeadLetterStore, DeliveryAttemptOutcome, InboundContext, InboundQueueOutcome, - ObserverTokenMintError, ObserverTokenMintOutcome, PendingDelivery, PendingDeliveryStore, - ProtocolHeadlessProvider, RelayWorkspace, TypedThreadMessage, MAX_DEAD_LETTERS, - MAX_DELIVERY_RETRIES, + relaycast_ws_spawn_token, requeue_dead_letter, resolve_exit_after_task, resolve_workspace, + retry_pending_delivery, save_dead_letters, seed_supplied_agent_token, send_broker_event, + sender_is_dashboard_label, should_clear_pending_delivery_for_event, + synthetic_delivery_read_ack_reason, AgentRuntime, DeadLetterEntry, DeadLetterStore, + DeliveryAttemptOutcome, InboundContext, InboundQueueOutcome, ObserverTokenMintError, + ObserverTokenMintOutcome, PendingDelivery, PendingDeliveryStore, ProtocolHeadlessProvider, + RelayWorkspace, TypedThreadMessage, MAX_DEAD_LETTERS, MAX_DELIVERY_RETRIES, }; use crate::dedup::DedupCache; use crate::relaycast::{ @@ -1164,6 +1164,44 @@ fn exit_after_task_instruction_appends_clean_exit_contract() { assert!(task.contains("output `/exit` on its own line")); } +#[test] +fn resolve_exit_after_task_maps_spawn_mode_and_explicit_flag() { + // Interactive / absent spawn_mode keeps the agent running. + assert!(!resolve_exit_after_task(None, None).expect("absent is valid")); + assert!(!resolve_exit_after_task(Some("interactive"), None).expect("interactive is valid")); + assert!(!resolve_exit_after_task(Some(""), None).expect("blank is valid")); + + // Every accepted task-exit synonym flips the flag on, case/spacing-insensitive. + for mode in [ + "task_exit", + "task-exit", + "single_shot", + "single-shot", + " Task_Exit ", + ] { + assert!( + resolve_exit_after_task(Some(mode), None).expect("task-exit synonym is valid"), + "spawn_mode '{mode}' should resolve to exit_after_task=true" + ); + } + + // An explicit exit_after_task=true wins even without a spawn_mode. + assert!(resolve_exit_after_task(None, Some(true)).expect("explicit flag is valid")); + // and does not override an interactive spawn_mode back off. + assert!(resolve_exit_after_task(Some("interactive"), Some(true)).expect("explicit flag wins")); + assert!(!resolve_exit_after_task(Some("interactive"), Some(false)).expect("both off")); +} + +#[test] +fn resolve_exit_after_task_rejects_unknown_spawn_mode() { + let error = resolve_exit_after_task(Some("detached"), None) + .expect_err("unknown spawn_mode must be rejected"); + assert!( + error.contains("unsupported spawnMode 'detached'"), + "error should name the bad mode; got {error}" + ); +} + #[test] fn relaycast_ws_spawn_token_extracts_agent_token() { let value = json!({ diff --git a/packages/harness-driver/src/protocol.ts b/packages/harness-driver/src/protocol.ts index 7fbf11460..86980c5f6 100644 --- a/packages/harness-driver/src/protocol.ts +++ b/packages/harness-driver/src/protocol.ts @@ -5,6 +5,13 @@ export type HeadlessProvider = 'claude' | 'opencode'; export type InboundDeliveryMode = 'auto_inject' | 'manual_flush'; export type SnapshotFormat = 'plain' | 'ansi'; +/** + * Requested spawn lifecycle. `task_exit`/`single_shot` (and their hyphenated + * spellings) make the agent exit once its task is complete; `interactive` + * keeps it running. + */ +export type SpawnMode = 'interactive' | 'task_exit' | 'task-exit' | 'single_shot' | 'single-shot'; + export interface RestartPolicy { enabled?: boolean; max_restarts?: number; @@ -27,6 +34,10 @@ export interface AgentSpec { shadow_of?: string; shadow_mode?: string; restart_policy?: RestartPolicy; + /** Request task-exit lifecycle; rides the spawn input to the broker. */ + spawn_mode?: SpawnMode; + /** Request task-exit lifecycle; rides the spawn input to the broker. */ + exit_after_task?: boolean; } export type MessageInjectionMode = 'wait' | 'steer'; diff --git a/packages/harness-driver/src/types.ts b/packages/harness-driver/src/types.ts index 70d7b0422..6ce7f9a82 100644 --- a/packages/harness-driver/src/types.ts +++ b/packages/harness-driver/src/types.ts @@ -9,6 +9,7 @@ import type { HeadlessProvider, MessageInjectionMode, RestartPolicy, + SpawnMode, } from './protocol.js'; import type { ResolvedHarnessConfig } from './harness.js'; @@ -21,7 +22,6 @@ export type JsonSchema = Record | boolean; * the spawn request reaches the broker, matching the actions surface. */ export type AgentResultSchema = JsonSchema | ZodLikeSchema | SafeParseSchema; -export type SpawnMode = 'interactive' | 'task_exit' | 'task-exit' | 'single_shot' | 'single-shot'; export interface SpawnPtyInput { name: string;