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