mlua_swarm/core/engine.rs
1//! `Engine` — the long-running stateful runtime plus the `with_state`
2//! helper (R1-R4 discipline).
3//!
4//! The engine owns the Domain side of the Data / Domain split:
5//! flow control (dispatch / verdict), state (`EngineState`), and the
6//! `submit_output` / `output_tail` surface that feeds it. Data-plane
7//! traffic (Big Response bodies) is delegated to the `output_store` module
8//! plus its paired `SpawnerLayer`s and passes through here without the
9//! engine core needing to grow.
10
11use crate::core::agent_context::{RUN_ID_KEY, STEP_CTX_KEY};
12use crate::core::config::EngineCfg;
13use crate::core::ctx::{Ctx, OperatorInfo, OperatorKind, SeniorBridge, SpawnHook};
14use crate::core::errors::EngineError;
15use crate::core::state::{
16 CapTokenRecord, DispatchOutcome, EngineState, Event, EventStream, OperatorSession, ResumeKey,
17 ResumePending, TaskSpec, TaskState, TaskStatus,
18};
19use crate::types::{
20 default_role_verb_table, now_unix, CapToken, Role, RoleVerbGate, RunId, SessionId, StepId,
21 TokenSigner, Verb,
22};
23use crate::worker::adapter::SpawnerAdapter;
24use serde_json::Value;
25use std::collections::HashMap;
26use std::sync::Arc;
27use std::time::{Duration, Instant};
28use tokio::sync::{broadcast, Mutex};
29
30/// Process-wide long-running runtime. Cheap to `clone()` — an `Arc`
31/// lives inside.
32#[derive(Clone)]
33pub struct Engine {
34 inner: Arc<EngineInner>,
35}
36
37struct EngineInner {
38 state: Mutex<EngineState>,
39 cfg: EngineCfg,
40 signer: TokenSigner,
41 gate: RoleVerbGate,
42 event_tx: broadcast::Sender<Event>,
43 /// ID-keyed bridge registry (register-by-ID design). `SeniorBridge`
44 /// and `SpawnHook` are registered by ID; sessions bind to those IDs
45 /// only. Persistence stores just the ID, and on reattach the caller
46 /// re-registers under the same ID to restore presence.
47 senior_bridges: tokio::sync::RwLock<HashMap<String, Arc<dyn SeniorBridge>>>,
48 spawn_hooks: tokio::sync::RwLock<HashMap<String, Arc<dyn SpawnHook>>>,
49 /// ID registry for full-spawn Operator backends (backends that take the
50 /// entire spawn via `execute`). Sibling to `senior_bridges` /
51 /// `spawn_hooks`. `OperatorDelegateMiddleware` looks these up via
52 /// `ctx` and, when `kind = MainAi` / `Composite`, bypasses
53 /// `inner.spawn` and calls `operator.execute` instead.
54 operators: tokio::sync::RwLock<HashMap<String, Arc<dyn crate::operator::Operator>>>,
55 /// Base and hint layer factories for the `SpawnerStack`. At
56 /// `service::linker::link` time, `compiled.router` is wrapped with
57 /// the base factories plus the hint factories resolved from
58 /// `blueprint.spawner_hints.layers`. This is the engine-side
59 /// counterpart to the discipline "Flow / Blueprint doesn't spell out
60 /// middleware implementations — it declares the capabilities it needs
61 /// as hint keys".
62 layer_registry: crate::middleware::LayerRegistry,
63}
64
65/// Renders a `TaskSpec.initial_directive` / `EngineState.prompts`
66/// `Value` down to the `String` shape that string-consuming boundaries
67/// require (issue #18). Strings pass through verbatim; anything else
68/// (Object / Array / Number / Bool / Null) is serde-stringified. This
69/// is the single canonical rendering — the coercion that used to sit
70/// inside `EngineDispatcher::dispatch` moved here and is invoked only
71/// at consumer boundaries: `WorkerPayload.prompt` (HTTP
72/// `/v1/worker/prompt`), `WorkerInvocation.prompt` (in-process
73/// spawners), the subprocess spawner's directive arg/stdin, and the
74/// WS Spawn frame text render (`operator_ws::session`). Everything
75/// upstream (Blueprint dispatch → engine state → `fetch_prompt` →
76/// `Operator::execute`) keeps the `Value` end-to-end.
77pub(crate) fn render_directive_to_string(v: &Value) -> String {
78 match v {
79 Value::String(s) => s.clone(),
80 other => other.to_string(),
81 }
82}
83
84impl Engine {
85 /// Backwards-compatible constructor that starts the engine without a
86 /// layer registry, preserving the signature already used by ~88
87 /// existing call sites. Use this when automatic middleware wrapping
88 /// at bind time is not needed. Callers such as `mlua-swarm-server` go through
89 /// `new_with_layers(cfg, registry)` to enable the hint-resolution path.
90 pub fn new(cfg: EngineCfg) -> Self {
91 Self::new_with_layers(cfg, crate::middleware::LayerRegistry::new())
92 }
93
94 /// Construct an `Engine` with an explicit `LayerRegistry`, enabling
95 /// hint-resolution: `spawner_hints.layers` declared on a `Blueprint`
96 /// are resolved against this registry when the spawner stack is bound
97 /// at `service::linker::link` time.
98 pub fn new_with_layers(
99 cfg: EngineCfg,
100 layer_registry: crate::middleware::LayerRegistry,
101 ) -> Self {
102 let (event_tx, _) = broadcast::channel(256);
103 let signer = TokenSigner::new(&cfg.token_secret);
104 Self {
105 inner: Arc::new(EngineInner {
106 state: Mutex::new(EngineState::new()),
107 cfg,
108 signer,
109 gate: default_role_verb_table(),
110 event_tx,
111 senior_bridges: tokio::sync::RwLock::new(HashMap::new()),
112 spawn_hooks: tokio::sync::RwLock::new(HashMap::new()),
113 operators: tokio::sync::RwLock::new(HashMap::new()),
114 layer_registry,
115 }),
116 }
117 }
118
119 /// Rebuild this `Engine` with a different `RoleVerbGate`. The gate is
120 /// treated as fixed-at-build-time, so this constructs a fresh
121 /// `EngineInner` (fresh empty `EngineState`) rather than mutating in
122 /// place — mainly a testing convenience for swapping gate rules.
123 pub fn with_gate(self, gate: RoleVerbGate) -> Self {
124 // The gate is fixed at build time — the intent is to build a fresh
125 // instance rather than mutating in place. As a testing convenience we
126 // do allow swapping the inner Arc. Simpler form: just rebuild
127 // Arc<EngineInner>.
128 let inner = Arc::new(EngineInner {
129 state: Mutex::new(EngineState::new()),
130 cfg: self.inner.cfg.clone(),
131 signer: self.inner.signer.clone(),
132 gate,
133 event_tx: self.inner.event_tx.clone(),
134 senior_bridges: tokio::sync::RwLock::new(HashMap::new()),
135 spawn_hooks: tokio::sync::RwLock::new(HashMap::new()),
136 operators: tokio::sync::RwLock::new(HashMap::new()),
137 layer_registry: self.inner.layer_registry.clone(),
138 });
139 Self { inner }
140 }
141
142 // ═══════════════════════════════════════════════════════════════════════
143 // Accessors. Production code drives execution through compile +
144 // `service::linker::link` + `dispatch_attempt_with(spawner)` inside
145 // `TaskLaunchService`; `Engine` itself is a pure execution surface — it
146 // does not own a BlueprintStore / EnhanceAdapter / Compiler, nor a
147 // global spawner (the spawner is carried per-request, never stashed on
148 // the engine).
149 // ═══════════════════════════════════════════════════════════════════════
150
151 /// Access the `EngineCfg` this engine was built with.
152 pub fn cfg(&self) -> &EngineCfg {
153 &self.inner.cfg
154 }
155
156 /// Expose the internal `LayerRegistry` — used when deriving a
157 /// sub-engine that needs the same registry re-injected. The
158 /// per-request sub-engine in `mlua-swarm-server` reads the parent engine's
159 /// registry through this accessor and passes it to
160 /// `Engine::new_with_layers(cfg, parent.layer_registry().clone())`.
161 pub fn layer_registry(&self) -> &crate::middleware::LayerRegistry {
162 &self.inner.layer_registry
163 }
164
165 /// Access the `TokenSigner` used to mint/verify `CapToken`s.
166 pub fn signer(&self) -> &TokenSigner {
167 &self.inner.signer
168 }
169
170 /// Clone a handle to the process-wide `Event` broadcast sender. Prefer
171 /// `subscribe` for a ready-to-use receiver.
172 pub fn event_tx(&self) -> broadcast::Sender<Event> {
173 self.inner.event_tx.clone()
174 }
175
176 /// Subscribe to the engine's `Event` broadcast stream.
177 pub fn subscribe(&self) -> EventStream {
178 self.inner.event_tx.subscribe()
179 }
180
181 // ═══════════════════════════════════════════════════════════════════════
182 // §7 with_state — single Mutex + R1-R4 (try_lock + bounded retry + max-hold panic)
183 // ═══════════════════════════════════════════════════════════════════════
184
185 /// The closure is a **sync** `FnOnce` — you cannot pass an async
186 /// closure, which enforces R3 at the type level. Exceeding `max_hold`
187 /// panics so that R4 violations surface immediately.
188 pub async fn with_state<F, R>(&self, op: &'static str, f: F) -> Result<R, EngineError>
189 where
190 F: FnOnce(&mut EngineState) -> R,
191 {
192 let cfg = &self.inner.cfg;
193
194 // R2: try_lock + bounded retry
195 let mut guard_opt = None;
196 for attempt in 0..=cfg.max_retry {
197 match self.inner.state.try_lock() {
198 Ok(g) => {
199 guard_opt = Some(g);
200 break;
201 }
202 Err(_) if cfg.try_only => return Err(EngineError::LockBusy(op)),
203 Err(_) => {
204 let backoff = cfg.backoff_ms_step * (attempt as u64 + 1);
205 tokio::time::sleep(Duration::from_millis(backoff)).await;
206 }
207 }
208 }
209 let mut guard = guard_opt.ok_or(EngineError::LockBusyAfterRetry(op))?;
210
211 // R4: max_hold guard
212 let start = Instant::now();
213 let result = f(&mut guard);
214 let elapsed_ms = start.elapsed().as_millis();
215 drop(guard);
216
217 if elapsed_ms > cfg.max_hold_ms {
218 panic!(
219 "Engine.with_state('{op}') held {elapsed_ms}ms > max {}ms — suspected R3 violation (long op inside lock)",
220 cfg.max_hold_ms
221 );
222 }
223 Ok(result)
224 }
225
226 // ═══════════════════════════════════════════════════════════════════════
227 // Token verify (= sig + expire + gate + uses_left)
228 // ═══════════════════════════════════════════════════════════════════════
229
230 /// Four steps: (1) signature verify, (2) expiry check, (3) role × verb
231 /// gate, (4) `uses_left` consume.
232 pub async fn verify_token(&self, token: &CapToken, verb: Verb) -> Result<(), EngineError> {
233 // (1) sig
234 if !self.inner.signer.verify_sig(token) {
235 return Err(EngineError::BadSignature);
236 }
237 // (2) expire
238 if token.is_expired(now_unix()) {
239 return Err(EngineError::TokenExpired);
240 }
241 // (3) role × verb gate
242 if !self.inner.gate.is_allowed(token.role, verb) {
243 return Err(EngineError::RoleViolation {
244 role: token.role,
245 verb,
246 });
247 }
248 // (4) server-side uses_left consume
249 let fp = token.fingerprint();
250 self.with_state("token.consume", move |s| {
251 let rec = s
252 .tokens
253 .get_mut(&fp)
254 .ok_or_else(|| EngineError::TokenNotFound(fp.clone()))?;
255 rec.consume()
256 .map_err(|_: crate::core::state::CapTokenConsumeError| {
257 EngineError::TokenUsesExhausted
258 })?;
259 Ok::<(), EngineError>(())
260 })
261 .await??;
262 Ok(())
263 }
264
265 /// `verify_token` plus the **task-ownership gate**.
266 ///
267 /// When a Worker-role token calls a state-touch verb (`fetch_prompt` /
268 /// `post_result` / `read_task_state` / `cancel_task` / `poll_task`),
269 /// the gate checks that `CapTokenRecord.task_id` matches the argument
270 /// `task_id`; a mismatch returns `EngineError::TokenTaskMismatch`.
271 /// Operator / Senior / Observer tokens are outside the ownership gate
272 /// and may touch any task.
273 ///
274 /// **Verbs exempt from the gate.** `start_task` and `dispatch_attempt`
275 /// stay outside so recursive swarming keeps working; depth is capped
276 /// by `max_spawn_depth`.
277 pub async fn verify_token_for_task(
278 &self,
279 token: &CapToken,
280 verb: Verb,
281 task_id: &StepId,
282 ) -> Result<(), EngineError> {
283 self.verify_token(token, verb).await?;
284 if token.role != Role::Worker {
285 return Ok(());
286 }
287 let fp = token.fingerprint();
288 let arg_tid = task_id.clone();
289 self.with_state("token.ownership_gate", move |s| {
290 let bound = s.tokens.get(&fp).and_then(|r| r.task_id.as_ref()).cloned();
291 match bound {
292 Some(t) if t == arg_tid => Ok(()),
293 Some(t) => Err(EngineError::TokenTaskMismatch {
294 bound: t.into_string(),
295 arg: arg_tid.into_string(),
296 }),
297 None => Err(EngineError::TokenNotFound(fp.clone())),
298 }
299 })
300 .await??;
301 Ok(())
302 }
303
304 /// Resolve the bound `task_id` from a Worker-role token. Used on the
305 /// simple `/v1/worker/submit` endpoint, where the worker POSTs with a
306 /// token but no `task_id`. Returns `Err` if the token role is not
307 /// Worker, or if no bound task is set.
308 pub async fn task_id_from_token(&self, token: &CapToken) -> Result<StepId, EngineError> {
309 if token.role != Role::Worker {
310 return Err(EngineError::RoleViolation {
311 role: token.role,
312 verb: Verb::PostResult,
313 });
314 }
315 let fp = token.fingerprint();
316 self.with_state("task_id_from_token", move |s| {
317 s.tokens
318 .get(&fp)
319 .and_then(|r| r.task_id.as_ref())
320 .cloned()
321 .ok_or_else(|| EngineError::TokenNotFound(fp.clone()))
322 })
323 .await?
324 }
325
326 /// Resolve a short worker handle (`wh-XXXXXXXX`) to the bound
327 /// `task_id`. Used on `/v1/worker/submit` when the Bearer is a short
328 /// handle string rather than a full `CapToken` JSON. A missing entry
329 /// returns `TokenNotFound`, i.e. "the handle is not in the store".
330 pub async fn task_id_from_handle(&self, handle: &str) -> Result<StepId, EngineError> {
331 let h = handle.to_string();
332 self.with_state("task_id_from_handle", move |s| {
333 let fp = s
334 .worker_handles
335 .get(&h)
336 .cloned()
337 .ok_or_else(|| EngineError::TokenNotFound(format!("handle={h}")))?;
338 s.tokens
339 .get(&fp)
340 .and_then(|r| r.task_id.as_ref())
341 .cloned()
342 .ok_or_else(|| EngineError::TokenNotFound(format!("fp={fp}")))
343 })
344 .await?
345 }
346
347 /// Submit a worker result via a short handle. Skips token verification
348 /// and updates `output_tail` `Final` + `task.last_result` directly in
349 /// a thin path. The caller is expected to have already resolved
350 /// `task_id` via `task_id_from_handle` — the handle's presence in
351 /// `worker_handles` means it was minted server-side and is therefore
352 /// trusted.
353 pub async fn submit_worker_result_trusted(
354 &self,
355 task_id: &StepId,
356 attempt: u32,
357 value: Value,
358 ok: bool,
359 ) -> Result<(), EngineError> {
360 let task_id_for_apply = task_id.clone();
361 let value_for_event = value.clone();
362 self.with_state("submit_worker_result_trusted.output", move |s| {
363 let ev = crate::worker::output::OutputEvent::Final {
364 content: crate::worker::output::ContentRef::Inline {
365 value: value_for_event,
366 },
367 ok,
368 };
369 s.output_store
370 .entry((task_id_for_apply.clone(), attempt))
371 .or_default()
372 .push(ev.clone());
373 s.push_event(crate::core::state::Event::WorkerOutput {
374 task_id: task_id_for_apply,
375 attempt,
376 event: ev,
377 });
378 })
379 .await?;
380 let task_id_for_result = task_id.clone();
381 let value_for_result = value.clone();
382 self.with_state("submit_worker_result_trusted.last_result", move |s| {
383 if let Some(t) = s.tasks.get_mut(&task_id_for_result) {
384 t.last_result = Some(value_for_result);
385 t.updated_at = now_unix();
386 }
387 })
388 .await?;
389 Ok(())
390 }
391
392 /// Mint a short handle and register it in the `worker_handles` map.
393 /// Called immediately after the worker-token mint inside
394 /// `dispatch_attempt_with`, and issues a handle bound to the same
395 /// token fingerprint. Format is `wh-<8 hex chars>` (11 chars total),
396 /// designed to remove the base64 copy-paste failure mode.
397 async fn mint_worker_handle(&self, worker_fp: String) -> Result<String, EngineError> {
398 // The handle is a sole bearer secret on the `/v1/worker/submit`
399 // short-handle path (`submit_worker_result_trusted` skips token
400 // verification), so it must be unguessable — OS RNG, not the
401 // predictable uid counter. 8 hex chars (~4B entropy) keeps the
402 // documented `wh-<8 hex>` wire shape; collision between live
403 // handles is negligible at in-process handle counts.
404 let short = crate::types::secure_hex(4);
405 let handle = format!("wh-{short}");
406 let h = handle.clone();
407 self.with_state("mint_worker_handle", move |s| {
408 s.worker_handles.insert(h, worker_fp);
409 })
410 .await?;
411 Ok(handle)
412 }
413
414 // ═══════════════════════════════════════════════════════════════════════
415 // Session API
416 // ═══════════════════════════════════════════════════════════════════════
417
418 /// Attach a new session with default `OperatorInfo` (`Automate`, no
419 /// bridges/hooks). Shorthand for `attach_with(.., OperatorInfo::default())`.
420 pub async fn attach(
421 &self,
422 operator_id: impl Into<String>,
423 role: Role,
424 ttl: Duration,
425 ) -> Result<CapToken, EngineError> {
426 self.attach_with(
427 operator_id,
428 role,
429 ttl,
430 crate::core::ctx::OperatorInfo::default(),
431 )
432 .await
433 }
434
435 // ═══════════════════════════════════════════════════════════════════════
436 // BridgeRegistry API.
437 // ═══════════════════════════════════════════════════════════════════════
438
439 /// Register a `SeniorBridge` under a name. An existing entry with the
440 /// same name is overwritten. On the persisted-session reattach path,
441 /// the caller re-registers under the same ID beforehand and the
442 /// bridge becomes effective again.
443 pub async fn register_senior_bridge(
444 &self,
445 id: impl Into<String>,
446 bridge: Arc<dyn SeniorBridge>,
447 ) {
448 self.inner
449 .senior_bridges
450 .write()
451 .await
452 .insert(id.into(), bridge);
453 }
454
455 /// Register a `SpawnHook` under a name. An existing entry with the
456 /// same name is overwritten.
457 pub async fn register_spawn_hook(&self, id: impl Into<String>, hook: Arc<dyn SpawnHook>) {
458 self.inner.spawn_hooks.write().await.insert(id.into(), hook);
459 }
460
461 /// Register an `Operator` (a spawn-body backend) under a name. An
462 /// existing entry with the same name is overwritten.
463 /// `OperatorDelegateMiddleware` looks this up via `ctx` and, when
464 /// `kind = MainAi` / `Composite`, bypasses `inner.spawn` and calls
465 /// `operator.execute` instead.
466 pub async fn register_operator(
467 &self,
468 id: impl Into<String>,
469 operator: Arc<dyn crate::operator::Operator>,
470 ) {
471 self.inner
472 .operators
473 .write()
474 .await
475 .insert(id.into(), operator);
476 }
477
478 /// Unregister a `SeniorBridge` by name (e.g. on WebSocket disconnect
479 /// or explicit teardown). A missing ID is a no-op.
480 pub async fn unregister_senior_bridge(&self, id: &str) {
481 self.inner.senior_bridges.write().await.remove(id);
482 }
483
484 /// Unregister a `SpawnHook` by name. A missing ID is a no-op.
485 pub async fn unregister_spawn_hook(&self, id: &str) {
486 self.inner.spawn_hooks.write().await.remove(id);
487 }
488
489 /// Unregister an `Operator` backend by name. A missing ID is a no-op.
490 pub async fn unregister_operator(&self, id: &str) {
491 self.inner.operators.write().await.remove(id);
492 }
493
494 /// Snapshot the list of registered `SpawnHook` IDs (for test
495 /// observation and debugging).
496 pub async fn list_spawn_hook_ids(&self) -> Vec<String> {
497 self.inner
498 .spawn_hooks
499 .read()
500 .await
501 .keys()
502 .cloned()
503 .collect()
504 }
505
506 /// Snapshot the list of registered `SeniorBridge` IDs.
507 pub async fn list_senior_bridge_ids(&self) -> Vec<String> {
508 self.inner
509 .senior_bridges
510 .read()
511 .await
512 .keys()
513 .cloned()
514 .collect()
515 }
516
517 /// Snapshot the list of registered `Operator` IDs.
518 pub async fn list_operator_ids(&self) -> Vec<String> {
519 self.inner.operators.read().await.keys().cloned().collect()
520 }
521
522 /// Attach specifying IDs directly. The caller is expected to have
523 /// pre-registered them via `register_senior_bridge` /
524 /// `register_spawn_hook` / `register_operator`. This is the canonical
525 /// path when persistence is in play.
526 ///
527 /// `kind` is the "Runtime Global" tier of the `OperatorKind` cascade
528 /// (stored verbatim on `OperatorSession.operator_kind`): `Some(_)` is
529 /// an explicit request (including `Some(OperatorKind::Automate)`) that
530 /// outranks the BP-level tiers; `None` leaves it unspecified so the
531 /// BP-level tiers / final default decide. See
532 /// `crate::core::ctx::collapse_operator_kind`.
533 #[allow(clippy::too_many_arguments)]
534 pub async fn attach_with_ids(
535 &self,
536 operator_id: impl Into<String>,
537 role: Role,
538 ttl: Duration,
539 kind: Option<OperatorKind>,
540 bridge_id: Option<String>,
541 hook_id: Option<String>,
542 operator_backend_id: Option<String>,
543 operator_kind_overrides: HashMap<String, OperatorKind>,
544 bp_agent_kinds: HashMap<String, OperatorKind>,
545 bp_global_kind: Option<OperatorKind>,
546 ) -> Result<CapToken, EngineError> {
547 let operator_id = operator_id.into();
548 let token = self
549 .inner
550 .signer
551 .session(operator_id.clone(), role, vec!["*".into()], ttl);
552 let session_id = SessionId::new();
553 let fp = token.fingerprint();
554 let now = now_unix();
555 let token_for_store = token.clone();
556
557 self.with_state("attach_with_ids", |s| {
558 s.tokens
559 .insert(fp.clone(), CapTokenRecord::from_token(token_for_store));
560 s.sessions.insert(
561 session_id.clone(),
562 OperatorSession {
563 id: session_id.clone(),
564 operator_id: operator_id.clone(),
565 role,
566 attached_at: now,
567 last_seen: now,
568 attached: true,
569 owned_task_ids: Vec::new(),
570 token_fp: fp.clone(),
571 operator_kind: kind,
572 runtime_agent_kinds: operator_kind_overrides,
573 bp_agent_kinds,
574 bp_global_kind,
575 bridge_id,
576 hook_id,
577 operator_backend_id,
578 },
579 );
580 s.push_event(Event::SessionAttached {
581 session_id: session_id.clone(),
582 role,
583 });
584 })
585 .await?;
586
587 let _ = self
588 .inner
589 .event_tx
590 .send(Event::SessionAttached { session_id, role });
591 Ok(token)
592 }
593
594 /// Build an `OperatorInfo` by looking up the session's registered IDs
595 /// on the `BridgeRegistry`, plus resolving the 4-tier `OperatorKind`
596 /// cascade for `agent_name` via `crate::core::ctx::collapse_operator_kind`.
597 /// Used when `dispatch_attempt` injects `Ctx`. An unresolved ID
598 /// (nothing registered) is silently `None` — the bridge / hook simply
599 /// does not fire and the default behaviour applies.
600 async fn resolve_operator_info(
601 &self,
602 session: &OperatorSession,
603 agent_name: &str,
604 ) -> OperatorInfo {
605 let senior_bridge = if let Some(id) = &session.bridge_id {
606 self.inner.senior_bridges.read().await.get(id).cloned()
607 } else {
608 None
609 };
610 let spawn_hook = if let Some(id) = &session.hook_id {
611 self.inner.spawn_hooks.read().await.get(id).cloned()
612 } else {
613 None
614 };
615 let operator = if let Some(id) = &session.operator_backend_id {
616 self.inner.operators.read().await.get(id).cloned()
617 } else {
618 None
619 };
620 let runtime_agent = session.runtime_agent_kinds.get(agent_name).copied();
621 // "Runtime Global" tier: `Some(_)` is always an explicit request
622 // (see the field doc on `OperatorSession.operator_kind`).
623 let runtime_global = session.operator_kind;
624 let bp_agent = session.bp_agent_kinds.get(agent_name).copied();
625 let bp_global = session.bp_global_kind;
626 let kind = crate::core::ctx::collapse_operator_kind(
627 runtime_agent,
628 runtime_global,
629 bp_agent,
630 bp_global,
631 );
632 OperatorInfo {
633 kind,
634 id: session.operator_id.clone(),
635 senior_bridge,
636 spawn_hook,
637 operator,
638 }
639 }
640
641 /// Convenience attach that takes an `OperatorInfo` (three
642 /// `Arc<dyn ...>` fields plus `kind`) **inline**.
643 ///
644 /// # Pipeline
645 ///
646 /// Each `Arc<dyn ...>` is auto-registered on the engine's registry
647 /// under a synthetic ID (`br-<hex>` / `hk-<hex>` / `ob-<hex>`), and
648 /// the session stores that synthetic ID. Subsequent `dispatch_attempt`
649 /// calls rebuild the `Arc`s from those IDs via
650 /// `resolve_operator_info`, and the three middlewares fire as usual.
651 ///
652 /// # ⚠ Non-persisted sessions only
653 ///
654 /// Because this API takes inline `Arc`s, the reattach path after
655 /// session persistence cannot rebuild them — the synthetic IDs are
656 /// not present in a freshly started process's registry. If you need
657 /// persistence, use [`Self::attach_with_ids`] with `register_*` calls
658 /// beforehand to go through **named IDs** instead.
659 ///
660 /// Handy for tests and short-lived in-process sessions. Production
661 /// WebSocket callbacks and the like should prefer `attach_with_ids`
662 /// as the canonical path.
663 pub async fn attach_with(
664 &self,
665 operator_id: impl Into<String>,
666 role: Role,
667 ttl: Duration,
668 operator_info: crate::core::ctx::OperatorInfo,
669 ) -> Result<CapToken, EngineError> {
670 let operator_id = operator_id.into();
671 // The caller always hands in a fully-formed `OperatorInfo`
672 // (including its `kind`), so it is stored as an explicit "Runtime
673 // Global" tier request (`Some(kind)`) — this path never persists
674 // BP-level tiers (both stay empty below), so `Some(kind)` resolves
675 // to the same `kind` at dispatch either way; see
676 // `OperatorSession.operator_kind` doc.
677 let kind = operator_info.kind;
678 // BridgeRegistry auto-register: when the caller hands in an
679 // `Arc<dyn>` directly, register it under a synthesised ID (the inline
680 // path aware of persistence). Callers who want to pre-register with a
681 // named ID should use `register_senior_bridge` / `register_spawn_hook`
682 // + `attach_with_ids`.
683 let bridge_id = if let Some(bridge) = operator_info.senior_bridge.clone() {
684 let id = format!("br-{}", crate::types::uid_hex(8));
685 self.inner
686 .senior_bridges
687 .write()
688 .await
689 .insert(id.clone(), bridge);
690 Some(id)
691 } else {
692 None
693 };
694 let hook_id = if let Some(hook) = operator_info.spawn_hook.clone() {
695 let id = format!("hk-{}", crate::types::uid_hex(8));
696 self.inner
697 .spawn_hooks
698 .write()
699 .await
700 .insert(id.clone(), hook);
701 Some(id)
702 } else {
703 None
704 };
705 let operator_backend_id = if let Some(operator) = operator_info.operator.clone() {
706 // `ob-` = operator-backend registry id. Renamed from `op-` in the
707 // issue #11 prefix reconciliation: `op-` used to collide with the
708 // WS operator sid shape (now unified into `S-<hex>` anyway), and a
709 // shared prefix across two unrelated registries made log filtering
710 // by prefix silently ambiguous.
711 let id = format!("ob-{}", crate::types::uid_hex(8));
712 self.inner
713 .operators
714 .write()
715 .await
716 .insert(id.clone(), operator);
717 Some(id)
718 } else {
719 None
720 };
721
722 let token = self
723 .inner
724 .signer
725 .session(operator_id.clone(), role, vec!["*".into()], ttl);
726 let session_id = SessionId::new();
727 let fp = token.fingerprint();
728 let now = now_unix();
729 let token_for_store = token.clone();
730
731 self.with_state("attach_with", |s| {
732 s.tokens
733 .insert(fp.clone(), CapTokenRecord::from_token(token_for_store));
734 s.sessions.insert(
735 session_id.clone(),
736 OperatorSession {
737 id: session_id.clone(),
738 operator_id,
739 role,
740 attached_at: now,
741 last_seen: now,
742 attached: true,
743 owned_task_ids: Vec::new(),
744 token_fp: fp.clone(),
745 operator_kind: Some(kind),
746 runtime_agent_kinds: HashMap::new(),
747 bp_agent_kinds: HashMap::new(),
748 bp_global_kind: None,
749 bridge_id,
750 hook_id,
751 operator_backend_id,
752 },
753 );
754 s.push_event(Event::SessionAttached {
755 session_id: session_id.clone(),
756 role,
757 });
758 })
759 .await?;
760
761 let _ = self
762 .inner
763 .event_tx
764 .send(Event::SessionAttached { session_id, role });
765 Ok(token)
766 }
767
768 /// Mark the session bound to `token` as detached (`attached = false`).
769 /// Tasks are left in place — a later `attach`/`attach_with_ids` call
770 /// carrying the same registered bridge/hook IDs can pick them back up.
771 pub async fn detach(&self, token: &CapToken) -> Result<(), EngineError> {
772 self.verify_token(token, Verb::DetachSession).await?;
773 let fp = token.fingerprint();
774 self.with_state("detach", move |s| {
775 let sid = s
776 .sessions
777 .iter()
778 .find(|(_, sess)| sess.token_fp == fp)
779 .map(|(id, _)| id.clone());
780 if let Some(sid) = sid {
781 if let Some(sess) = s.sessions.get_mut(&sid) {
782 sess.attached = false;
783 }
784 s.push_event(Event::SessionDetached {
785 session_id: sid.clone(),
786 });
787 let _ = sid;
788 }
789 })
790 .await?;
791 Ok(())
792 }
793
794 /// Refresh the session's `last_seen` timestamp and mark it `attached`.
795 /// Called periodically by an attached client to avoid being flipped to
796 /// detached by `start_detach_loop`.
797 pub async fn heartbeat(&self, token: &CapToken) -> Result<(), EngineError> {
798 self.verify_token(token, Verb::Heartbeat).await?;
799 let now = now_unix();
800 let fp = token.fingerprint();
801 self.with_state("heartbeat", move |s| {
802 if let Some(sess) = s.sessions.values_mut().find(|sess| sess.token_fp == fp) {
803 sess.last_seen = now;
804 sess.attached = true;
805 }
806 })
807 .await?;
808 Ok(())
809 }
810
811 // ═══════════════════════════════════════════════════════════════════════
812 // Task lifecycle
813 // ═══════════════════════════════════════════════════════════════════════
814
815 /// Create a new `TaskState` from `spec` and register its initial
816 /// prompt. When the calling token is a Worker (i.e. this is a
817 /// recursive spawn), the new task inherits `parent.spawn_depth + 1`
818 /// and is rejected with `SpawnDepthExceeded` once `max_spawn_depth` is
819 /// hit; an Operator-issued call starts at depth 0.
820 pub async fn start_task(
821 &self,
822 token: &CapToken,
823 spec: TaskSpec,
824 ) -> Result<StepId, EngineError> {
825 self.verify_token(token, Verb::StartTask).await?;
826 let task_id = StepId::new();
827 let initial_directive = spec.initial_directive.clone();
828 let task_id_clone = task_id.clone();
829 let fp = token.fingerprint();
830 let max_depth = self.inner.cfg.max_spawn_depth;
831 self.with_state("start_task", move |s| {
832 // Recursive swarm depth gate (recursion guard):
833 // Worker tokens carry CapTokenRecord.parent_task_id. Give the
834 // child parent's spawn_depth + 1; if it exceeds `max`, raise an
835 // error. Operator tokens (parent_task_id=None) start at depth 0.
836 let parent_depth_opt = s
837 .tokens
838 .get(&fp)
839 .and_then(|rec| rec.task_id.as_ref())
840 .and_then(|tid| s.tasks.get(tid))
841 .map(|t| t.spawn_depth);
842 let depth = match parent_depth_opt {
843 Some(d) => {
844 if d + 1 >= max_depth {
845 return Err(EngineError::SpawnDepthExceeded {
846 current: d + 1,
847 max: max_depth,
848 });
849 }
850 d + 1
851 }
852 None => 0,
853 };
854
855 let mut task = TaskState::new(task_id_clone.clone(), spec);
856 task.spawn_depth = depth;
857 s.tasks.insert(task_id_clone.clone(), task);
858 s.prompts
859 .insert((task_id_clone.clone(), 1), initial_directive);
860 // Link to the owner session (only Operator tokens match; Worker tokens have no session).
861 if let Some(sess) = s.sessions.values_mut().find(|sess| sess.token_fp == fp) {
862 sess.owned_task_ids.push(task_id_clone.clone());
863 }
864 s.push_event(Event::TaskCreated {
865 task_id: task_id_clone.clone(),
866 });
867 Ok::<(), EngineError>(())
868 })
869 .await??;
870 let _ = self.inner.event_tx.send(Event::TaskCreated {
871 task_id: task_id.clone(),
872 });
873 Ok(task_id)
874 }
875
876 /// Fetch a snapshot of `TaskState` for `task_id`, subject to the
877 /// task-ownership gate (see `verify_token_for_task`).
878 pub async fn read_task_state(
879 &self,
880 token: &CapToken,
881 task_id: &StepId,
882 ) -> Result<TaskState, EngineError> {
883 self.verify_token_for_task(token, Verb::ReadTaskState, task_id)
884 .await?;
885 let task_id = task_id.clone();
886 self.with_state("read_task_state", move |s| {
887 s.tasks
888 .get(&task_id)
889 .cloned()
890 .ok_or_else(|| EngineError::TaskNotFound(task_id.to_string()))
891 })
892 .await?
893 }
894
895 /// Mark `task_id` as `Cancelled` and wake any caller blocked in
896 /// `poll_task` for it.
897 pub async fn cancel_task(&self, token: &CapToken, task_id: &StepId) -> Result<(), EngineError> {
898 self.verify_token_for_task(token, Verb::CancelTask, task_id)
899 .await?;
900 let tid = task_id.clone();
901 self.with_state("cancel_task", move |s| {
902 let task = s
903 .tasks
904 .get_mut(&tid)
905 .ok_or_else(|| EngineError::TaskNotFound(tid.to_string()))?;
906 task.status = TaskStatus::Cancelled;
907 task.updated_at = now_unix();
908 s.push_event(Event::TaskCancelled {
909 task_id: tid.clone(),
910 });
911 Ok::<(), EngineError>(())
912 })
913 .await??;
914 self.wake_task(task_id).await?;
915 Ok(())
916 }
917
918 /// Dispatch a single attempt through the given `spawner`.
919 ///
920 /// The lock is only held for snapshot capture; the actual spawn and
921 /// completion await happen outside the lock (R3 discipline).
922 ///
923 /// Sits on the Domain side of the Data / Domain split. The dispatch
924 /// path itself does not touch big response bodies — those flow through
925 /// the Data plane (`output_store` module + sink / input_inject
926 /// `SpawnerLayer`s) around this method.
927 ///
928 /// The caller does the compile plus `service::linker::link` and
929 /// carries the same stack through each dispatch. Because the spawner
930 /// is passed per-request rather than looked up from engine-global
931 /// state, parallel requests against a single `Engine` instance
932 /// (different Blueprints, different spawners) do not race.
933 ///
934 /// `run_id`, when `Some` (issue #13 run_id propagation —
935 /// `EngineDispatcher` threads it in from its `RunContext`), is
936 /// inserted into `Ctx.meta.runtime["run_id"]` (a plain JSON string)
937 /// alongside `worker_handle`, so `Operator::execute` implementations
938 /// (e.g. `WSOperatorSession`) can read it back and surface it to the
939 /// worker (Spawn directive / prompt). `None` (every pre-existing
940 /// caller / test) omits the key entirely — unchanged behavior.
941 pub async fn dispatch_attempt_with(
942 &self,
943 token: &CapToken,
944 task_id: &StepId,
945 spawner: &Arc<dyn SpawnerAdapter>,
946 run_id: Option<&RunId>,
947 ) -> Result<DispatchOutcome, EngineError> {
948 self.verify_token(token, Verb::DispatchAttempt).await?;
949 let task_id = task_id.clone();
950
951 // 1) Under the lock: increment the attempt number, mark Running, snapshot the
952 // prompt, and pull `operator_info` from the session so we can inject it into Ctx.
953 let fp = token.fingerprint();
954 let tid_for_prep = task_id.clone();
955 let (attempt, agent, session_snapshot, step_ctx) = self
956 .with_state("dispatch.prep", move |s| {
957 let task = s
958 .tasks
959 .get_mut(&tid_for_prep)
960 .ok_or_else(|| EngineError::TaskNotFound(tid_for_prep.to_string()))?;
961 task.attempt += 1;
962 task.status = TaskStatus::Running;
963 task.updated_at = now_unix();
964 // The spawner pulls the prompt via engine.fetch_prompt. In prep,
965 // if the prompts table has no entry for this attempt yet,
966 // fall back and insert `initial_directive` so the subsequent
967 // fetch_prompt succeeds.
968 let attempt = task.attempt;
969 let initial = task.spec.initial_directive.clone();
970 s.prompts
971 .entry((tid_for_prep.clone(), attempt))
972 .or_insert(initial);
973 let task = s
974 .tasks
975 .get(&tid_for_prep)
976 .ok_or_else(|| EngineError::TaskNotFound(tid_for_prep.to_string()))?;
977 let agent = task.spec.agent.clone();
978 // GH #21 Phase 2: re-read `TaskSpec.step_ctx` on EVERY
979 // attempt (not cached once at start_task) so retries and
980 // Run-rekicks all carry the Step tier through to Ctx —
981 // see TaskSpec.step_ctx's doc.
982 let step_ctx = task.spec.step_ctx.clone();
983 // Session snapshot (looked up by token nonce). When no session
984 // exists (worker token invoked directly / test injection), fall
985 // back to None → default OperatorInfo.
986 let sess_clone = s
987 .sessions
988 .values()
989 .find(|sess| sess.token_fp == fp)
990 .cloned();
991 Ok::<_, EngineError>((attempt, agent, sess_clone, step_ctx))
992 })
993 .await??;
994 // BridgeRegistry lookup + per-agent OperatorKind cascade.
995 let operator_info = match session_snapshot {
996 Some(sess) => self.resolve_operator_info(&sess, &agent).await,
997 None => OperatorInfo::default(),
998 };
999
1000 // 2) Outside the lock: worker token mint + spawn.
1001 //
1002 // Session-style mint (max_uses=None). Within one attempt the worker is
1003 // expected to hit `verify_token + fetch_prompt + fetch_data + post_result`
1004 // multiple times in order, so `one_time` would exhaust the token on the
1005 // very first verb. Capability is guarded by (a) the role × verb gate and
1006 // (b) the short TTL (1800s).
1007 let worker_token = self.inner.signer.session(
1008 format!("worker-of-{task_id}"),
1009 Role::Worker,
1010 vec!["*".into()],
1011 Duration::from_secs(1800),
1012 );
1013 let worker_fp = worker_token.fingerprint();
1014 let task_id_for_worker = task_id.clone();
1015 let worker_token_for_store = worker_token.clone();
1016 self.with_state("dispatch.mint_worker", move |s| {
1017 s.tokens.insert(
1018 worker_fp,
1019 CapTokenRecord::from_worker_token(worker_token_for_store, task_id_for_worker),
1020 );
1021 })
1022 .await?;
1023
1024 // Mint a short handle (`wh-XXXXXXXX`) and register it in worker_handles.
1025 // Used by the simplified Bearer path for SubAgents (short-handle form
1026 // avoids base64 copy-paste incidents).
1027 let worker_handle = self.mint_worker_handle(worker_token.fingerprint()).await?;
1028
1029 let mut ctx = Ctx::new(task_id.clone(), attempt, agent.clone());
1030 ctx.operator = operator_info; // activates MainAIMiddleware / Senior bridge
1031 ctx.meta
1032 .runtime
1033 .insert("worker_handle".to_string(), Value::String(worker_handle));
1034 if let Some(rid) = run_id {
1035 ctx.meta
1036 .runtime
1037 .insert(RUN_ID_KEY.to_string(), Value::String(rid.to_string()));
1038 }
1039 // GH #21 Phase 2: the Step tier's resolved context bundle (from
1040 // `TaskSpec.step_ctx`, re-read every attempt above) — consumed by
1041 // `AgentContextMiddleware`, which unpacks its keys ahead of the
1042 // Agent / BP-global tiers.
1043 if let Some(step_ctx) = step_ctx {
1044 ctx.meta.runtime.insert(STEP_CTX_KEY.to_string(), step_ctx);
1045 }
1046
1047 let worker = spawner
1048 .spawn(self, &ctx, task_id.clone(), attempt, worker_token)
1049 .await
1050 .map_err(|e| EngineError::DispatchFailed(e.to_string()))?;
1051
1052 // 3) Outside the lock: await worker.join() (signal-only). WorkerError is
1053 // stringified. The value is fetched via output_tail (sink path).
1054 let signal_result: Result<(), String> = worker.join().await.map_err(|e| e.to_string());
1055
1056 // Pull the last Final from output_tail and use it as the value.
1057 let value_ok: Result<(Value, bool), String> = match signal_result {
1058 Ok(()) => {
1059 let tail = self.output_tail(&task_id, attempt).await;
1060 let last_final = tail.iter().rev().find_map(|ev| match ev {
1061 crate::worker::output::OutputEvent::Final { content, ok } => {
1062 Some((content.clone(), *ok))
1063 }
1064 _ => None,
1065 });
1066 match last_final {
1067 Some((crate::worker::output::ContentRef::Inline { value }, ok)) => {
1068 Ok((value, ok))
1069 }
1070 Some((
1071 crate::worker::output::ContentRef::FileRef {
1072 path,
1073 mime,
1074 size_hint,
1075 },
1076 ok,
1077 )) => Ok((
1078 serde_json::json!({
1079 "file_ref": path.to_string_lossy(),
1080 "mime": mime,
1081 "size_hint": size_hint,
1082 }),
1083 ok,
1084 )),
1085 None => Err("no Final in output_tail".to_string()),
1086 }
1087 }
1088 Err(msg) => Err(msg),
1089 };
1090
1091 // 4) Under the lock: apply (split the borrow scope so push_event and task mut can co-exist).
1092 let outcome = self
1093 .with_state("dispatch.apply", |s| {
1094 if !s.tasks.contains_key(&task_id) {
1095 return Err(EngineError::TaskNotFound(task_id.to_string()));
1096 }
1097 match value_ok {
1098 Ok((value, ok)) => {
1099 let pass = ok;
1100 {
1101 let task = s.tasks.get_mut(&task_id).unwrap();
1102 task.last_result = Some(value.clone());
1103 task.updated_at = now_unix();
1104 task.status = if pass {
1105 TaskStatus::Pass
1106 } else {
1107 TaskStatus::Blocked
1108 };
1109 }
1110 s.push_event(Event::TaskAttemptCompleted {
1111 task_id: task_id.clone(),
1112 attempt,
1113 result: value.clone(),
1114 });
1115 if pass {
1116 s.push_event(Event::TaskPass {
1117 task_id: task_id.clone(),
1118 result: value.clone(),
1119 });
1120 Ok::<_, EngineError>(DispatchOutcome::Pass(value))
1121 } else {
1122 s.push_event(Event::TaskBlocked {
1123 task_id: task_id.clone(),
1124 result: value.clone(),
1125 });
1126 Ok(DispatchOutcome::Blocked(value))
1127 }
1128 }
1129 Err(msg) => {
1130 let task = s.tasks.get_mut(&task_id).unwrap();
1131 task.status = TaskStatus::Blocked;
1132 task.updated_at = now_unix();
1133 Err(EngineError::DispatchFailed(msg))
1134 }
1135 }
1136 })
1137 .await??;
1138
1139 // event broadcast (outside the lock — push_event feeds the in-memory tail; broadcast is a separate path).
1140 let _ = self.inner.event_tx.send(Event::TaskAttemptCompleted {
1141 task_id: task_id.clone(),
1142 attempt,
1143 result: match &outcome {
1144 DispatchOutcome::Pass(v) | DispatchOutcome::Blocked(v) => v.clone(),
1145 _ => Value::Null,
1146 },
1147 });
1148
1149 // Wake any callers waiting in poll_task.
1150 self.wake_task(&task_id).await?;
1151
1152 Ok(outcome)
1153 }
1154
1155 // ═══════════════════════════════════════════════════════════════════════
1156 // Worker-side API (= prompt / data fetch + result post)
1157 // ═══════════════════════════════════════════════════════════════════════
1158
1159 /// Fetch the directive/prompt `Value` for `task_id`'s current attempt.
1160 /// Falls back to `initial_directive` when no prompt has been recorded
1161 /// yet for that attempt. Returns the `Value` end-to-end (issue #18);
1162 /// the render down to `String` happens only at the two consumer
1163 /// boundaries — the Worker HTTP path (`fetch_worker_payload*` →
1164 /// `WorkerPayload.prompt: String`) and the WS Spawn frame text
1165 /// render (`operator_ws::session`).
1166 pub async fn fetch_prompt(
1167 &self,
1168 token: &CapToken,
1169 task_id: &StepId,
1170 ) -> Result<Value, EngineError> {
1171 self.verify_token_for_task(token, Verb::FetchPrompt, task_id)
1172 .await?;
1173 let task_id = task_id.clone();
1174 self.with_state("fetch_prompt", move |s| {
1175 let task = s
1176 .tasks
1177 .get(&task_id)
1178 .ok_or_else(|| EngineError::TaskNotFound(task_id.to_string()))?;
1179 s.prompts
1180 .get(&(task_id.clone(), task.attempt.max(1)))
1181 .cloned()
1182 .ok_or_else(|| {
1183 EngineError::ResourceNotFound(format!(
1184 "prompt({}, attempt={})",
1185 task_id, task.attempt
1186 ))
1187 })
1188 })
1189 .await?
1190 }
1191
1192 /// Combined fetch for `HTTP /v1/worker/prompt`: returns `prompt` +
1193 /// (optional) `system` + `agent` + `attempt` in a single round trip.
1194 /// The verb gate reuses `FetchPrompt` — same semantics as "the worker
1195 /// pulls its task input".
1196 ///
1197 /// `system` is the value written by `OperatorSpawner::spawn` through
1198 /// `bake_worker_system_prompt` when it ran; otherwise `None` (no
1199 /// profile present, or the bake never happened).
1200 pub async fn fetch_worker_payload(
1201 &self,
1202 token: &CapToken,
1203 task_id: &StepId,
1204 ) -> Result<crate::types::WorkerPayload, EngineError> {
1205 self.verify_token_for_task(token, Verb::FetchPrompt, task_id)
1206 .await?;
1207 let task_id_clone = task_id.clone();
1208 self.with_state("fetch_worker_payload", move |s| {
1209 let task = s
1210 .tasks
1211 .get(&task_id_clone)
1212 .ok_or_else(|| EngineError::TaskNotFound(task_id_clone.to_string()))?;
1213 let attempt = task.attempt.max(1);
1214 let prompt = s
1215 .prompts
1216 .get(&(task_id_clone.clone(), attempt))
1217 .cloned()
1218 .ok_or_else(|| {
1219 EngineError::ResourceNotFound(format!(
1220 "prompt({}, attempt={})",
1221 task_id_clone, attempt
1222 ))
1223 })?;
1224 let system = s
1225 .systems
1226 .get(&(task_id_clone.clone(), attempt))
1227 .cloned()
1228 .unwrap_or(None);
1229 let agent = task.spec.agent.clone();
1230 let context = s
1231 .agent_contexts
1232 .get(&(task_id_clone.clone(), attempt))
1233 .cloned();
1234 Ok::<_, EngineError>(crate::types::WorkerPayload {
1235 task_id: task_id_clone.clone(),
1236 attempt,
1237 agent,
1238 prompt: render_directive_to_string(&prompt),
1239 system,
1240 context,
1241 })
1242 })
1243 .await?
1244 }
1245
1246 /// Fetch a worker payload via a short handle. Skips token verification
1247 /// and returns `prompt` + `system` + `agent` + `attempt` in a thin
1248 /// path. The caller is expected to have already resolved `task_id`
1249 /// via `task_id_from_handle` — the handle's presence in
1250 /// `worker_handles` means it was minted server-side and is therefore
1251 /// trusted.
1252 pub async fn fetch_worker_payload_trusted(
1253 &self,
1254 task_id: &StepId,
1255 ) -> Result<crate::types::WorkerPayload, EngineError> {
1256 let task_id_clone = task_id.clone();
1257 self.with_state("fetch_worker_payload_trusted", move |s| {
1258 let task = s
1259 .tasks
1260 .get(&task_id_clone)
1261 .ok_or_else(|| EngineError::TaskNotFound(task_id_clone.to_string()))?;
1262 let attempt = task.attempt.max(1);
1263 let prompt = s
1264 .prompts
1265 .get(&(task_id_clone.clone(), attempt))
1266 .cloned()
1267 .ok_or_else(|| {
1268 EngineError::ResourceNotFound(format!(
1269 "prompt({}, attempt={})",
1270 task_id_clone, attempt
1271 ))
1272 })?;
1273 let system = s
1274 .systems
1275 .get(&(task_id_clone.clone(), attempt))
1276 .cloned()
1277 .unwrap_or(None);
1278 let agent = task.spec.agent.clone();
1279 let context = s
1280 .agent_contexts
1281 .get(&(task_id_clone.clone(), attempt))
1282 .cloned();
1283 Ok::<_, EngineError>(crate::types::WorkerPayload {
1284 task_id: task_id_clone.clone(),
1285 attempt,
1286 agent,
1287 prompt: render_directive_to_string(&prompt),
1288 system,
1289 context,
1290 })
1291 })
1292 .await?
1293 }
1294
1295 /// Read the current attempt number for a task (server-side lookup, no
1296 /// token verification). Used on `HTTP /v1/worker/result` when the
1297 /// worker omits `attempt` and the server has to fill it in.
1298 pub async fn task_attempt(&self, task_id: &StepId) -> Result<u32, EngineError> {
1299 let task_id = task_id.clone();
1300 self.with_state("task_attempt", move |s| {
1301 s.tasks
1302 .get(&task_id)
1303 .map(|t| t.attempt)
1304 .ok_or_else(|| EngineError::TaskNotFound(task_id.to_string()))
1305 })
1306 .await?
1307 }
1308
1309 /// Server-side admin API that lets `OperatorSpawner::spawn` bake the
1310 /// rendered `system_prompt` into engine state. There is no verb gate
1311 /// — the only expected caller is inside the spawner. SubAgents fetch
1312 /// this alongside the prompt on the `/v1/worker/prompt` path.
1313 pub async fn bake_worker_system_prompt(
1314 &self,
1315 task_id: &StepId,
1316 attempt: u32,
1317 system: Option<String>,
1318 ) -> Result<(), EngineError> {
1319 let task_id = task_id.clone();
1320 self.with_state("bake_worker_system_prompt", move |s| {
1321 s.systems.insert((task_id, attempt), system);
1322 })
1323 .await?;
1324 Ok(())
1325 }
1326
1327 /// Fetch an arbitrary named resource previously stored via
1328 /// `set_resource`. Not task-scoped — any valid token with the
1329 /// `FetchData` verb may read any key.
1330 pub async fn fetch_data(&self, token: &CapToken, key: &str) -> Result<Value, EngineError> {
1331 self.verify_token(token, Verb::FetchData).await?;
1332 let key = key.to_string();
1333 self.with_state("fetch_data", move |s| {
1334 s.resources
1335 .get(&key)
1336 .cloned()
1337 .ok_or(EngineError::ResourceNotFound(key))
1338 })
1339 .await?
1340 }
1341
1342 // ───────────────────────────────────────────────────────────────────────
1343 // Output path.
1344 // ───────────────────────────────────────────────────────────────────────
1345
1346 /// Send one output event from inside a `SpawnerAdapter` or worker.
1347 /// Structuring is assumed to be complete by the time we cross the
1348 /// `SpawnerAdapter` boundary; this API just appends to the
1349 /// `OutputStore`, pushes to the `EventLog`, and (for `Final`) emits
1350 /// the `TaskAttemptCompleted` event.
1351 ///
1352 /// This is Domain-side plumbing: it feeds the engine's verdict flow,
1353 /// not the Data-plane store in the `output_store` module. It also
1354 /// does not wake the dispatch path — that is done through the
1355 /// spawner's completion oneshot when the worker terminates.
1356 pub async fn submit_output(
1357 &self,
1358 token: &crate::types::CapToken,
1359 task_id: &StepId,
1360 attempt: u32,
1361 event: crate::worker::output::OutputEvent,
1362 ) -> Result<(), EngineError> {
1363 self.verify_token_for_task(token, crate::types::Verb::EmitOutput, task_id)
1364 .await?;
1365 let task_id_for_apply = task_id.clone();
1366 let event_clone = event.clone();
1367 self.with_state("submit_output", move |s| {
1368 s.output_store
1369 .entry((task_id_for_apply.clone(), attempt))
1370 .or_default()
1371 .push(event_clone.clone());
1372 s.push_event(crate::core::state::Event::WorkerOutput {
1373 task_id: task_id_for_apply,
1374 attempt,
1375 event: event_clone,
1376 });
1377 })
1378 .await?;
1379 Ok(())
1380 }
1381
1382 /// Snapshot the entire output tail for a given `(task_id, attempt)`.
1383 /// Used by the dispatch path when pulling `Final`, and by observers
1384 /// reading the trace.
1385 pub async fn output_tail(
1386 &self,
1387 task_id: &StepId,
1388 attempt: u32,
1389 ) -> Vec<crate::worker::output::OutputEvent> {
1390 let key = (task_id.clone(), attempt);
1391 self.with_state("output_tail", move |s| {
1392 s.output_store.get(&key).cloned().unwrap_or_default()
1393 })
1394 .await
1395 .unwrap_or_default()
1396 }
1397
1398 /// Record an interim `last_result` for `task_id` without changing its
1399 /// `status`. Distinct from the terminal `Final` output event handled
1400 /// through `submit_output` / `dispatch_attempt_with`.
1401 pub async fn post_result(
1402 &self,
1403 token: &CapToken,
1404 task_id: &StepId,
1405 result: Value,
1406 ) -> Result<(), EngineError> {
1407 self.verify_token_for_task(token, Verb::PostResult, task_id)
1408 .await?;
1409 let task_id = task_id.clone();
1410 let result_clone = result.clone();
1411 self.with_state("post_result", move |s| {
1412 let task = s
1413 .tasks
1414 .get_mut(&task_id)
1415 .ok_or_else(|| EngineError::TaskNotFound(task_id.to_string()))?;
1416 task.last_result = Some(result_clone);
1417 task.updated_at = now_unix();
1418 Ok::<(), EngineError>(())
1419 })
1420 .await??;
1421 Ok(())
1422 }
1423
1424 /// Store a named resource value, retrievable later via `fetch_data`.
1425 /// No token is required — this is a server-side/admin-style setter
1426 /// (mirrors `bake_worker_system_prompt`).
1427 pub async fn set_resource(
1428 &self,
1429 key: impl Into<String>,
1430 value: Value,
1431 ) -> Result<(), EngineError> {
1432 let key = key.into();
1433 self.with_state("set_resource", move |s| {
1434 s.resources.insert(key, value);
1435 })
1436 .await?;
1437 Ok(())
1438 }
1439
1440 // ═══════════════════════════════════════════════════════════════════════
1441 // Senior suspend / resume
1442 // ═══════════════════════════════════════════════════════════════════════
1443
1444 /// Ask a question of the Senior, mark the task `Suspended`, and
1445 /// return a `ResumeKey`. The suspended state persists until another
1446 /// task calls `resume(key, answer)`.
1447 ///
1448 /// Resume-side waiting is `Notify`-based, so a caller (typically
1449 /// MainAI) can detach, reattach from a different process, and still
1450 /// pull the answer out via `await_resume(key, timeout)` — the answer
1451 /// is stored inside `EngineState`.
1452 pub async fn query_senior(
1453 &self,
1454 token: &CapToken,
1455 task_id: &StepId,
1456 question: Value,
1457 ) -> Result<ResumeKey, EngineError> {
1458 self.verify_token(token, Verb::QuerySenior).await?;
1459 let task_id = task_id.clone();
1460 let key = ResumeKey::for_senior(&task_id);
1461 let task_notify = self
1462 .with_state("query_senior.notify_ensure", |s| {
1463 s.ensure_task_notify(&task_id)
1464 })
1465 .await?;
1466
1467 let key_clone = key.clone();
1468 let task_id_inner = task_id.clone();
1469 let question_clone = question.clone();
1470 self.with_state("query_senior.suspend", move |s| {
1471 let task = s
1472 .tasks
1473 .get_mut(&task_id_inner)
1474 .ok_or_else(|| EngineError::TaskNotFound(task_id_inner.to_string()))?;
1475 task.status = TaskStatus::Suspended;
1476 task.suspended_on = Some(key_clone.clone());
1477 task.updated_at = now_unix();
1478 s.pending_resumes
1479 .insert(key_clone.clone(), ResumePending::new());
1480 s.push_event(Event::SeniorQueried {
1481 task_id: task_id_inner.clone(),
1482 question: question_clone.clone(),
1483 });
1484 s.push_event(Event::TaskSuspended {
1485 task_id: task_id_inner.clone(),
1486 key: key_clone.clone(),
1487 });
1488 Ok::<(), EngineError>(())
1489 })
1490 .await??;
1491
1492 // Notify callers waiting for a task status change (Running → Suspended).
1493 task_notify.notify_waiters();
1494
1495 let _ = self
1496 .inner
1497 .event_tx
1498 .send(Event::SeniorQueried { task_id, question });
1499 Ok(key)
1500 }
1501
1502 /// Store the answer for a `ResumeKey` in `EngineState` and wake the
1503 /// waiting caller via `Notify`. Also flips the suspended task's
1504 /// status back to `Running` and fires the per-task notifier.
1505 pub async fn resume(&self, key: ResumeKey, answer: Value) -> Result<(), EngineError> {
1506 let answer_for_state = answer.clone();
1507 let answer_for_event = answer.clone();
1508 let key_clone = key.clone();
1509 let (notify, task_notify, task_id_opt) = self
1510 .with_state("resume.set", move |s| {
1511 let pending = s
1512 .pending_resumes
1513 .get_mut(&key_clone)
1514 .ok_or(EngineError::ResumeKeyNotFound)?;
1515 pending.answer = Some(answer_for_state);
1516 let notify = pending.notify.clone();
1517
1518 let task_id = s
1519 .tasks
1520 .iter()
1521 .find(|(_, t)| t.suspended_on.as_ref() == Some(&key_clone))
1522 .map(|(id, _)| id.clone());
1523
1524 let task_notify = task_id.as_ref().map(|tid| s.ensure_task_notify(tid));
1525
1526 if let Some(tid) = &task_id {
1527 if let Some(task) = s.tasks.get_mut(tid) {
1528 task.suspended_on = None;
1529 task.status = TaskStatus::Running;
1530 task.updated_at = now_unix();
1531 }
1532 s.push_event(Event::TaskResumed {
1533 task_id: tid.clone(),
1534 key: key_clone.clone(),
1535 });
1536 s.push_event(Event::SeniorAnswered {
1537 task_id: tid.clone(),
1538 answer: answer_for_event.clone(),
1539 });
1540 }
1541 Ok::<_, EngineError>((notify, task_notify, task_id))
1542 })
1543 .await??;
1544
1545 // Outside the lock: notify_waiters for both the ResumePending and task-status waits.
1546 notify.notify_waiters();
1547 if let Some(n) = task_notify {
1548 n.notify_waiters();
1549 }
1550
1551 if let Some(tid) = task_id_opt {
1552 let _ = self
1553 .inner
1554 .event_tx
1555 .send(Event::TaskResumed { task_id: tid, key });
1556 }
1557 Ok(())
1558 }
1559
1560 /// Wait for the resume answer. Even if the caller (an Operator)
1561 /// detached and reattached, the answer is available immediately here
1562 /// — if it was already stored, this returns without waiting on the
1563 /// notifier.
1564 ///
1565 /// `timeout = Duration::ZERO` performs an instant check without
1566 /// waiting.
1567 pub async fn await_resume(
1568 &self,
1569 key: ResumeKey,
1570 timeout: Duration,
1571 ) -> Result<Value, EngineError> {
1572 // (1) Under the lock: clone the notify handle and check for an existing answer.
1573 let key_clone = key.clone();
1574 let (notify, existing) = self
1575 .with_state("await_resume.snapshot", move |s| {
1576 let pending = s
1577 .pending_resumes
1578 .get(&key_clone)
1579 .ok_or(EngineError::ResumeKeyNotFound)?;
1580 Ok::<_, EngineError>((pending.notify.clone(), pending.answer.clone()))
1581 })
1582 .await??;
1583
1584 // (2) If an answer has already been stored, return immediately (detach / reattach pattern).
1585 if let Some(v) = existing {
1586 return Ok(v);
1587 }
1588
1589 // (3) Outside the lock: wait on the notify with a timeout.
1590 if timeout.is_zero() {
1591 return Err(EngineError::PollTimeout);
1592 }
1593 let waited = tokio::time::timeout(timeout, notify.notified()).await;
1594 if waited.is_err() {
1595 return Err(EngineError::PollTimeout);
1596 }
1597
1598 // (4) Under the lock: re-read the answer (should be present now that we were notified).
1599 let key_clone = key.clone();
1600 self.with_state("await_resume.read", move |s| {
1601 let pending = s
1602 .pending_resumes
1603 .get(&key_clone)
1604 .ok_or(EngineError::ResumeKeyNotFound)?;
1605 pending
1606 .answer
1607 .clone()
1608 .ok_or_else(|| EngineError::Internal("notified but answer missing".into()))
1609 })
1610 .await?
1611 }
1612
1613 // ═══════════════════════════════════════════════════════════════════════
1614 // poll_task — the "wait" path that waits for task-status changes (works for long-poll and regular wait).
1615 // ═══════════════════════════════════════════════════════════════════════
1616
1617 /// Wait until the task's status **transitions to terminal or
1618 /// `Suspended`**, then return the latest `TaskState`. Returns
1619 /// immediately if the task is already in a terminal state.
1620 /// Exceeding the timeout returns `EngineError::PollTimeout`.
1621 ///
1622 /// A `hold` of `Duration::from_secs(0)` returns a snapshot immediately
1623 /// (no wait). Larger holds — tens of minutes up to days — are fine;
1624 /// the wait state is kept in memory inside the engine and does not
1625 /// degrade.
1626 pub async fn poll_task(
1627 &self,
1628 token: &CapToken,
1629 task_id: &StepId,
1630 hold: Duration,
1631 ) -> Result<TaskState, EngineError> {
1632 self.verify_token_for_task(token, Verb::PollTask, task_id)
1633 .await?;
1634 let task_id_inner = task_id.clone();
1635
1636 // (1) Under the lock: take a snapshot and clone task_notify.
1637 let (state, notify) = self
1638 .with_state("poll_task.snapshot", move |s| {
1639 let task = s
1640 .tasks
1641 .get(&task_id_inner)
1642 .cloned()
1643 .ok_or_else(|| EngineError::TaskNotFound(task_id_inner.to_string()))?;
1644 let notify = s.ensure_task_notify(&task_id_inner);
1645 Ok::<_, EngineError>((task, notify))
1646 })
1647 .await??;
1648
1649 // (2) Immediate-return condition: already terminal / Suspended (nothing left to wait on).
1650 if matches!(
1651 state.status,
1652 TaskStatus::Pass | TaskStatus::Blocked | TaskStatus::Cancelled | TaskStatus::Suspended
1653 ) {
1654 return Ok(state);
1655 }
1656 if hold.is_zero() {
1657 return Ok(state);
1658 }
1659
1660 // (3) Outside the lock: wait on Notify with a timeout.
1661 let waited = tokio::time::timeout(hold, notify.notified()).await;
1662 if waited.is_err() {
1663 return Err(EngineError::PollTimeout);
1664 }
1665
1666 // (4) Under the lock: take a fresh snapshot.
1667 let task_id_inner = task_id.clone();
1668 self.with_state("poll_task.reread", move |s| {
1669 s.tasks
1670 .get(&task_id_inner)
1671 .cloned()
1672 .ok_or_else(|| EngineError::TaskNotFound(task_id_inner.to_string()))
1673 })
1674 .await?
1675 }
1676
1677 // ═══════════════════════════════════════════════════════════════════════
1678 // Background: heartbeat miss → detach loop
1679 // ═══════════════════════════════════════════════════════════════════════
1680
1681 /// Background loop that scans sessions every `heartbeat_interval` and
1682 /// flips `attached = false` on any session whose `last_seen` exceeds
1683 /// `heartbeat_miss_threshold * interval`.
1684 ///
1685 /// The tasks themselves are kept (assuming
1686 /// `keepalive_on_idle = true`), so another client can reattach with
1687 /// the same token and resume immediately. Dropping the returned
1688 /// `JoinHandle` does not stop the loop — the handle exists so callers
1689 /// who want to abort can hold onto it.
1690 pub fn start_detach_loop(&self) -> tokio::task::JoinHandle<()> {
1691 let engine = self.clone();
1692 let cfg = self.inner.cfg.long_hold.clone();
1693 let interval = cfg.heartbeat_interval;
1694 let miss_secs = cfg.heartbeat_interval.as_secs() * cfg.heartbeat_miss_threshold as u64;
1695
1696 tokio::spawn(async move {
1697 let mut ticker = tokio::time::interval(interval);
1698 ticker.tick().await; // first tick is immediate
1699 loop {
1700 ticker.tick().await;
1701 let now = now_unix();
1702 let detached = engine
1703 .with_state("detach_loop.scan", |s| {
1704 let mut detached = Vec::new();
1705 for (sid, sess) in s.sessions.iter_mut() {
1706 if !sess.attached {
1707 continue;
1708 }
1709 if now.saturating_sub(sess.last_seen) >= miss_secs {
1710 sess.attached = false;
1711 detached.push(sid.clone());
1712 }
1713 }
1714 for sid in &detached {
1715 s.push_event(Event::SessionDetached {
1716 session_id: sid.clone(),
1717 });
1718 }
1719 detached
1720 })
1721 .await
1722 .unwrap_or_default();
1723 for sid in detached {
1724 let _ = engine
1725 .inner
1726 .event_tx
1727 .send(Event::SessionDetached { session_id: sid });
1728 }
1729 }
1730 })
1731 }
1732
1733 /// Helper: wake a task whose status has changed. Called from the
1734 /// method body outside the lock.
1735 async fn wake_task(&self, task_id: &StepId) -> Result<(), EngineError> {
1736 let task_id = task_id.clone();
1737 let notify_opt = self
1738 .with_state("wake_task.get_notify", move |s| {
1739 s.task_notifies.get(&task_id).cloned()
1740 })
1741 .await?;
1742 if let Some(n) = notify_opt {
1743 n.notify_waiters();
1744 }
1745 Ok(())
1746 }
1747}
1748
1749// ─── UT: issue #14 — token store keyed by fingerprint, not nonce ────────────
1750#[cfg(test)]
1751mod token_fingerprint_store_tests {
1752 use super::*;
1753
1754 /// A token that was never attached fails verify with a `TokenNotFound`
1755 /// that carries the fingerprint — never the nonce. The error string can
1756 /// surface in HTTP error bodies, so this is the secret-hygiene contract.
1757 #[tokio::test]
1758 async fn verify_unknown_token_reports_fingerprint_not_nonce() {
1759 let engine = Engine::new(EngineCfg::default());
1760 // Signed by the engine's own signer (sig passes) but never inserted
1761 // into the store — verify must fail at step (4), the store lookup.
1762 let token = engine.signer().session(
1763 "ghost",
1764 Role::Operator,
1765 vec!["*".into()],
1766 Duration::from_secs(60),
1767 );
1768 let err = engine
1769 .verify_token(&token, Verb::ReadTaskState)
1770 .await
1771 .expect_err("token is not in the store");
1772 let msg = err.to_string();
1773 assert!(
1774 msg.contains(&token.fingerprint()),
1775 "error must carry the fingerprint: {msg}"
1776 );
1777 assert!(
1778 !msg.contains(&token.nonce),
1779 "error must not leak the nonce: {msg}"
1780 );
1781 }
1782
1783 /// attach → verify → heartbeat → detach all resolve the session /
1784 /// token record through fingerprint keys (mint/verify lifecycle
1785 /// regression guard for the issue #14 key migration).
1786 #[tokio::test]
1787 async fn attach_verify_heartbeat_detach_cycle_with_fp_keying() {
1788 let engine = Engine::new(EngineCfg::default());
1789 let token = engine
1790 .attach("op-1", Role::Operator, Duration::from_secs(60))
1791 .await
1792 .expect("attach");
1793 engine
1794 .verify_token(&token, Verb::ReadTaskState)
1795 .await
1796 .expect("verify consumes via fp key");
1797 engine
1798 .heartbeat(&token)
1799 .await
1800 .expect("heartbeat finds the session by fp");
1801 engine
1802 .detach(&token)
1803 .await
1804 .expect("detach finds the session by fp");
1805 }
1806}
1807
1808// ─── UT: `OperatorKind` "Runtime Global" tier — `Option` semantics ─────────
1809//
1810// Regression coverage for the "explicit Automate is indistinguishable from
1811// unspecified" defect: `OperatorSession.operator_kind` (and the
1812// `attach_with_ids` `kind` parameter it stores) is `Option<OperatorKind>`,
1813// so `Some(Automate)` is an explicit Runtime Global request that must
1814// outrank `bp_global`, while `None` must let `bp_global` decide. Exercises
1815// the real `resolve_operator_info` cascade path (not just
1816// `collapse_operator_kind` in isolation), attaching via `attach_with_ids`
1817// exactly as `TaskLaunchService::launch` does.
1818#[cfg(test)]
1819mod resolve_operator_info_runtime_global_tests {
1820 use super::*;
1821
1822 async fn attach_and_resolve(
1823 runtime_global: Option<OperatorKind>,
1824 bp_global: Option<OperatorKind>,
1825 ) -> OperatorInfo {
1826 let engine = Engine::new(EngineCfg::default());
1827 let token = engine
1828 .attach_with_ids(
1829 "ut-op",
1830 Role::Operator,
1831 Duration::from_secs(30),
1832 runtime_global,
1833 None,
1834 None,
1835 None,
1836 HashMap::new(),
1837 HashMap::new(),
1838 bp_global,
1839 )
1840 .await
1841 .expect("attach_with_ids ok");
1842 let session = engine
1843 .with_state("test.find_session", |s| {
1844 s.sessions
1845 .values()
1846 .find(|sess| sess.token_fp == token.fingerprint())
1847 .cloned()
1848 })
1849 .await
1850 .expect("with_state ok")
1851 .expect("session present after attach_with_ids");
1852 engine.resolve_operator_info(&session, "agent-x").await
1853 }
1854
1855 #[tokio::test]
1856 async fn explicit_some_automate_outranks_bp_global_main_ai() {
1857 // Runtime Global explicitly requests Automate; bp_global is MainAi.
1858 // The explicit `Some(Automate)` must win — this is exactly the case
1859 // the old `== OperatorKind::default()` convention got wrong (it
1860 // could not tell "explicitly Automate" from "unspecified" and would
1861 // have let `bp_global` (MainAi) take over instead).
1862 let info =
1863 attach_and_resolve(Some(OperatorKind::Automate), Some(OperatorKind::MainAi)).await;
1864 assert_eq!(
1865 info.kind,
1866 OperatorKind::Automate,
1867 "explicit Some(Automate) runtime_global must outrank bp_global MainAi"
1868 );
1869 }
1870
1871 #[tokio::test]
1872 async fn none_lets_bp_global_main_ai_win() {
1873 // Runtime Global left unspecified (`None`); bp_global is MainAi.
1874 // With nothing more specific set, `bp_global` must decide.
1875 let info = attach_and_resolve(None, Some(OperatorKind::MainAi)).await;
1876 assert_eq!(
1877 info.kind,
1878 OperatorKind::MainAi,
1879 "None runtime_global must let bp_global MainAi win"
1880 );
1881 }
1882}
1883
1884/// issue #13 run_id propagation: `dispatch_attempt_with`'s `run_id` param
1885/// must land in `Ctx.meta.runtime["run_id"]` (the same slot pattern as the
1886/// pre-existing `worker_handle`), or be omitted entirely when `None`. Same
1887/// `CtxProbe` shape as `middleware::worker_binding`'s test module — an
1888/// inner `SpawnerAdapter` that snapshots the `Ctx` it was called with and
1889/// fails the spawn (only the ctx snapshot matters here).
1890#[cfg(test)]
1891mod dispatch_attempt_with_run_id_tests {
1892 use super::*;
1893 use crate::worker::adapter::{SpawnError, SpawnerAdapter};
1894 use crate::worker::Worker;
1895 use std::sync::Mutex as StdMutex;
1896
1897 struct CtxProbe {
1898 seen: Arc<StdMutex<Option<Ctx>>>,
1899 }
1900
1901 #[async_trait::async_trait]
1902 impl SpawnerAdapter for CtxProbe {
1903 async fn spawn(
1904 &self,
1905 _engine: &Engine,
1906 ctx: &Ctx,
1907 _task_id: StepId,
1908 _attempt: u32,
1909 _token: CapToken,
1910 ) -> Result<Box<dyn Worker>, SpawnError> {
1911 *self.seen.lock().unwrap() = Some(ctx.clone());
1912 Err(SpawnError::Internal("probe stop".into()))
1913 }
1914 }
1915
1916 async fn dispatch_with_probe(run_id: Option<&RunId>) -> Ctx {
1917 let engine = Engine::new(EngineCfg::default());
1918 let token = engine
1919 .attach("ut-op", Role::Operator, Duration::from_secs(30))
1920 .await
1921 .expect("attach");
1922 let tid = engine
1923 .start_task(
1924 &token,
1925 TaskSpec {
1926 agent: "probe".into(),
1927 initial_directive: "hi".into(),
1928 step_ctx: None,
1929 },
1930 )
1931 .await
1932 .expect("start_task");
1933 let seen: Arc<StdMutex<Option<Ctx>>> = Arc::new(StdMutex::new(None));
1934 let spawner: Arc<dyn SpawnerAdapter> = Arc::new(CtxProbe { seen: seen.clone() });
1935 // The probe always errors the spawn (`SpawnError::Internal`); we
1936 // only care about the `Ctx` snapshot it captured, so the dispatch
1937 // outcome itself (`Err`) is discarded.
1938 let _ = engine
1939 .dispatch_attempt_with(&token, &tid, &spawner, run_id)
1940 .await;
1941 let captured = seen.lock().unwrap().clone();
1942 captured.expect("inner ctx captured")
1943 }
1944
1945 #[tokio::test]
1946 async fn run_id_lands_in_ctx_meta_runtime_when_some() {
1947 let run_id = RunId::new();
1948 let observed = dispatch_with_probe(Some(&run_id)).await;
1949 assert_eq!(
1950 observed.meta.runtime.get("run_id").and_then(|v| v.as_str()),
1951 Some(run_id.as_str()),
1952 "ctx.meta.runtime[\"run_id\"] must carry the run_id passed to dispatch_attempt_with"
1953 );
1954 }
1955
1956 #[tokio::test]
1957 async fn run_id_key_absent_when_none() {
1958 let observed = dispatch_with_probe(None).await;
1959 assert!(
1960 !observed.meta.runtime.contains_key("run_id"),
1961 "no run_id key must be injected when dispatch_attempt_with is called with None"
1962 );
1963 }
1964}
1965
1966/// GH #21 Phase 2: `TaskSpec.step_ctx` must land in
1967/// `Ctx.meta.runtime[STEP_CTX_KEY]` — re-read from the spec on EVERY
1968/// attempt (the prep closure re-reads `task.spec.step_ctx` every call, not
1969/// caching it once at `start_task`), so a retry (attempt 2) carries it
1970/// too. Same `CtxProbe` shape as `dispatch_attempt_with_run_id_tests`.
1971#[cfg(test)]
1972mod dispatch_attempt_with_step_ctx_tests {
1973 use super::*;
1974 use crate::worker::adapter::{SpawnError, SpawnerAdapter};
1975 use crate::worker::Worker;
1976 use std::sync::Mutex as StdMutex;
1977
1978 struct CtxProbe {
1979 seen: Arc<StdMutex<Option<Ctx>>>,
1980 }
1981
1982 #[async_trait::async_trait]
1983 impl SpawnerAdapter for CtxProbe {
1984 async fn spawn(
1985 &self,
1986 _engine: &Engine,
1987 ctx: &Ctx,
1988 _task_id: StepId,
1989 _attempt: u32,
1990 _token: CapToken,
1991 ) -> Result<Box<dyn Worker>, SpawnError> {
1992 *self.seen.lock().unwrap() = Some(ctx.clone());
1993 Err(SpawnError::Internal("probe stop".into()))
1994 }
1995 }
1996
1997 #[tokio::test]
1998 async fn step_ctx_lands_in_ctx_meta_runtime_on_attempt_1_and_2() {
1999 let engine = Engine::new(EngineCfg::default());
2000 let token = engine
2001 .attach("ut-op", Role::Operator, Duration::from_secs(30))
2002 .await
2003 .expect("attach");
2004 let tid = engine
2005 .start_task(
2006 &token,
2007 TaskSpec {
2008 agent: "probe".into(),
2009 initial_directive: "hi".into(),
2010 step_ctx: Some(serde_json::json!({ "work_dir": "/step" })),
2011 },
2012 )
2013 .await
2014 .expect("start_task");
2015 let seen: Arc<StdMutex<Option<Ctx>>> = Arc::new(StdMutex::new(None));
2016 let spawner: Arc<dyn SpawnerAdapter> = Arc::new(CtxProbe { seen: seen.clone() });
2017
2018 // The probe always errors the spawn; only the ctx snapshot matters.
2019 let _ = engine
2020 .dispatch_attempt_with(&token, &tid, &spawner, None)
2021 .await;
2022 let first = seen
2023 .lock()
2024 .unwrap()
2025 .clone()
2026 .expect("attempt 1 ctx captured");
2027 assert_eq!(
2028 first.meta.runtime.get(STEP_CTX_KEY),
2029 Some(&serde_json::json!({ "work_dir": "/step" })),
2030 "attempt 1 must carry TaskSpec.step_ctx in ctx.meta.runtime[STEP_CTX_KEY]"
2031 );
2032
2033 let _ = engine
2034 .dispatch_attempt_with(&token, &tid, &spawner, None)
2035 .await;
2036 let second = seen
2037 .lock()
2038 .unwrap()
2039 .clone()
2040 .expect("attempt 2 ctx captured");
2041 assert_eq!(
2042 second.meta.runtime.get(STEP_CTX_KEY),
2043 Some(&serde_json::json!({ "work_dir": "/step" })),
2044 "attempt 2 (retry) must ALSO carry TaskSpec.step_ctx — prep re-reads the spec every attempt"
2045 );
2046 }
2047
2048 #[tokio::test]
2049 async fn step_ctx_key_absent_when_none() {
2050 let engine = Engine::new(EngineCfg::default());
2051 let token = engine
2052 .attach("ut-op", Role::Operator, Duration::from_secs(30))
2053 .await
2054 .expect("attach");
2055 let tid = engine
2056 .start_task(
2057 &token,
2058 TaskSpec {
2059 agent: "probe".into(),
2060 initial_directive: "hi".into(),
2061 step_ctx: None,
2062 },
2063 )
2064 .await
2065 .expect("start_task");
2066 let seen: Arc<StdMutex<Option<Ctx>>> = Arc::new(StdMutex::new(None));
2067 let spawner: Arc<dyn SpawnerAdapter> = Arc::new(CtxProbe { seen: seen.clone() });
2068 let _ = engine
2069 .dispatch_attempt_with(&token, &tid, &spawner, None)
2070 .await;
2071 let observed = seen.lock().unwrap().clone().expect("ctx captured");
2072 assert!(
2073 !observed.meta.runtime.contains_key(STEP_CTX_KEY),
2074 "no step_ctx key must be injected when TaskSpec.step_ctx is None"
2075 );
2076 }
2077}
2078
2079// ─── issue #18: `TaskSpec.initial_directive` `Value` pass-through ──────────
2080#[cfg(test)]
2081mod initial_directive_value_passthrough_tests {
2082 use super::*;
2083
2084 async fn seeded_engine(initial_directive: Value) -> (Engine, CapToken, StepId) {
2085 let engine = Engine::new(EngineCfg::default());
2086 let op_token = engine
2087 .attach("ut-op", Role::Operator, Duration::from_secs(30))
2088 .await
2089 .expect("attach");
2090 let task_id = engine
2091 .start_task(
2092 &op_token,
2093 TaskSpec {
2094 agent: "planner".to_string(),
2095 initial_directive,
2096 step_ctx: None,
2097 },
2098 )
2099 .await
2100 .expect("start_task");
2101 (engine, op_token, task_id)
2102 }
2103
2104 /// Mint + register a `Role::Worker` token the same way
2105 /// `dispatch_attempt_with` does — `fetch_prompt` is worker-verb-gated.
2106 async fn mint_worker_token(engine: &Engine, task_id: &StepId) -> CapToken {
2107 let worker_token = engine.signer().session(
2108 format!("worker-of-{task_id}"),
2109 Role::Worker,
2110 vec!["*".into()],
2111 Duration::from_secs(600),
2112 );
2113 let fp = worker_token.fingerprint();
2114 let record = CapTokenRecord::from_worker_token(worker_token.clone(), task_id.clone());
2115 engine
2116 .with_state("test.mint_worker", move |s| {
2117 s.tokens.insert(fp, record);
2118 })
2119 .await
2120 .expect("mint worker token");
2121 worker_token
2122 }
2123
2124 /// `EngineDispatcher::dispatch` no longer stringifies the evaluated
2125 /// `Step.in` value before seeding `TaskSpec.initial_directive` — an
2126 /// Object seed must round-trip through `start_task` /
2127 /// `read_task_state` byte-for-byte as the same `Value::Object`, not a
2128 /// JSON-stringified `Value::String`.
2129 #[tokio::test]
2130 async fn object_seed_passes_through_task_spec_unchanged() {
2131 let seed = serde_json::json!({"key": "value"});
2132 let (engine, token, task_id) = seeded_engine(seed.clone()).await;
2133 let state = engine
2134 .read_task_state(&token, &task_id)
2135 .await
2136 .expect("read_task_state");
2137 assert_eq!(
2138 state.spec.initial_directive, seed,
2139 "TaskSpec.initial_directive must equal the raw Object seed, not a stringified copy"
2140 );
2141 }
2142
2143 /// `Engine::fetch_prompt` returns the `Value` end-to-end (issue #18):
2144 /// an Object seed stays a `Value::Object` and is not stringified in
2145 /// the engine layer. The Worker HTTP boundary
2146 /// (`fetch_worker_payload*`) is what performs the render down to a
2147 /// JSON literal `String` for `WorkerPayload.prompt`.
2148 #[tokio::test]
2149 async fn object_seed_passes_through_fetch_prompt_as_value() {
2150 let seed = serde_json::json!({"key": "value"});
2151 let (engine, _token, task_id) = seeded_engine(seed.clone()).await;
2152 let worker_token = mint_worker_token(&engine, &task_id).await;
2153 let prompt = engine
2154 .fetch_prompt(&worker_token, &task_id)
2155 .await
2156 .expect("fetch_prompt");
2157 assert_eq!(
2158 prompt, seed,
2159 "fetch_prompt must return the raw Object Value, not a stringified copy"
2160 );
2161 }
2162
2163 /// The Worker HTTP boundary is the render point: `fetch_worker_payload*`
2164 /// coerces the stored `Value` down to `WorkerPayload.prompt: String`
2165 /// (JSON-literal shape for non-strings). Verifies the boundary render
2166 /// stays intact for an Object seed.
2167 #[tokio::test]
2168 async fn object_seed_renders_as_json_literal_at_worker_payload_boundary() {
2169 let seed = serde_json::json!({"key": "value"});
2170 let (engine, _token, task_id) = seeded_engine(seed).await;
2171 let worker_token = mint_worker_token(&engine, &task_id).await;
2172 let payload = engine
2173 .fetch_worker_payload(&worker_token, &task_id)
2174 .await
2175 .expect("fetch_worker_payload");
2176 assert_eq!(
2177 payload.prompt, r#"{"key":"value"}"#,
2178 "WorkerPayload.prompt must be the JSON literal String render of the Value seed"
2179 );
2180 }
2181
2182 /// A `String` seed is unaffected — still passes through verbatim, both
2183 /// as the `TaskSpec.initial_directive` `Value` and as the Worker
2184 /// `fetch_prompt` return (issue #18 Invariant 2).
2185 #[tokio::test]
2186 async fn string_seed_passes_through_unchanged() {
2187 let (engine, token, task_id) = seeded_engine(serde_json::json!("do the thing")).await;
2188 let state = engine
2189 .read_task_state(&token, &task_id)
2190 .await
2191 .expect("read_task_state");
2192 assert_eq!(
2193 state.spec.initial_directive,
2194 serde_json::json!("do the thing")
2195 );
2196 let worker_token = mint_worker_token(&engine, &task_id).await;
2197 let prompt = engine
2198 .fetch_prompt(&worker_token, &task_id)
2199 .await
2200 .expect("fetch_prompt");
2201 assert_eq!(prompt, serde_json::json!("do the thing"));
2202 }
2203}