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
98/// Renders a [`crate::worker::output::ContentRef`] down to the `Value` shape
99/// the BP-chain / `DispatchOutcome` consume. `Inline` passes its `value`
100/// through verbatim; `FileRef` is stringified into the same
101/// `{"file_ref", "mime", "size_hint"}` shape `materialize_final_submission`
102/// uses for its own file-materialize projection — one canonical
103/// stringification, not two independently-maintained copies (GH #36 ST1:
104/// shared by both the `Final`-pull and the `Artifact`-parts fold in
105/// [`Engine::dispatch_attempt_with`]'s doc).
106fn content_ref_to_value(content: crate::worker::output::ContentRef) -> Value {
107 match content {
108 crate::worker::output::ContentRef::Inline { value } => value,
109 crate::worker::output::ContentRef::FileRef {
110 path,
111 mime,
112 size_hint,
113 } => serde_json::json!({
114 "file_ref": path.to_string_lossy(),
115 "mime": mime,
116 "size_hint": size_hint,
117 }),
118 }
119}
120
121/// [`Engine::dispatch_attempt_with`]'s Final-pull assembly (GH #36 ST1:
122/// named multi-part worker output), factored out as a pure function of the
123/// output-event tail so it is unit-testable without a live `Engine` /
124/// spawner.
125///
126/// Finds the LAST `Final` event in `tail` (mirrors the pre-GH-#36 pull:
127/// "last Final wins" if more than one was ever appended) and folds every
128/// `Artifact` event in the SAME tail WHOSE NAME APPEARS IN `staged_names`
129/// into a `"parts"` object keyed by `Artifact.name` — walked in tail (=
130/// event-append) order, so a name staged more than once within the attempt
131/// is last-write-wins (`Map` insert semantics, not an accumulating list;
132/// `Engine::stage_worker_artifact_trusted`'s doc). `staged_names` is the
133/// WORKER's own opt-in allowlist (`EngineState.worker_artifact_names`'s
134/// doc) — an `Artifact` on the tail whose name is NOT in `staged_names`
135/// (e.g. `AfterRunAuditMiddleware`'s `"audit:<step_ref>"` sidecar finding)
136/// is left alone, exactly as before GH #36; this is what keeps an audited
137/// step's BP-chain value byte-identical when the worker itself never
138/// staged a part.
139///
140/// At least one matching part: the returned value is `{"out": <final
141/// value>, "parts": {<name>: <value>, ...}}`. Zero matching parts: the
142/// returned value is the plain final value, unchanged from the pre-GH-#36
143/// shape — this is the back-compat guarantee, not an incidental default.
144///
145/// `None` when `tail` carries no `Final` at all (the caller's pre-existing
146/// "no Final in output_tail" error path).
147fn fold_final_and_parts(
148 tail: &[crate::worker::output::OutputEvent],
149 staged_names: &[String],
150) -> Option<(Value, bool)> {
151 let (final_content, ok) = tail.iter().rev().find_map(|ev| match ev {
152 crate::worker::output::OutputEvent::Final { content, ok } => Some((content.clone(), *ok)),
153 _ => None,
154 })?;
155 let final_value = content_ref_to_value(final_content);
156
157 let mut parts = serde_json::Map::new();
158 for ev in tail {
159 if let crate::worker::output::OutputEvent::Artifact { name, content } = ev {
160 if staged_names.iter().any(|staged| staged == name) {
161 parts.insert(name.clone(), content_ref_to_value(content.clone()));
162 }
163 }
164 }
165
166 let value = if parts.is_empty() {
167 final_value
168 } else {
169 serde_json::json!({ "out": final_value, "parts": Value::Object(parts) })
170 };
171 Some((value, ok))
172}
173
174impl Engine {
175 /// Backwards-compatible constructor that starts the engine without a
176 /// layer registry, preserving the signature already used by ~88
177 /// existing call sites. Use this when automatic middleware wrapping
178 /// at bind time is not needed. Callers such as `mlua-swarm-server` go through
179 /// `new_with_layers(cfg, registry)` to enable the hint-resolution path.
180 pub fn new(cfg: EngineCfg) -> Self {
181 Self::new_with_layers(cfg, crate::middleware::LayerRegistry::new())
182 }
183
184 /// Construct an `Engine` with an explicit `LayerRegistry`, enabling
185 /// hint-resolution: `spawner_hints.layers` declared on a `Blueprint`
186 /// are resolved against this registry when the spawner stack is bound
187 /// at `service::linker::link` time.
188 pub fn new_with_layers(
189 cfg: EngineCfg,
190 layer_registry: crate::middleware::LayerRegistry,
191 ) -> Self {
192 let (event_tx, _) = broadcast::channel(256);
193 let signer = TokenSigner::new(&cfg.token_secret);
194 Self {
195 inner: Arc::new(EngineInner {
196 state: Mutex::new(EngineState::new()),
197 cfg,
198 signer,
199 gate: default_role_verb_table(),
200 event_tx,
201 senior_bridges: tokio::sync::RwLock::new(HashMap::new()),
202 spawn_hooks: tokio::sync::RwLock::new(HashMap::new()),
203 operators: tokio::sync::RwLock::new(HashMap::new()),
204 layer_registry,
205 data_store: std::sync::RwLock::new(None),
206 }),
207 }
208 }
209
210 /// Rebuild this `Engine` with a different `RoleVerbGate`. The gate is
211 /// treated as fixed-at-build-time, so this constructs a fresh
212 /// `EngineInner` (fresh empty `EngineState`) rather than mutating in
213 /// place — mainly a testing convenience for swapping gate rules.
214 pub fn with_gate(self, gate: RoleVerbGate) -> Self {
215 // The gate is fixed at build time — the intent is to build a fresh
216 // instance rather than mutating in place. As a testing convenience we
217 // do allow swapping the inner Arc. Simpler form: just rebuild
218 // Arc<EngineInner>.
219 let inner = Arc::new(EngineInner {
220 state: Mutex::new(EngineState::new()),
221 cfg: self.inner.cfg.clone(),
222 signer: self.inner.signer.clone(),
223 gate,
224 event_tx: self.inner.event_tx.clone(),
225 senior_bridges: tokio::sync::RwLock::new(HashMap::new()),
226 spawn_hooks: tokio::sync::RwLock::new(HashMap::new()),
227 operators: tokio::sync::RwLock::new(HashMap::new()),
228 layer_registry: self.inner.layer_registry.clone(),
229 data_store: std::sync::RwLock::new(None),
230 });
231 Self { inner }
232 }
233
234 // ═══════════════════════════════════════════════════════════════════════
235 // Accessors. Production code drives execution through compile +
236 // `service::linker::link` + `dispatch_attempt_with(spawner)` inside
237 // `TaskLaunchService`; `Engine` itself is a pure execution surface — it
238 // does not own a BlueprintStore / EnhanceAdapter / Compiler, nor a
239 // global spawner (the spawner is carried per-request, never stashed on
240 // the engine).
241 // ═══════════════════════════════════════════════════════════════════════
242
243 /// Access the `EngineCfg` this engine was built with.
244 pub fn cfg(&self) -> &EngineCfg {
245 &self.inner.cfg
246 }
247
248 /// Expose the internal `LayerRegistry` — used when deriving a
249 /// sub-engine that needs the same registry re-injected. The
250 /// per-request sub-engine in `mlua-swarm-server` reads the parent engine's
251 /// registry through this accessor and passes it to
252 /// `Engine::new_with_layers(cfg, parent.layer_registry().clone())`.
253 pub fn layer_registry(&self) -> &crate::middleware::LayerRegistry {
254 &self.inner.layer_registry
255 }
256
257 /// Access the `TokenSigner` used to mint/verify `CapToken`s.
258 pub fn signer(&self) -> &TokenSigner {
259 &self.inner.signer
260 }
261
262 /// Clone a handle to the process-wide `Event` broadcast sender. Prefer
263 /// `subscribe` for a ready-to-use receiver.
264 pub fn event_tx(&self) -> broadcast::Sender<Event> {
265 self.inner.event_tx.clone()
266 }
267
268 /// Subscribe to the engine's `Event` broadcast stream.
269 pub fn subscribe(&self) -> EventStream {
270 self.inner.event_tx.subscribe()
271 }
272
273 /// Wires the Data-plane [`crate::store::output::OutputStore`] backend
274 /// used by `submit_output` / `submit_worker_result_trusted`'s
275 /// submit-time projection sink (subtask-4 / ST2 rework — see
276 /// `submit_output`'s doc). Synchronous (a plain `std::sync::RwLock`
277 /// write) so a caller can wire it up at boot from a non-`async`
278 /// context (`mlua-swarm-server`'s router builder passes the same
279 /// `Arc` it hands to its `AppState.data_store`, so `POST
280 /// /v1/data/emit` and every worker's ordinary `/v1/worker/submit` land
281 /// in the one store). Calling this more than once replaces the
282 /// previous backend; not calling it at all (the default) preserves
283 /// pre-subtask-4 behavior exactly — `submit_output` only touches the
284 /// Domain-plane `EngineState.output_store` HashMap.
285 pub fn set_output_store(&self, store: Arc<dyn crate::store::output::OutputStore>) {
286 let mut guard = self
287 .inner
288 .data_store
289 .write()
290 .unwrap_or_else(|poisoned| poisoned.into_inner());
291 *guard = Some(store);
292 }
293
294 /// Clones the currently-wired Data-plane store handle, if any. Kept
295 /// private and side-effect-free (no lock held past this call) —
296 /// callers (`materialize_final_submission`) do their actual `.append`
297 /// work outside of any lock.
298 fn output_store_backend(&self) -> Option<Arc<dyn crate::store::output::OutputStore>> {
299 self.inner
300 .data_store
301 .read()
302 .unwrap_or_else(|poisoned| poisoned.into_inner())
303 .clone()
304 }
305
306 // ═══════════════════════════════════════════════════════════════════════
307 // §7 with_state — single Mutex + R1-R4 (try_lock + bounded retry + max-hold panic)
308 // ═══════════════════════════════════════════════════════════════════════
309
310 /// The closure is a **sync** `FnOnce` — you cannot pass an async
311 /// closure, which enforces R3 at the type level. Exceeding `max_hold`
312 /// panics so that R4 violations surface immediately.
313 pub async fn with_state<F, R>(&self, op: &'static str, f: F) -> Result<R, EngineError>
314 where
315 F: FnOnce(&mut EngineState) -> R,
316 {
317 let cfg = &self.inner.cfg;
318
319 // R2: try_lock + bounded retry
320 let mut guard_opt = None;
321 for attempt in 0..=cfg.max_retry {
322 match self.inner.state.try_lock() {
323 Ok(g) => {
324 guard_opt = Some(g);
325 break;
326 }
327 Err(_) if cfg.try_only => return Err(EngineError::LockBusy(op)),
328 Err(_) => {
329 let backoff = cfg.backoff_ms_step * (attempt as u64 + 1);
330 tokio::time::sleep(Duration::from_millis(backoff)).await;
331 }
332 }
333 }
334 let mut guard = guard_opt.ok_or(EngineError::LockBusyAfterRetry(op))?;
335
336 // R4: max_hold guard
337 let start = Instant::now();
338 let result = f(&mut guard);
339 let elapsed_ms = start.elapsed().as_millis();
340 drop(guard);
341
342 if elapsed_ms > cfg.max_hold_ms {
343 panic!(
344 "Engine.with_state('{op}') held {elapsed_ms}ms > max {}ms — suspected R3 violation (long op inside lock)",
345 cfg.max_hold_ms
346 );
347 }
348 Ok(result)
349 }
350
351 // ═══════════════════════════════════════════════════════════════════════
352 // Token verify (= sig + expire + gate + uses_left)
353 // ═══════════════════════════════════════════════════════════════════════
354
355 /// Four steps: (1) signature verify, (2) expiry check, (3) role × verb
356 /// gate, (4) `uses_left` consume.
357 pub async fn verify_token(&self, token: &CapToken, verb: Verb) -> Result<(), EngineError> {
358 // (1) sig
359 if !self.inner.signer.verify_sig(token) {
360 return Err(EngineError::BadSignature);
361 }
362 // (2) expire
363 if token.is_expired(now_unix()) {
364 return Err(EngineError::TokenExpired);
365 }
366 // (3) role × verb gate
367 if !self.inner.gate.is_allowed(token.role, verb) {
368 return Err(EngineError::RoleViolation {
369 role: token.role,
370 verb,
371 });
372 }
373 // (4) server-side uses_left consume
374 let fp = token.fingerprint();
375 self.with_state("token.consume", move |s| {
376 let rec = s
377 .tokens
378 .get_mut(&fp)
379 .ok_or_else(|| EngineError::TokenNotFound(fp.clone()))?;
380 rec.consume()
381 .map_err(|_: crate::core::state::CapTokenConsumeError| {
382 EngineError::TokenUsesExhausted
383 })?;
384 Ok::<(), EngineError>(())
385 })
386 .await??;
387 Ok(())
388 }
389
390 /// `verify_token` plus the **task-ownership gate**.
391 ///
392 /// When a Worker-role token calls a state-touch verb (`fetch_prompt` /
393 /// `post_result` / `read_task_state` / `cancel_task` / `poll_task`),
394 /// the gate checks that `CapTokenRecord.task_id` matches the argument
395 /// `task_id`; a mismatch returns `EngineError::TokenTaskMismatch`.
396 /// Operator / Senior / Observer tokens are outside the ownership gate
397 /// and may touch any task.
398 ///
399 /// **Verbs exempt from the gate.** `start_task` and `dispatch_attempt`
400 /// stay outside so recursive swarming keeps working; depth is capped
401 /// by `max_spawn_depth`.
402 pub async fn verify_token_for_task(
403 &self,
404 token: &CapToken,
405 verb: Verb,
406 task_id: &StepId,
407 ) -> Result<(), EngineError> {
408 self.verify_token(token, verb).await?;
409 if token.role != Role::Worker {
410 return Ok(());
411 }
412 let fp = token.fingerprint();
413 let arg_tid = task_id.clone();
414 self.with_state("token.ownership_gate", move |s| {
415 let bound = s.tokens.get(&fp).and_then(|r| r.task_id.as_ref()).cloned();
416 match bound {
417 Some(t) if t == arg_tid => Ok(()),
418 Some(t) => Err(EngineError::TokenTaskMismatch {
419 bound: t.into_string(),
420 arg: arg_tid.into_string(),
421 }),
422 None => Err(EngineError::TokenNotFound(fp.clone())),
423 }
424 })
425 .await??;
426 Ok(())
427 }
428
429 /// Resolve the bound `task_id` from a Worker-role token. Used on the
430 /// simple `/v1/worker/submit` endpoint, where the worker POSTs with a
431 /// token but no `task_id`. Returns `Err` if the token role is not
432 /// Worker, or if no bound task is set.
433 pub async fn task_id_from_token(&self, token: &CapToken) -> Result<StepId, EngineError> {
434 if token.role != Role::Worker {
435 return Err(EngineError::RoleViolation {
436 role: token.role,
437 verb: Verb::PostResult,
438 });
439 }
440 let fp = token.fingerprint();
441 self.with_state("task_id_from_token", move |s| {
442 s.tokens
443 .get(&fp)
444 .and_then(|r| r.task_id.as_ref())
445 .cloned()
446 .ok_or_else(|| EngineError::TokenNotFound(fp.clone()))
447 })
448 .await?
449 }
450
451 /// Resolve a short worker handle (`wh-XXXXXXXX`) to the bound
452 /// `task_id`. Used on `/v1/worker/submit` when the Bearer is a short
453 /// handle string rather than a full `CapToken` JSON. A missing entry
454 /// returns `TokenNotFound`, i.e. "the handle is not in the store".
455 pub async fn task_id_from_handle(&self, handle: &str) -> Result<StepId, EngineError> {
456 let h = handle.to_string();
457 self.with_state("task_id_from_handle", move |s| {
458 let fp = s
459 .worker_handles
460 .get(&h)
461 .cloned()
462 .ok_or_else(|| EngineError::TokenNotFound(format!("handle={h}")))?;
463 s.tokens
464 .get(&fp)
465 .and_then(|r| r.task_id.as_ref())
466 .cloned()
467 .ok_or_else(|| EngineError::TokenNotFound(format!("fp={fp}")))
468 })
469 .await?
470 }
471
472 /// Submit a worker result via a short handle. Skips token verification
473 /// and updates `output_tail` `Final` + `task.last_result` directly in
474 /// a thin path. The caller is expected to have already resolved
475 /// `task_id` via `task_id_from_handle` — the handle's presence in
476 /// `worker_handles` means it was minted server-side and is therefore
477 /// trusted.
478 pub async fn submit_worker_result_trusted(
479 &self,
480 task_id: &StepId,
481 attempt: u32,
482 value: Value,
483 ok: bool,
484 ) -> Result<(), EngineError> {
485 let task_id_for_apply = task_id.clone();
486 let value_for_event = value.clone();
487 self.with_state("submit_worker_result_trusted.output", move |s| {
488 let ev = crate::worker::output::OutputEvent::Final {
489 content: crate::worker::output::ContentRef::Inline {
490 value: value_for_event,
491 },
492 ok,
493 };
494 s.output_store
495 .entry((task_id_for_apply.clone(), attempt))
496 .or_default()
497 .push(ev.clone());
498 s.push_event(crate::core::state::Event::WorkerOutput {
499 task_id: task_id_for_apply,
500 attempt,
501 event: ev,
502 });
503 })
504 .await?;
505 let task_id_for_result = task_id.clone();
506 let value_for_result = value.clone();
507 self.with_state("submit_worker_result_trusted.last_result", move |s| {
508 if let Some(t) = s.tasks.get_mut(&task_id_for_result) {
509 t.last_result = Some(value_for_result);
510 t.updated_at = now_unix();
511 }
512 })
513 .await?;
514 // subtask-4 / ST2 rework: this path always submits a `Final` (there
515 // is no other event kind on `/v1/worker/submit`), so the
516 // submit-time projection sink always fires — see
517 // `materialize_final_submission`'s doc and `submit_output`'s
518 // Invariants (fail-open, never turns a would-have-succeeded submit
519 // into a failure).
520 let content = crate::worker::output::ContentRef::Inline { value };
521 self.materialize_final_submission(task_id, attempt, &content, ok)
522 .await;
523 Ok(())
524 }
525
526 /// Stage a named `Artifact` from a worker via a short handle (GH #36
527 /// ST1: named multi-part worker output). Trusted analog of
528 /// [`Self::submit_worker_result_trusted`] for `OutputEvent::Artifact`:
529 /// skips token verification for the same reason (the caller already
530 /// resolved `task_id` via `task_id_from_handle`, so the handle's
531 /// presence in `worker_handles` is itself the trust boundary).
532 ///
533 /// Appends to the same per-`(task_id, attempt)` `output_store` tail
534 /// [`Self::dispatch_attempt_with`]'s Final-pull later folds into
535 /// `{"out": <final>, "parts": {<name>: <value>, ...}}` (see that
536 /// method's doc for the fold semantics — event order, last-write-wins
537 /// per name), AND records `name` in `EngineState.worker_artifact_names`
538 /// — the fold's allowlist of the WORKER's own staged parts, as opposed
539 /// to every `Artifact` that happens to land on the shared tail (e.g. an
540 /// audit sidecar finding; see that field's doc). Also dual-writes to
541 /// the Data-plane `OutputStore` the same way [`Self::submit_output`]'s
542 /// `Artifact` arm does, via [`Self::materialize_artifact_submission`]
543 /// (the artifact's own `name` is its Data-plane key, no
544 /// canonicalization — see that method's doc).
545 pub async fn stage_worker_artifact_trusted(
546 &self,
547 task_id: &StepId,
548 attempt: u32,
549 name: String,
550 value: Value,
551 ) -> Result<(), EngineError> {
552 let content = crate::worker::output::ContentRef::Inline { value };
553 let task_id_for_apply = task_id.clone();
554 let name_for_apply = name.clone();
555 let content_for_apply = content.clone();
556 self.with_state("stage_worker_artifact_trusted.output", move |s| {
557 let ev = crate::worker::output::OutputEvent::Artifact {
558 name: name_for_apply.clone(),
559 content: content_for_apply,
560 };
561 s.output_store
562 .entry((task_id_for_apply.clone(), attempt))
563 .or_default()
564 .push(ev.clone());
565 s.worker_artifact_names
566 .entry((task_id_for_apply.clone(), attempt))
567 .or_default()
568 .push(name_for_apply);
569 s.push_event(crate::core::state::Event::WorkerOutput {
570 task_id: task_id_for_apply,
571 attempt,
572 event: ev,
573 });
574 })
575 .await?;
576 self.materialize_artifact_submission(task_id, attempt, &name, &content)
577 .await;
578 Ok(())
579 }
580
581 /// GH #36 ST1: the set of `Artifact` names staged for `(task_id,
582 /// attempt)` via [`Self::stage_worker_artifact_trusted`] — see
583 /// `EngineState.worker_artifact_names`'s doc. Used by
584 /// [`Self::dispatch_attempt_with`]'s Final-pull to distinguish a
585 /// worker's own named parts from any other `Artifact` producer on the
586 /// same tail.
587 async fn worker_artifact_names_for(&self, task_id: &StepId, attempt: u32) -> Vec<String> {
588 let key = (task_id.clone(), attempt);
589 self.with_state("worker_artifact_names_for", move |s| {
590 s.worker_artifact_names
591 .get(&key)
592 .cloned()
593 .unwrap_or_default()
594 })
595 .await
596 .unwrap_or_default()
597 }
598
599 /// Mint a short handle and register it in the `worker_handles` map.
600 /// Called immediately after the worker-token mint inside
601 /// `dispatch_attempt_with`, and issues a handle bound to the same
602 /// token fingerprint. Format is `wh-<8 hex chars>` (11 chars total),
603 /// designed to remove the base64 copy-paste failure mode.
604 async fn mint_worker_handle(&self, worker_fp: String) -> Result<String, EngineError> {
605 // The handle is a sole bearer secret on the `/v1/worker/submit`
606 // short-handle path (`submit_worker_result_trusted` skips token
607 // verification), so it must be unguessable — OS RNG, not the
608 // predictable uid counter. 8 hex chars (~4B entropy) keeps the
609 // documented `wh-<8 hex>` wire shape; collision between live
610 // handles is negligible at in-process handle counts.
611 let short = crate::types::secure_hex(4);
612 let handle = format!("wh-{short}");
613 let h = handle.clone();
614 self.with_state("mint_worker_handle", move |s| {
615 s.worker_handles.insert(h, worker_fp);
616 })
617 .await?;
618 Ok(handle)
619 }
620
621 // ═══════════════════════════════════════════════════════════════════════
622 // Session API
623 // ═══════════════════════════════════════════════════════════════════════
624
625 /// Attach a new session with default `OperatorInfo` (`Automate`, no
626 /// bridges/hooks). Shorthand for `attach_with(.., OperatorInfo::default())`.
627 pub async fn attach(
628 &self,
629 operator_id: impl Into<String>,
630 role: Role,
631 ttl: Duration,
632 ) -> Result<CapToken, EngineError> {
633 self.attach_with(
634 operator_id,
635 role,
636 ttl,
637 crate::core::ctx::OperatorInfo::default(),
638 )
639 .await
640 }
641
642 // ═══════════════════════════════════════════════════════════════════════
643 // BridgeRegistry API.
644 // ═══════════════════════════════════════════════════════════════════════
645
646 /// Register a `SeniorBridge` under a name. An existing entry with the
647 /// same name is overwritten. On the persisted-session reattach path,
648 /// the caller re-registers under the same ID beforehand and the
649 /// bridge becomes effective again.
650 pub async fn register_senior_bridge(
651 &self,
652 id: impl Into<String>,
653 bridge: Arc<dyn SeniorBridge>,
654 ) {
655 self.inner
656 .senior_bridges
657 .write()
658 .await
659 .insert(id.into(), bridge);
660 }
661
662 /// Register a `SpawnHook` under a name. An existing entry with the
663 /// same name is overwritten.
664 pub async fn register_spawn_hook(&self, id: impl Into<String>, hook: Arc<dyn SpawnHook>) {
665 self.inner.spawn_hooks.write().await.insert(id.into(), hook);
666 }
667
668 /// Register an `Operator` (a spawn-body backend) under a name. An
669 /// existing entry with the same name is overwritten.
670 /// `OperatorDelegateMiddleware` looks this up via `ctx` and, when
671 /// `kind = MainAi` / `Composite`, bypasses `inner.spawn` and calls
672 /// `operator.execute` instead.
673 pub async fn register_operator(
674 &self,
675 id: impl Into<String>,
676 operator: Arc<dyn crate::operator::Operator>,
677 ) {
678 self.inner
679 .operators
680 .write()
681 .await
682 .insert(id.into(), operator);
683 }
684
685 /// Unregister a `SeniorBridge` by name (e.g. on WebSocket disconnect
686 /// or explicit teardown). A missing ID is a no-op.
687 pub async fn unregister_senior_bridge(&self, id: &str) {
688 self.inner.senior_bridges.write().await.remove(id);
689 }
690
691 /// Unregister a `SpawnHook` by name. A missing ID is a no-op.
692 pub async fn unregister_spawn_hook(&self, id: &str) {
693 self.inner.spawn_hooks.write().await.remove(id);
694 }
695
696 /// Unregister an `Operator` backend by name. A missing ID is a no-op.
697 pub async fn unregister_operator(&self, id: &str) {
698 self.inner.operators.write().await.remove(id);
699 }
700
701 /// Snapshot the list of registered `SpawnHook` IDs (for test
702 /// observation and debugging).
703 pub async fn list_spawn_hook_ids(&self) -> Vec<String> {
704 self.inner
705 .spawn_hooks
706 .read()
707 .await
708 .keys()
709 .cloned()
710 .collect()
711 }
712
713 /// Snapshot the list of registered `SeniorBridge` IDs.
714 pub async fn list_senior_bridge_ids(&self) -> Vec<String> {
715 self.inner
716 .senior_bridges
717 .read()
718 .await
719 .keys()
720 .cloned()
721 .collect()
722 }
723
724 /// Snapshot the list of registered `Operator` IDs.
725 pub async fn list_operator_ids(&self) -> Vec<String> {
726 self.inner.operators.read().await.keys().cloned().collect()
727 }
728
729 /// Attach specifying IDs directly. The caller is expected to have
730 /// pre-registered them via `register_senior_bridge` /
731 /// `register_spawn_hook` / `register_operator`. This is the canonical
732 /// path when persistence is in play.
733 ///
734 /// `kind` is the "Runtime Global" tier of the `OperatorKind` cascade
735 /// (stored verbatim on `OperatorSession.operator_kind`): `Some(_)` is
736 /// an explicit request (including `Some(OperatorKind::Automate)`) that
737 /// outranks the BP-level tiers; `None` leaves it unspecified so the
738 /// BP-level tiers / final default decide. See
739 /// `crate::core::ctx::collapse_operator_kind`.
740 #[allow(clippy::too_many_arguments)]
741 pub async fn attach_with_ids(
742 &self,
743 operator_id: impl Into<String>,
744 role: Role,
745 ttl: Duration,
746 kind: Option<OperatorKind>,
747 bridge_id: Option<String>,
748 hook_id: Option<String>,
749 operator_backend_id: Option<String>,
750 operator_kind_overrides: HashMap<String, OperatorKind>,
751 bp_agent_kinds: HashMap<String, OperatorKind>,
752 bp_global_kind: Option<OperatorKind>,
753 ) -> Result<CapToken, EngineError> {
754 let operator_id = operator_id.into();
755 let token = self
756 .inner
757 .signer
758 .session(operator_id.clone(), role, vec!["*".into()], ttl);
759 let session_id = SessionId::new();
760 let fp = token.fingerprint();
761 let now = now_unix();
762 let token_for_store = token.clone();
763
764 self.with_state("attach_with_ids", |s| {
765 s.tokens
766 .insert(fp.clone(), CapTokenRecord::from_token(token_for_store));
767 s.sessions.insert(
768 session_id.clone(),
769 OperatorSession {
770 id: session_id.clone(),
771 operator_id: operator_id.clone(),
772 role,
773 attached_at: now,
774 last_seen: now,
775 attached: true,
776 owned_task_ids: Vec::new(),
777 token_fp: fp.clone(),
778 operator_kind: kind,
779 runtime_agent_kinds: operator_kind_overrides,
780 bp_agent_kinds,
781 bp_global_kind,
782 bridge_id,
783 hook_id,
784 operator_backend_id,
785 },
786 );
787 s.push_event(Event::SessionAttached {
788 session_id: session_id.clone(),
789 role,
790 });
791 })
792 .await?;
793
794 let _ = self
795 .inner
796 .event_tx
797 .send(Event::SessionAttached { session_id, role });
798 Ok(token)
799 }
800
801 /// Build an `OperatorInfo` by looking up the session's registered IDs
802 /// on the `BridgeRegistry`, plus resolving the 4-tier `OperatorKind`
803 /// cascade for `agent_name` via `crate::core::ctx::collapse_operator_kind`.
804 /// Used when `dispatch_attempt` injects `Ctx`. An unresolved ID
805 /// (nothing registered) is silently `None` — the bridge / hook simply
806 /// does not fire and the default behaviour applies.
807 async fn resolve_operator_info(
808 &self,
809 session: &OperatorSession,
810 agent_name: &str,
811 ) -> OperatorInfo {
812 let senior_bridge = if let Some(id) = &session.bridge_id {
813 self.inner.senior_bridges.read().await.get(id).cloned()
814 } else {
815 None
816 };
817 let spawn_hook = if let Some(id) = &session.hook_id {
818 self.inner.spawn_hooks.read().await.get(id).cloned()
819 } else {
820 None
821 };
822 let operator = if let Some(id) = &session.operator_backend_id {
823 self.inner.operators.read().await.get(id).cloned()
824 } else {
825 None
826 };
827 let runtime_agent = session.runtime_agent_kinds.get(agent_name).copied();
828 // "Runtime Global" tier: `Some(_)` is always an explicit request
829 // (see the field doc on `OperatorSession.operator_kind`).
830 let runtime_global = session.operator_kind;
831 let bp_agent = session.bp_agent_kinds.get(agent_name).copied();
832 let bp_global = session.bp_global_kind;
833 let kind = crate::core::ctx::collapse_operator_kind(
834 runtime_agent,
835 runtime_global,
836 bp_agent,
837 bp_global,
838 );
839 OperatorInfo {
840 kind,
841 id: session.operator_id.clone(),
842 senior_bridge,
843 spawn_hook,
844 operator,
845 }
846 }
847
848 /// Convenience attach that takes an `OperatorInfo` (three
849 /// `Arc<dyn ...>` fields plus `kind`) **inline**.
850 ///
851 /// # Pipeline
852 ///
853 /// Each `Arc<dyn ...>` is auto-registered on the engine's registry
854 /// under a synthetic ID (`br-<hex>` / `hk-<hex>` / `ob-<hex>`), and
855 /// the session stores that synthetic ID. Subsequent `dispatch_attempt`
856 /// calls rebuild the `Arc`s from those IDs via
857 /// `resolve_operator_info`, and the three middlewares fire as usual.
858 ///
859 /// # ⚠ Non-persisted sessions only
860 ///
861 /// Because this API takes inline `Arc`s, the reattach path after
862 /// session persistence cannot rebuild them — the synthetic IDs are
863 /// not present in a freshly started process's registry. If you need
864 /// persistence, use [`Self::attach_with_ids`] with `register_*` calls
865 /// beforehand to go through **named IDs** instead.
866 ///
867 /// Handy for tests and short-lived in-process sessions. Production
868 /// WebSocket callbacks and the like should prefer `attach_with_ids`
869 /// as the canonical path.
870 pub async fn attach_with(
871 &self,
872 operator_id: impl Into<String>,
873 role: Role,
874 ttl: Duration,
875 operator_info: crate::core::ctx::OperatorInfo,
876 ) -> Result<CapToken, EngineError> {
877 let operator_id = operator_id.into();
878 // The caller always hands in a fully-formed `OperatorInfo`
879 // (including its `kind`), so it is stored as an explicit "Runtime
880 // Global" tier request (`Some(kind)`) — this path never persists
881 // BP-level tiers (both stay empty below), so `Some(kind)` resolves
882 // to the same `kind` at dispatch either way; see
883 // `OperatorSession.operator_kind` doc.
884 let kind = operator_info.kind;
885 // BridgeRegistry auto-register: when the caller hands in an
886 // `Arc<dyn>` directly, register it under a synthesised ID (the inline
887 // path aware of persistence). Callers who want to pre-register with a
888 // named ID should use `register_senior_bridge` / `register_spawn_hook`
889 // + `attach_with_ids`.
890 let bridge_id = if let Some(bridge) = operator_info.senior_bridge.clone() {
891 let id = format!("br-{}", crate::types::uid_hex(8));
892 self.inner
893 .senior_bridges
894 .write()
895 .await
896 .insert(id.clone(), bridge);
897 Some(id)
898 } else {
899 None
900 };
901 let hook_id = if let Some(hook) = operator_info.spawn_hook.clone() {
902 let id = format!("hk-{}", crate::types::uid_hex(8));
903 self.inner
904 .spawn_hooks
905 .write()
906 .await
907 .insert(id.clone(), hook);
908 Some(id)
909 } else {
910 None
911 };
912 let operator_backend_id = if let Some(operator) = operator_info.operator.clone() {
913 // `ob-` = operator-backend registry id. Renamed from `op-` in the
914 // issue #11 prefix reconciliation: `op-` used to collide with the
915 // WS operator sid shape (now unified into `S-<hex>` anyway), and a
916 // shared prefix across two unrelated registries made log filtering
917 // by prefix silently ambiguous.
918 let id = format!("ob-{}", crate::types::uid_hex(8));
919 self.inner
920 .operators
921 .write()
922 .await
923 .insert(id.clone(), operator);
924 Some(id)
925 } else {
926 None
927 };
928
929 let token = self
930 .inner
931 .signer
932 .session(operator_id.clone(), role, vec!["*".into()], ttl);
933 let session_id = SessionId::new();
934 let fp = token.fingerprint();
935 let now = now_unix();
936 let token_for_store = token.clone();
937
938 self.with_state("attach_with", |s| {
939 s.tokens
940 .insert(fp.clone(), CapTokenRecord::from_token(token_for_store));
941 s.sessions.insert(
942 session_id.clone(),
943 OperatorSession {
944 id: session_id.clone(),
945 operator_id,
946 role,
947 attached_at: now,
948 last_seen: now,
949 attached: true,
950 owned_task_ids: Vec::new(),
951 token_fp: fp.clone(),
952 operator_kind: Some(kind),
953 runtime_agent_kinds: HashMap::new(),
954 bp_agent_kinds: HashMap::new(),
955 bp_global_kind: None,
956 bridge_id,
957 hook_id,
958 operator_backend_id,
959 },
960 );
961 s.push_event(Event::SessionAttached {
962 session_id: session_id.clone(),
963 role,
964 });
965 })
966 .await?;
967
968 let _ = self
969 .inner
970 .event_tx
971 .send(Event::SessionAttached { session_id, role });
972 Ok(token)
973 }
974
975 /// Mark the session bound to `token` as detached (`attached = false`).
976 /// Tasks are left in place — a later `attach`/`attach_with_ids` call
977 /// carrying the same registered bridge/hook IDs can pick them back up.
978 pub async fn detach(&self, token: &CapToken) -> Result<(), EngineError> {
979 self.verify_token(token, Verb::DetachSession).await?;
980 let fp = token.fingerprint();
981 self.with_state("detach", move |s| {
982 let sid = s
983 .sessions
984 .iter()
985 .find(|(_, sess)| sess.token_fp == fp)
986 .map(|(id, _)| id.clone());
987 if let Some(sid) = sid {
988 if let Some(sess) = s.sessions.get_mut(&sid) {
989 sess.attached = false;
990 }
991 s.push_event(Event::SessionDetached {
992 session_id: sid.clone(),
993 });
994 let _ = sid;
995 }
996 })
997 .await?;
998 Ok(())
999 }
1000
1001 /// Refresh the session's `last_seen` timestamp and mark it `attached`.
1002 /// Called periodically by an attached client to avoid being flipped to
1003 /// detached by `start_detach_loop`.
1004 pub async fn heartbeat(&self, token: &CapToken) -> Result<(), EngineError> {
1005 self.verify_token(token, Verb::Heartbeat).await?;
1006 let now = now_unix();
1007 let fp = token.fingerprint();
1008 self.with_state("heartbeat", move |s| {
1009 if let Some(sess) = s.sessions.values_mut().find(|sess| sess.token_fp == fp) {
1010 sess.last_seen = now;
1011 sess.attached = true;
1012 }
1013 })
1014 .await?;
1015 Ok(())
1016 }
1017
1018 // ═══════════════════════════════════════════════════════════════════════
1019 // Task lifecycle
1020 // ═══════════════════════════════════════════════════════════════════════
1021
1022 /// Create a new `TaskState` from `spec` and register its initial
1023 /// prompt. When the calling token is a Worker (i.e. this is a
1024 /// recursive spawn), the new task inherits `parent.spawn_depth + 1`
1025 /// and is rejected with `SpawnDepthExceeded` once `max_spawn_depth` is
1026 /// hit; an Operator-issued call starts at depth 0.
1027 pub async fn start_task(
1028 &self,
1029 token: &CapToken,
1030 spec: TaskSpec,
1031 ) -> Result<StepId, EngineError> {
1032 self.verify_token(token, Verb::StartTask).await?;
1033 let task_id = StepId::new();
1034 let initial_directive = spec.initial_directive.clone();
1035 let task_id_clone = task_id.clone();
1036 let fp = token.fingerprint();
1037 let max_depth = self.inner.cfg.max_spawn_depth;
1038 self.with_state("start_task", move |s| {
1039 // Recursive swarm depth gate (recursion guard):
1040 // Worker tokens carry CapTokenRecord.parent_task_id. Give the
1041 // child parent's spawn_depth + 1; if it exceeds `max`, raise an
1042 // error. Operator tokens (parent_task_id=None) start at depth 0.
1043 let parent_depth_opt = s
1044 .tokens
1045 .get(&fp)
1046 .and_then(|rec| rec.task_id.as_ref())
1047 .and_then(|tid| s.tasks.get(tid))
1048 .map(|t| t.spawn_depth);
1049 let depth = match parent_depth_opt {
1050 Some(d) => {
1051 if d + 1 >= max_depth {
1052 return Err(EngineError::SpawnDepthExceeded {
1053 current: d + 1,
1054 max: max_depth,
1055 });
1056 }
1057 d + 1
1058 }
1059 None => 0,
1060 };
1061
1062 let mut task = TaskState::new(task_id_clone.clone(), spec);
1063 task.spawn_depth = depth;
1064 s.tasks.insert(task_id_clone.clone(), task);
1065 s.prompts
1066 .insert((task_id_clone.clone(), 1), initial_directive);
1067 // Link to the owner session (only Operator tokens match; Worker tokens have no session).
1068 if let Some(sess) = s.sessions.values_mut().find(|sess| sess.token_fp == fp) {
1069 sess.owned_task_ids.push(task_id_clone.clone());
1070 }
1071 s.push_event(Event::TaskCreated {
1072 task_id: task_id_clone.clone(),
1073 });
1074 Ok::<(), EngineError>(())
1075 })
1076 .await??;
1077 let _ = self.inner.event_tx.send(Event::TaskCreated {
1078 task_id: task_id.clone(),
1079 });
1080 Ok(task_id)
1081 }
1082
1083 /// Fetch a snapshot of `TaskState` for `task_id`, subject to the
1084 /// task-ownership gate (see `verify_token_for_task`).
1085 pub async fn read_task_state(
1086 &self,
1087 token: &CapToken,
1088 task_id: &StepId,
1089 ) -> Result<TaskState, EngineError> {
1090 self.verify_token_for_task(token, Verb::ReadTaskState, task_id)
1091 .await?;
1092 let task_id = task_id.clone();
1093 self.with_state("read_task_state", move |s| {
1094 s.tasks
1095 .get(&task_id)
1096 .cloned()
1097 .ok_or_else(|| EngineError::TaskNotFound(task_id.to_string()))
1098 })
1099 .await?
1100 }
1101
1102 /// Mark `task_id` as `Cancelled` and wake any caller blocked in
1103 /// `poll_task` for it.
1104 pub async fn cancel_task(&self, token: &CapToken, task_id: &StepId) -> Result<(), EngineError> {
1105 self.verify_token_for_task(token, Verb::CancelTask, task_id)
1106 .await?;
1107 let tid = task_id.clone();
1108 self.with_state("cancel_task", move |s| {
1109 let task = s
1110 .tasks
1111 .get_mut(&tid)
1112 .ok_or_else(|| EngineError::TaskNotFound(tid.to_string()))?;
1113 task.status = TaskStatus::Cancelled;
1114 task.updated_at = now_unix();
1115 s.push_event(Event::TaskCancelled {
1116 task_id: tid.clone(),
1117 });
1118 Ok::<(), EngineError>(())
1119 })
1120 .await??;
1121 self.wake_task(task_id).await?;
1122 Ok(())
1123 }
1124
1125 /// Dispatch a single attempt through the given `spawner`.
1126 ///
1127 /// The lock is only held for snapshot capture; the actual spawn and
1128 /// completion await happen outside the lock (R3 discipline).
1129 ///
1130 /// Sits on the Domain side of the Data / Domain split. The dispatch
1131 /// path itself does not touch big response bodies — those flow through
1132 /// the Data plane (`output_store` module + sink / input_inject
1133 /// `SpawnerLayer`s) around this method.
1134 ///
1135 /// The caller does the compile plus `service::linker::link` and
1136 /// carries the same stack through each dispatch. Because the spawner
1137 /// is passed per-request rather than looked up from engine-global
1138 /// state, parallel requests against a single `Engine` instance
1139 /// (different Blueprints, different spawners) do not race.
1140 ///
1141 /// `run_id`, when `Some` (issue #13 run_id propagation —
1142 /// `EngineDispatcher` threads it in from its `RunContext`), is
1143 /// inserted into `Ctx.meta.runtime["run_id"]` (a plain JSON string)
1144 /// alongside `worker_handle`, so `Operator::execute` implementations
1145 /// (e.g. `WSOperatorSession`) can read it back and surface it to the
1146 /// worker (Spawn directive / prompt). `None` (every pre-existing
1147 /// caller / test) omits the key entirely — unchanged behavior.
1148 pub async fn dispatch_attempt_with(
1149 &self,
1150 token: &CapToken,
1151 task_id: &StepId,
1152 spawner: &Arc<dyn SpawnerAdapter>,
1153 run_id: Option<&RunId>,
1154 ) -> Result<DispatchOutcome, EngineError> {
1155 self.verify_token(token, Verb::DispatchAttempt).await?;
1156 let task_id = task_id.clone();
1157
1158 // 1) Under the lock: increment the attempt number, mark Running, snapshot the
1159 // prompt, and pull `operator_info` from the session so we can inject it into Ctx.
1160 let fp = token.fingerprint();
1161 let tid_for_prep = task_id.clone();
1162 let (attempt, agent, session_snapshot, step_ctx) = self
1163 .with_state("dispatch.prep", move |s| {
1164 let task = s
1165 .tasks
1166 .get_mut(&tid_for_prep)
1167 .ok_or_else(|| EngineError::TaskNotFound(tid_for_prep.to_string()))?;
1168 task.attempt += 1;
1169 task.status = TaskStatus::Running;
1170 task.updated_at = now_unix();
1171 // The spawner pulls the prompt via engine.fetch_prompt. In prep,
1172 // if the prompts table has no entry for this attempt yet,
1173 // fall back and insert `initial_directive` so the subsequent
1174 // fetch_prompt succeeds.
1175 let attempt = task.attempt;
1176 let initial = task.spec.initial_directive.clone();
1177 s.prompts
1178 .entry((tid_for_prep.clone(), attempt))
1179 .or_insert(initial);
1180 let task = s
1181 .tasks
1182 .get(&tid_for_prep)
1183 .ok_or_else(|| EngineError::TaskNotFound(tid_for_prep.to_string()))?;
1184 let agent = task.spec.agent.clone();
1185 // GH #21 Phase 2: re-read `TaskSpec.step_ctx` on EVERY
1186 // attempt (not cached once at start_task) so retries and
1187 // Run-rekicks all carry the Step tier through to Ctx —
1188 // see TaskSpec.step_ctx's doc.
1189 let step_ctx = task.spec.step_ctx.clone();
1190 // Session snapshot (looked up by token nonce). When no session
1191 // exists (worker token invoked directly / test injection), fall
1192 // back to None → default OperatorInfo.
1193 let sess_clone = s
1194 .sessions
1195 .values()
1196 .find(|sess| sess.token_fp == fp)
1197 .cloned();
1198 Ok::<_, EngineError>((attempt, agent, sess_clone, step_ctx))
1199 })
1200 .await??;
1201 // BridgeRegistry lookup + per-agent OperatorKind cascade.
1202 let operator_info = match session_snapshot {
1203 Some(sess) => self.resolve_operator_info(&sess, &agent).await,
1204 None => OperatorInfo::default(),
1205 };
1206
1207 // 2) Outside the lock: worker token mint + spawn.
1208 //
1209 // Session-style mint (max_uses=None). Within one attempt the worker is
1210 // expected to hit `verify_token + fetch_prompt + fetch_data + post_result`
1211 // multiple times in order, so `one_time` would exhaust the token on the
1212 // very first verb. Capability is guarded by (a) the role × verb gate and
1213 // (b) the short TTL (1800s).
1214 let worker_token = self.inner.signer.session(
1215 format!("worker-of-{task_id}"),
1216 Role::Worker,
1217 vec!["*".into()],
1218 Duration::from_secs(1800),
1219 );
1220 let worker_fp = worker_token.fingerprint();
1221 let task_id_for_worker = task_id.clone();
1222 let worker_token_for_store = worker_token.clone();
1223 self.with_state("dispatch.mint_worker", move |s| {
1224 s.tokens.insert(
1225 worker_fp,
1226 CapTokenRecord::from_worker_token(worker_token_for_store, task_id_for_worker),
1227 );
1228 })
1229 .await?;
1230
1231 // Mint a short handle (`wh-XXXXXXXX`) and register it in worker_handles.
1232 // Used by the simplified Bearer path for SubAgents (short-handle form
1233 // avoids base64 copy-paste incidents).
1234 let worker_handle = self.mint_worker_handle(worker_token.fingerprint()).await?;
1235
1236 let mut ctx = Ctx::new(task_id.clone(), attempt, agent.clone());
1237 ctx.operator = operator_info; // activates MainAIMiddleware / Senior bridge
1238 ctx.meta
1239 .runtime
1240 .insert("worker_handle".to_string(), Value::String(worker_handle));
1241 if let Some(rid) = run_id {
1242 ctx.meta
1243 .runtime
1244 .insert(RUN_ID_KEY.to_string(), Value::String(rid.to_string()));
1245 }
1246 // GH #21 Phase 2: the Step tier's resolved context bundle (from
1247 // `TaskSpec.step_ctx`, re-read every attempt above) — consumed by
1248 // `AgentContextMiddleware`, which unpacks its keys ahead of the
1249 // Agent / BP-global tiers.
1250 if let Some(step_ctx) = step_ctx {
1251 ctx.meta.runtime.insert(STEP_CTX_KEY.to_string(), step_ctx);
1252 }
1253
1254 let worker = spawner
1255 .spawn(self, &ctx, task_id.clone(), attempt, worker_token)
1256 .await
1257 .map_err(|e| EngineError::DispatchFailed(e.to_string()))?;
1258
1259 // 3) Outside the lock: await worker.join() (signal-only). WorkerError is
1260 // stringified. The value is fetched via output_tail (sink path).
1261 let signal_result: Result<(), String> = worker.join().await.map_err(|e| e.to_string());
1262
1263 // Pull the last Final from output_tail and use it as the value. GH
1264 // #36 ST1 (named multi-part worker output): also fold every
1265 // `Artifact` the WORKER ITSELF staged on the same tail (via
1266 // `stage_worker_artifact_trusted` / `POST /v1/worker/artifact`)
1267 // into a `"parts"` object keyed by name — event order,
1268 // last-write-wins per name (a name staged twice overwrites,
1269 // mirroring `HashMap`/`Map` insert semantics, not an accumulating
1270 // list). `worker_artifact_names_for` is the allowlist that scopes
1271 // this to the worker's own opt-in parts — an `Artifact` some OTHER
1272 // producer appended to this same tail (e.g.
1273 // `AfterRunAuditMiddleware`'s `"audit:<step_ref>"` sidecar finding)
1274 // is left untouched (see `fold_final_and_parts`'s doc). When at
1275 // least one part was staged, the BP-chain value becomes `{"out":
1276 // <final value>, "parts": {...}}`; zero parts staged (the
1277 // pre-GH-#36 case, and every non-opt-in step) leaves the value
1278 // exactly the plain `Final` value, byte-identical to before this
1279 // change.
1280 let value_ok: Result<(Value, bool), String> = match signal_result {
1281 Ok(()) => {
1282 let tail = self.output_tail(&task_id, attempt).await;
1283 let staged_names = self.worker_artifact_names_for(&task_id, attempt).await;
1284 fold_final_and_parts(&tail, &staged_names)
1285 .ok_or_else(|| "no Final in output_tail".to_string())
1286 }
1287 Err(msg) => Err(msg),
1288 };
1289
1290 // 4) Under the lock: apply (split the borrow scope so push_event and task mut can co-exist).
1291 let outcome = self
1292 .with_state("dispatch.apply", |s| {
1293 if !s.tasks.contains_key(&task_id) {
1294 return Err(EngineError::TaskNotFound(task_id.to_string()));
1295 }
1296 match value_ok {
1297 Ok((value, ok)) => {
1298 let pass = ok;
1299 {
1300 let task = s.tasks.get_mut(&task_id).unwrap();
1301 task.last_result = Some(value.clone());
1302 task.updated_at = now_unix();
1303 task.status = if pass {
1304 TaskStatus::Pass
1305 } else {
1306 TaskStatus::Blocked
1307 };
1308 }
1309 s.push_event(Event::TaskAttemptCompleted {
1310 task_id: task_id.clone(),
1311 attempt,
1312 result: value.clone(),
1313 });
1314 if pass {
1315 s.push_event(Event::TaskPass {
1316 task_id: task_id.clone(),
1317 result: value.clone(),
1318 });
1319 Ok::<_, EngineError>(DispatchOutcome::Pass(value))
1320 } else {
1321 s.push_event(Event::TaskBlocked {
1322 task_id: task_id.clone(),
1323 result: value.clone(),
1324 });
1325 Ok(DispatchOutcome::Blocked(value))
1326 }
1327 }
1328 Err(msg) => {
1329 let task = s.tasks.get_mut(&task_id).unwrap();
1330 task.status = TaskStatus::Blocked;
1331 task.updated_at = now_unix();
1332 Err(EngineError::DispatchFailed(msg))
1333 }
1334 }
1335 })
1336 .await??;
1337
1338 // event broadcast (outside the lock — push_event feeds the in-memory tail; broadcast is a separate path).
1339 let _ = self.inner.event_tx.send(Event::TaskAttemptCompleted {
1340 task_id: task_id.clone(),
1341 attempt,
1342 result: match &outcome {
1343 DispatchOutcome::Pass(v) | DispatchOutcome::Blocked(v) => v.clone(),
1344 _ => Value::Null,
1345 },
1346 });
1347
1348 // Wake any callers waiting in poll_task.
1349 self.wake_task(&task_id).await?;
1350
1351 Ok(outcome)
1352 }
1353
1354 // ═══════════════════════════════════════════════════════════════════════
1355 // Worker-side API (= prompt / data fetch + result post)
1356 // ═══════════════════════════════════════════════════════════════════════
1357
1358 /// Fetch the directive/prompt `Value` for `task_id`'s current attempt.
1359 /// Falls back to `initial_directive` when no prompt has been recorded
1360 /// yet for that attempt. Returns the `Value` end-to-end (issue #18);
1361 /// the render down to `String` happens only at the two consumer
1362 /// boundaries — the Worker HTTP path (`fetch_worker_payload*` →
1363 /// `WorkerPayload.prompt: String`) and the WS Spawn frame text
1364 /// render (`operator_ws::session`).
1365 pub async fn fetch_prompt(
1366 &self,
1367 token: &CapToken,
1368 task_id: &StepId,
1369 ) -> Result<Value, EngineError> {
1370 self.verify_token_for_task(token, Verb::FetchPrompt, task_id)
1371 .await?;
1372 let task_id = task_id.clone();
1373 self.with_state("fetch_prompt", move |s| {
1374 let task = s
1375 .tasks
1376 .get(&task_id)
1377 .ok_or_else(|| EngineError::TaskNotFound(task_id.to_string()))?;
1378 s.prompts
1379 .get(&(task_id.clone(), task.attempt.max(1)))
1380 .cloned()
1381 .ok_or_else(|| {
1382 EngineError::ResourceNotFound(format!(
1383 "prompt({}, attempt={})",
1384 task_id, task.attempt
1385 ))
1386 })
1387 })
1388 .await?
1389 }
1390
1391 /// Combined fetch for `HTTP /v1/worker/prompt`: returns `prompt` +
1392 /// (optional) `system` + `agent` + `attempt` in a single round trip.
1393 /// The verb gate reuses `FetchPrompt` — same semantics as "the worker
1394 /// pulls its task input".
1395 ///
1396 /// `system` is the value written by `OperatorSpawner::spawn` through
1397 /// `bake_worker_system_prompt` when it ran; otherwise `None` (no
1398 /// profile present, or the bake never happened).
1399 pub async fn fetch_worker_payload(
1400 &self,
1401 token: &CapToken,
1402 task_id: &StepId,
1403 ) -> Result<crate::types::WorkerPayload, EngineError> {
1404 self.verify_token_for_task(token, Verb::FetchPrompt, task_id)
1405 .await?;
1406 let task_id_clone = task_id.clone();
1407 let mut payload = self
1408 .with_state("fetch_worker_payload", move |s| {
1409 let task = s
1410 .tasks
1411 .get(&task_id_clone)
1412 .ok_or_else(|| EngineError::TaskNotFound(task_id_clone.to_string()))?;
1413 let attempt = task.attempt.max(1);
1414 let prompt = s
1415 .prompts
1416 .get(&(task_id_clone.clone(), attempt))
1417 .cloned()
1418 .ok_or_else(|| {
1419 EngineError::ResourceNotFound(format!(
1420 "prompt({}, attempt={})",
1421 task_id_clone, attempt
1422 ))
1423 })?;
1424 let system = s
1425 .systems
1426 .get(&(task_id_clone.clone(), attempt))
1427 .cloned()
1428 .unwrap_or(None);
1429 let agent = task.spec.agent.clone();
1430 let context = s
1431 .agent_ctx
1432 .get(&(task_id_clone.clone(), attempt))
1433 .map(|e| e.view.clone());
1434 Ok::<_, EngineError>(crate::types::WorkerPayload {
1435 task_id: task_id_clone.clone(),
1436 attempt,
1437 agent,
1438 prompt: render_directive_to_string(&prompt),
1439 system,
1440 context,
1441 system_ref: None,
1442 })
1443 })
1444 .await??;
1445 self.apply_system_ref_threshold(&mut payload).await?;
1446 Ok(payload)
1447 }
1448
1449 /// Fetch a worker payload via a short handle. Skips token verification
1450 /// and returns `prompt` + `system` + `agent` + `attempt` in a thin
1451 /// path. The caller is expected to have already resolved `task_id`
1452 /// via `task_id_from_handle` — the handle's presence in
1453 /// `worker_handles` means it was minted server-side and is therefore
1454 /// trusted.
1455 pub async fn fetch_worker_payload_trusted(
1456 &self,
1457 task_id: &StepId,
1458 ) -> Result<crate::types::WorkerPayload, EngineError> {
1459 let task_id_clone = task_id.clone();
1460 let mut payload = self
1461 .with_state("fetch_worker_payload_trusted", move |s| {
1462 let task = s
1463 .tasks
1464 .get(&task_id_clone)
1465 .ok_or_else(|| EngineError::TaskNotFound(task_id_clone.to_string()))?;
1466 let attempt = task.attempt.max(1);
1467 let prompt = s
1468 .prompts
1469 .get(&(task_id_clone.clone(), attempt))
1470 .cloned()
1471 .ok_or_else(|| {
1472 EngineError::ResourceNotFound(format!(
1473 "prompt({}, attempt={})",
1474 task_id_clone, attempt
1475 ))
1476 })?;
1477 let system = s
1478 .systems
1479 .get(&(task_id_clone.clone(), attempt))
1480 .cloned()
1481 .unwrap_or(None);
1482 let agent = task.spec.agent.clone();
1483 let context = s
1484 .agent_ctx
1485 .get(&(task_id_clone.clone(), attempt))
1486 .map(|e| e.view.clone());
1487 Ok::<_, EngineError>(crate::types::WorkerPayload {
1488 task_id: task_id_clone.clone(),
1489 attempt,
1490 agent,
1491 prompt: render_directive_to_string(&prompt),
1492 system,
1493 context,
1494 system_ref: None,
1495 })
1496 })
1497 .await??;
1498 self.apply_system_ref_threshold(&mut payload).await?;
1499 Ok(payload)
1500 }
1501
1502 /// GH #31: shared threshold-branch tail for
1503 /// [`Self::fetch_worker_payload`] / [`Self::fetch_worker_payload_trusted`].
1504 /// Both build a raw `WorkerPayload` inside `with_state` with `system`
1505 /// populated as before and `system_ref: None`; this runs *outside* any
1506 /// lock (R3 — `SystemRefMode::File`'s `tokio::fs` write is a genuine
1507 /// `.await`, which `with_state`'s sync-closure contract forbids inside
1508 /// the lock) and rewrites `payload.system` / `payload.system_ref` in
1509 /// place per `SystemRefConfig.threshold_bytes`: over-threshold clears
1510 /// `system` and populates `system_ref`; at-or-under-threshold leaves
1511 /// `system` as-is and `system_ref` stays `None`. A no-op when
1512 /// `payload.system` is already `None` (no `system_prompt` was baked).
1513 async fn apply_system_ref_threshold(
1514 &self,
1515 payload: &mut crate::types::WorkerPayload,
1516 ) -> Result<(), EngineError> {
1517 let Some(rendered) = payload.system.take() else {
1518 return Ok(());
1519 };
1520 let cfg = self.cfg().system_ref.clone();
1521 if rendered.len() <= cfg.threshold_bytes {
1522 payload.system = Some(rendered);
1523 return Ok(());
1524 }
1525 use sha2::Digest;
1526 let size_bytes = rendered.len() as u64;
1527 let sha256 = hex::encode(sha2::Sha256::digest(rendered.as_bytes()));
1528 let task_id = &payload.task_id;
1529 let attempt = payload.attempt;
1530 let system_ref = match cfg.mode {
1531 crate::types::SystemRefMode::Http => crate::types::SystemRef {
1532 // The engine has no knowledge of scheme/host here — see
1533 // `SystemRefMode::Http`'s doc for who fills that in.
1534 uri: format!("/v1/worker/prompt/system?task_id={task_id}&attempt={attempt}"),
1535 sha256,
1536 size_bytes,
1537 mode: crate::types::SystemRefMode::Http,
1538 },
1539 crate::types::SystemRefMode::File => {
1540 tokio::fs::create_dir_all(&cfg.store_dir).await?;
1541 let path = cfg.store_dir.join(format!("{task_id}-{attempt}.md"));
1542 tokio::fs::write(&path, rendered.as_bytes()).await?;
1543 crate::types::SystemRef {
1544 uri: format!("file://{}", path.display()),
1545 sha256,
1546 size_bytes,
1547 mode: crate::types::SystemRefMode::File,
1548 }
1549 }
1550 };
1551 payload.system = None;
1552 payload.system_ref = Some(system_ref);
1553 Ok(())
1554 }
1555
1556 /// Returns the effective [`mlua_swarm_schema::ContextPolicy`]
1557 /// `AgentContextMiddleware` resolved and snapshotted for `(task_id,
1558 /// attempt)` at spawn time (the same policy already applied to that
1559 /// key's `EngineState.agent_ctx` entry's `.view`, GH #23 fold).
1560 /// Pass-all (`ContextPolicy::default()`) when no entry exists — either
1561 /// a pre-ST5 spawn, or a spawner stack that never layered
1562 /// `AgentContextMiddleware` (fail-open, mirroring [`Self::output_tail`]'s
1563 /// "no entry = empty default" convention).
1564 ///
1565 /// `crates/mlua-swarm-server/src/worker.rs`'s `GET /v1/worker/prompt`
1566 /// handler reads this back to filter `WorkerPayload.context.steps` via
1567 /// `ContextPolicy::allows_step`, without re-deriving the policy from
1568 /// the Blueprint at fetch time (`projection-adapter` ST5).
1569 pub async fn context_policy_for(
1570 &self,
1571 task_id: &StepId,
1572 attempt: u32,
1573 ) -> mlua_swarm_schema::ContextPolicy {
1574 let key = (task_id.clone(), attempt);
1575 self.with_state("context_policy_for", move |s| {
1576 s.agent_ctx
1577 .get(&key)
1578 .map(|e| e.policy.clone())
1579 .unwrap_or_default()
1580 })
1581 .await
1582 .unwrap_or_default()
1583 }
1584
1585 /// GH #23: returns the Blueprint-wide
1586 /// [`crate::core::step_naming::StepNaming`] table snapshotted for
1587 /// `task_id` (the same `Arc` `crate::blueprint::EngineDispatcher::dispatch`
1588 /// stashed into `EngineState.step_namings` at dispatch time —
1589 /// `Self::start_task`'s `StepId`, not the `TaskId` work item). `None`
1590 /// when no entry exists — either the dispatcher was never given a
1591 /// `StepNaming` (`EngineDispatcher::with_step_naming` not called) or
1592 /// the lock could not be acquired; callers are expected to fall back
1593 /// to the pre-GH-#23 runtime union rule in that case (subtask-2/3
1594 /// consumers).
1595 pub async fn step_naming_for(
1596 &self,
1597 task_id: &StepId,
1598 ) -> Option<Arc<crate::core::step_naming::StepNaming>> {
1599 let key = task_id.clone();
1600 self.with_state("step_naming_for", move |s| {
1601 s.step_namings.get(&key).cloned()
1602 })
1603 .await
1604 .ok()
1605 .flatten()
1606 }
1607
1608 /// GH #27 (follow-up to #23): returns the Blueprint-wide
1609 /// [`crate::core::projection_placement::ProjectionPlacement`] resolver
1610 /// snapshotted for `task_id` (the same `Arc`
1611 /// `crate::blueprint::EngineDispatcher::dispatch` stashed into
1612 /// `EngineState.projection_placements` at dispatch time — mirroring
1613 /// [`Self::step_naming_for`]'s contract exactly). `None` when no entry
1614 /// exists — either the dispatcher was never given a
1615 /// `ProjectionPlacement` (`EngineDispatcher::with_projection_placement`
1616 /// not called) or the lock could not be acquired; callers are expected
1617 /// to fall back to `ProjectionPlacement::default()` (byte-compat with
1618 /// the pre-#27 hardcoded layout) in that case.
1619 pub async fn projection_placement_for(
1620 &self,
1621 task_id: &StepId,
1622 ) -> Option<Arc<crate::core::projection_placement::ProjectionPlacement>> {
1623 let key = task_id.clone();
1624 self.with_state("projection_placement_for", move |s| {
1625 s.projection_placements.get(&key).cloned()
1626 })
1627 .await
1628 .ok()
1629 .flatten()
1630 }
1631
1632 /// Returns the [`crate::core::agent_context::AgentContextView`]
1633 /// snapshotted for `(task_id, attempt)`, if `AgentContextMiddleware`
1634 /// stashed one — the same lookup [`Self::fetch_worker_payload`] /
1635 /// [`Self::fetch_worker_payload_trusted`] perform inline, exposed
1636 /// standalone for callers that only need the view (not a full
1637 /// `WorkerPayload`) — e.g. the HTTP debug-plane `GET
1638 /// /v1/tasks/:id/runs/:run/steps*` handlers resolving a
1639 /// materialized-file root for a step *other than* the one currently
1640 /// fetching its own prompt (`projection-adapter` ST5).
1641 pub async fn agent_context_for(
1642 &self,
1643 task_id: &StepId,
1644 attempt: u32,
1645 ) -> Option<crate::core::agent_context::AgentContextView> {
1646 let key = (task_id.clone(), attempt);
1647 self.with_state("agent_context_for", move |s| {
1648 s.agent_ctx.get(&key).map(|e| e.view.clone())
1649 })
1650 .await
1651 .ok()
1652 .flatten()
1653 }
1654
1655 /// Read the current attempt number for a task (server-side lookup, no
1656 /// token verification). Used on `HTTP /v1/worker/result` when the
1657 /// worker omits `attempt` and the server has to fill it in.
1658 pub async fn task_attempt(&self, task_id: &StepId) -> Result<u32, EngineError> {
1659 let task_id = task_id.clone();
1660 self.with_state("task_attempt", move |s| {
1661 s.tasks
1662 .get(&task_id)
1663 .map(|t| t.attempt)
1664 .ok_or_else(|| EngineError::TaskNotFound(task_id.to_string()))
1665 })
1666 .await?
1667 }
1668
1669 /// Server-side admin API that lets `OperatorSpawner::spawn` bake the
1670 /// rendered `system_prompt` into engine state. There is no verb gate
1671 /// — the only expected caller is inside the spawner. SubAgents fetch
1672 /// this alongside the prompt on the `/v1/worker/prompt` path.
1673 pub async fn bake_worker_system_prompt(
1674 &self,
1675 task_id: &StepId,
1676 attempt: u32,
1677 system: Option<String>,
1678 ) -> Result<(), EngineError> {
1679 let task_id = task_id.clone();
1680 self.with_state("bake_worker_system_prompt", move |s| {
1681 // GH #31: record this agent's most-recently-baked render size
1682 // before `system` is moved into `s.systems.insert` below. Same
1683 // `s.tasks.get(&task_id)` → `.spec.agent` lookup pattern
1684 // `fetch_worker_payload` uses (see its doc for why this keying
1685 // is load-bearing for a later `bp_doctor` route).
1686 if let Some(rendered) = system.as_ref() {
1687 if let Some(agent) = s.tasks.get(&task_id).map(|t| t.spec.agent.clone()) {
1688 s.agent_render_sizes.insert(agent, rendered.len());
1689 }
1690 }
1691 s.systems.insert((task_id, attempt), system);
1692 })
1693 .await?;
1694 Ok(())
1695 }
1696
1697 /// GH #31: the most-recently-baked `system_prompt` render size (in
1698 /// bytes) observed for `agent_name`, if `bake_worker_system_prompt` has
1699 /// ever recorded one — last-write-wins across every `(task_id,
1700 /// attempt)` dispatch of that agent. `None` when no `system_prompt`
1701 /// has ever been baked for this agent name. Read by the `bp_doctor`
1702 /// route this subtask's follow-up adds.
1703 pub async fn agent_last_rendered_size(&self, agent_name: &str) -> Option<usize> {
1704 let agent_name = agent_name.to_string();
1705 self.with_state("agent_last_rendered_size", move |s| {
1706 s.agent_render_sizes.get(&agent_name).copied()
1707 })
1708 .await
1709 .ok()
1710 .flatten()
1711 }
1712
1713 /// GH #31: plain read-through of the baked `system` string for
1714 /// `(task_id, attempt)` from `EngineState.systems`, with no threshold
1715 /// branching. Backs `GET /v1/worker/prompt/system` (the `Http`-mode
1716 /// fetch target `system_ref.uri` points at) — that route needs the
1717 /// exact raw bytes to serve as the response body for the client's
1718 /// sha256 verification, not a `WorkerPayload`-wrapped value.
1719 ///
1720 /// Distinct from `apply_system_ref_threshold` (private, mutates an
1721 /// already-built `WorkerPayload` in place after full construction):
1722 /// this accessor has no threshold logic and is `pub` so
1723 /// `mlua-swarm-server`'s `worker` module can call it directly.
1724 ///
1725 /// Returns `Ok(None)` if no baked system exists for that `(task_id,
1726 /// attempt)` (either the task/attempt has no entry in `s.systems`, or
1727 /// the entry is present but stores `None`) — the caller maps this to
1728 /// a 404.
1729 pub async fn raw_system_prompt(
1730 &self,
1731 task_id: &StepId,
1732 attempt: u32,
1733 ) -> Result<Option<String>, EngineError> {
1734 let task_id = task_id.clone();
1735 self.with_state("raw_system_prompt", move |s| {
1736 s.systems.get(&(task_id, attempt)).cloned().unwrap_or(None)
1737 })
1738 .await
1739 }
1740
1741 /// Fetch an arbitrary named resource previously stored via
1742 /// `set_resource`. Not task-scoped — any valid token with the
1743 /// `FetchData` verb may read any key.
1744 pub async fn fetch_data(&self, token: &CapToken, key: &str) -> Result<Value, EngineError> {
1745 self.verify_token(token, Verb::FetchData).await?;
1746 let key = key.to_string();
1747 self.with_state("fetch_data", move |s| {
1748 s.resources
1749 .get(&key)
1750 .cloned()
1751 .ok_or(EngineError::ResourceNotFound(key))
1752 })
1753 .await?
1754 }
1755
1756 // ───────────────────────────────────────────────────────────────────────
1757 // Output path.
1758 // ───────────────────────────────────────────────────────────────────────
1759
1760 /// Send one output event from inside a `SpawnerAdapter` or worker.
1761 /// Structuring is assumed to be complete by the time we cross the
1762 /// `SpawnerAdapter` boundary; this API just appends to the
1763 /// `OutputStore`, pushes to the `EventLog`, and (for `Final`) emits
1764 /// the `TaskAttemptCompleted` event.
1765 ///
1766 /// This is Domain-side plumbing: it feeds the engine's verdict flow,
1767 /// not the Data-plane store in the `output_store` module. It also
1768 /// does not wake the dispatch path — that is done through the
1769 /// spawner's completion oneshot when the worker terminates.
1770 ///
1771 /// # Submit-time projection sink (subtask-4 / ST2 rework)
1772 ///
1773 /// A `Final` event additionally fans out to the submit-time projection
1774 /// sink ([`Self::materialize_final_submission`]): (a) when
1775 /// [`Self::set_output_store`] has wired a Data-plane
1776 /// [`crate::store::output::OutputStore`], the event is dual-written
1777 /// there (`producer_agent` = `TaskState.spec.agent`, resolved to its
1778 /// GH #23 canonical projection name — see below), and (b) when this
1779 /// task's spawn ran through `AgentContextMiddleware` (so
1780 /// `EngineState.agent_ctx` has a `.view.work_dir` / `.view.project_root`
1781 /// for it), the value is additionally materialized to the
1782 /// [`crate::core::projection_placement::ProjectionPlacement`]
1783 /// resolver's target (byte-compat default layout
1784 /// `<root>/workspace/tasks/<task_id>/ctx/<canonical_agent>.md`) — see
1785 /// `crate::core::projection`'s module doc.
1786 ///
1787 /// **GH #23 subtask-2 (canonical sink):** both writes above key off the
1788 /// canonical name — `Engine::step_naming_for(task_id)`'s
1789 /// `StepNaming::canonical_of_producer(producer_agent)` when a table was
1790 /// snapshotted for this task (`EngineDispatcher::with_step_naming`),
1791 /// else `producer_agent` unchanged (fail-open, byte-identical to
1792 /// pre-GH-#23 behavior — see [`crate::core::step_naming`]'s module
1793 /// doc).
1794 ///
1795 /// **Invariants** (Subtask 4): (1) this sink is fail-open — an
1796 /// unresolved root, an unconfigured `OutputStore`, or either one
1797 /// erroring, only logs a `tracing::warn!` and never turns this
1798 /// `Ok(())` into an `Err`; (2) the wired `OutputStore` stays the single
1799 /// source of truth for cross-step queries — the materialized file is a
1800 /// projection of it, not a second store; (3) core does not depend on
1801 /// `mlua-swarm-server` — everything this sink touches
1802 /// (`crate::store::output` / `crate::core::projection`) already lives
1803 /// in this crate.
1804 ///
1805 /// # `Artifact` dual-write (GH #34 subtask-3 gap fix)
1806 ///
1807 /// An `Artifact` event ALSO fans out to the Data-plane, via
1808 /// [`Self::materialize_artifact_submission`] — general-form: every
1809 /// `Artifact` submitted through this API dual-writes, no name-prefix
1810 /// gate. Unlike `Final`, the dual-write key is the artifact's own
1811 /// `name` field, verbatim — NOT resolved through the GH #23 canonical
1812 /// `StepNaming` table. An artifact's `name` IS its identity (mirrors
1813 /// [`crate::store::output::OutputStore::get_latest_by_name`]'s doc),
1814 /// so no canonicalization applies. Same fail-open discipline as
1815 /// `Final` (Invariant 1 above), but `Artifact` does NOT drive the
1816 /// file-materialize half (b) — artifact findings (e.g.
1817 /// `AfterRunAuditMiddleware`'s `"audit:<step_ref>"`) are observational
1818 /// sidecar data, not a step's own submission a work_dir/project_root
1819 /// projection needs to track. `Progress` / `Partial` events are
1820 /// unaffected — no behavior change.
1821 pub async fn submit_output(
1822 &self,
1823 token: &crate::types::CapToken,
1824 task_id: &StepId,
1825 attempt: u32,
1826 event: crate::worker::output::OutputEvent,
1827 ) -> Result<(), EngineError> {
1828 self.verify_token_for_task(token, crate::types::Verb::EmitOutput, task_id)
1829 .await?;
1830 let task_id_for_apply = task_id.clone();
1831 let event_clone = event.clone();
1832 self.with_state("submit_output", move |s| {
1833 s.output_store
1834 .entry((task_id_for_apply.clone(), attempt))
1835 .or_default()
1836 .push(event_clone.clone());
1837 s.push_event(crate::core::state::Event::WorkerOutput {
1838 task_id: task_id_for_apply,
1839 attempt,
1840 event: event_clone,
1841 });
1842 })
1843 .await?;
1844 match &event {
1845 crate::worker::output::OutputEvent::Final { content, ok } => {
1846 self.materialize_final_submission(task_id, attempt, content, *ok)
1847 .await;
1848 }
1849 crate::worker::output::OutputEvent::Artifact { name, content } => {
1850 self.materialize_artifact_submission(task_id, attempt, name, content)
1851 .await;
1852 }
1853 _ => {}
1854 }
1855 Ok(())
1856 }
1857
1858 /// Submit-time projection sink (subtask-4 / ST2 rework) shared by
1859 /// [`Self::submit_output`] and [`Self::submit_worker_result_trusted`].
1860 /// Best-effort / fail-open throughout (see `submit_output`'s doc
1861 /// Invariants): every failure path only `tracing::warn!`s and returns.
1862 ///
1863 /// Reads `(producer_agent, view)` via one read-only [`Self::with_state`]
1864 /// call — `producer_agent` off `TaskState.spec.agent`, `view` (the
1865 /// full [`crate::core::agent_context::AgentContextView`]) off
1866 /// `EngineState.agent_ctx[(task_id, attempt)]`, the same snapshot
1867 /// `crate::middleware::agent_context::AgentContextMiddleware` writes at
1868 /// spawn time — then does its actual (dual-write / file-write) work
1869 /// *outside* that lock, so a slow disk write or Data-plane store call
1870 /// never holds up unrelated `Engine::with_state` callers. `root` itself
1871 /// is resolved from `view` AFTER the lock via
1872 /// [`crate::core::projection_placement::ProjectionPlacement::resolve_root`]
1873 /// (GH #27, follow-up to #23) — the SAME resolver
1874 /// [`Self::step_naming_for`]'s sibling accessor
1875 /// [`Self::projection_placement_for`] snapshotted at dispatch time, so
1876 /// this sink's root-preference / fallback order is identical to the
1877 /// server read-back and the spawn-time pointer.
1878 async fn materialize_final_submission(
1879 &self,
1880 task_id: &StepId,
1881 attempt: u32,
1882 content: &crate::worker::output::ContentRef,
1883 ok: bool,
1884 ) {
1885 let task_id_for_lookup = task_id.clone();
1886 let lookup = self
1887 .with_state("materialize_final_submission.lookup", move |s| {
1888 let producer_agent = s
1889 .tasks
1890 .get(&task_id_for_lookup)
1891 .map(|t| t.spec.agent.clone());
1892 let view = s
1893 .agent_ctx
1894 .get(&(task_id_for_lookup.clone(), attempt))
1895 .map(|e| e.view.clone());
1896 (producer_agent, view)
1897 })
1898 .await;
1899 let (producer_agent, view) = match lookup {
1900 Ok(pair) => pair,
1901 Err(err) => {
1902 tracing::warn!(
1903 %task_id,
1904 error = %err,
1905 "submit-time projection sink: state lookup failed; skipping (fail-open)"
1906 );
1907 return;
1908 }
1909 };
1910 let Some(producer_agent) = producer_agent else {
1911 // Defensive only: `task_id` is always a just-looked-up task at
1912 // every real call site. No task, no addressable producer name
1913 // — nothing to project.
1914 return;
1915 };
1916 let placement = self
1917 .projection_placement_for(task_id)
1918 .await
1919 .unwrap_or_default();
1920 let root = view.and_then(|v| placement.resolve_root(&v));
1921
1922 // GH #23 subtask-2: resolve `producer_agent` to its canonical
1923 // projection name via the Blueprint-wide `StepNaming` table
1924 // snapshotted at dispatch time (`Engine::step_naming_for`). Both
1925 // write paths below ((a) data-plane, (b) file stem) use the
1926 // *canonical* name — `StepNaming::canonical_of_producer` returns
1927 // `producer_agent` unchanged for undeclared steps (byte-identical
1928 // to pre-GH-#23 behavior), and `None` (no table for this
1929 // `task_id`, e.g. a spawn that never went through
1930 // `EngineDispatcher::with_step_naming`) is a defensive fail-open
1931 // to the raw `producer_agent`, same discipline as the rest of this
1932 // sink.
1933 let canonical_agent = self
1934 .step_naming_for(task_id)
1935 .await
1936 .and_then(|naming| {
1937 naming
1938 .canonical_of_producer(&producer_agent)
1939 .map(str::to_string)
1940 })
1941 .unwrap_or_else(|| producer_agent.clone());
1942
1943 // (a) Data-plane dual-write, when an OutputStore backend is wired.
1944 if let Some(store) = self.output_store_backend() {
1945 if let Err(err) = store
1946 .append(
1947 task_id.as_str(),
1948 attempt,
1949 &canonical_agent,
1950 crate::worker::output::OutputEvent::Final {
1951 content: content.clone(),
1952 ok,
1953 },
1954 Vec::new(),
1955 )
1956 .await
1957 {
1958 tracing::warn!(
1959 %task_id,
1960 agent = %producer_agent,
1961 canonical = %canonical_agent,
1962 error = %err,
1963 "submit-time projection sink: OutputStore dual-write failed (fail-open)"
1964 );
1965 }
1966 }
1967
1968 // (b) File materialize, when a root resolved.
1969 let Some(root) = root else {
1970 tracing::warn!(
1971 %task_id,
1972 agent = %producer_agent,
1973 canonical = %canonical_agent,
1974 "submit-time projection sink: no work_dir/project_root resolved; skipping file materialize (fail-open)"
1975 );
1976 return;
1977 };
1978 let value = match content {
1979 crate::worker::output::ContentRef::Inline { value } => value.clone(),
1980 crate::worker::output::ContentRef::FileRef {
1981 path,
1982 mime,
1983 size_hint,
1984 } => serde_json::json!({
1985 "file_ref": path.to_string_lossy(),
1986 "mime": mime,
1987 "size_hint": size_hint,
1988 }),
1989 };
1990 let key = crate::core::projection::ProjectionKey {
1991 task_id: task_id.to_string(),
1992 run_id: None,
1993 step: Some(canonical_agent.clone()),
1994 path: None,
1995 };
1996 let adapter = crate::core::projection::FileProjectionAdapter::with_placement(
1997 root,
1998 (*placement).clone(),
1999 );
2000 if let Err(err) = adapter.materialize_submission(&key, &value, attempt, ok) {
2001 tracing::warn!(
2002 %task_id,
2003 agent = %producer_agent,
2004 canonical = %canonical_agent,
2005 error = %err,
2006 "submit-time projection sink: file materialize failed (fail-open)"
2007 );
2008 }
2009 }
2010
2011 /// Submit-time projection sink for `OutputEvent::Artifact` (GH #34
2012 /// subtask-3 gap fix). Data-plane-only sibling of
2013 /// [`Self::materialize_final_submission`]: when [`Self::set_output_store`]
2014 /// has wired a Data-plane [`crate::store::output::OutputStore`], the
2015 /// artifact dual-writes there under its own `name`, verbatim — general
2016 /// form, every `Artifact` submitted via [`Self::submit_output`]
2017 /// materializes this way, no name-prefix gate (see `submit_output`'s
2018 /// doc, "`Artifact` dual-write" section, for the full rationale).
2019 ///
2020 /// Deliberately does NOT resolve `producer_agent` / `StepNaming` /
2021 /// `AgentContextView` / a `root` the way `materialize_final_submission`
2022 /// does — an artifact's `name` already IS the Data-plane key (no
2023 /// canonicalization needed), and this sink does not drive the
2024 /// file-materialize half, so none of that lookup is needed here. Same
2025 /// fail-open discipline: an unconfigured `OutputStore`, or the write
2026 /// erroring, only logs a `tracing::warn!` and never surfaces to the
2027 /// caller (`submit_output` already committed the domain-plane append
2028 /// before calling this sink).
2029 async fn materialize_artifact_submission(
2030 &self,
2031 task_id: &StepId,
2032 attempt: u32,
2033 name: &str,
2034 content: &crate::worker::output::ContentRef,
2035 ) {
2036 let Some(store) = self.output_store_backend() else {
2037 return;
2038 };
2039 if let Err(err) = store
2040 .append(
2041 task_id.as_str(),
2042 attempt,
2043 name,
2044 crate::worker::output::OutputEvent::Artifact {
2045 name: name.to_string(),
2046 content: content.clone(),
2047 },
2048 Vec::new(),
2049 )
2050 .await
2051 {
2052 tracing::warn!(
2053 %task_id,
2054 artifact = %name,
2055 error = %err,
2056 "submit-time projection sink: OutputStore dual-write failed for Artifact (fail-open)"
2057 );
2058 }
2059 }
2060
2061 /// Snapshot the entire output tail for a given `(task_id, attempt)`.
2062 /// Used by the dispatch path when pulling `Final`, and by observers
2063 /// reading the trace.
2064 pub async fn output_tail(
2065 &self,
2066 task_id: &StepId,
2067 attempt: u32,
2068 ) -> Vec<crate::worker::output::OutputEvent> {
2069 let key = (task_id.clone(), attempt);
2070 self.with_state("output_tail", move |s| {
2071 s.output_store.get(&key).cloned().unwrap_or_default()
2072 })
2073 .await
2074 .unwrap_or_default()
2075 }
2076
2077 /// Record an interim `last_result` for `task_id` without changing its
2078 /// `status`. Distinct from the terminal `Final` output event handled
2079 /// through `submit_output` / `dispatch_attempt_with`.
2080 pub async fn post_result(
2081 &self,
2082 token: &CapToken,
2083 task_id: &StepId,
2084 result: Value,
2085 ) -> Result<(), EngineError> {
2086 self.verify_token_for_task(token, Verb::PostResult, task_id)
2087 .await?;
2088 let task_id = task_id.clone();
2089 let result_clone = result.clone();
2090 self.with_state("post_result", move |s| {
2091 let task = s
2092 .tasks
2093 .get_mut(&task_id)
2094 .ok_or_else(|| EngineError::TaskNotFound(task_id.to_string()))?;
2095 task.last_result = Some(result_clone);
2096 task.updated_at = now_unix();
2097 Ok::<(), EngineError>(())
2098 })
2099 .await??;
2100 Ok(())
2101 }
2102
2103 /// Store a named resource value, retrievable later via `fetch_data`.
2104 /// No token is required — this is a server-side/admin-style setter
2105 /// (mirrors `bake_worker_system_prompt`).
2106 pub async fn set_resource(
2107 &self,
2108 key: impl Into<String>,
2109 value: Value,
2110 ) -> Result<(), EngineError> {
2111 let key = key.into();
2112 self.with_state("set_resource", move |s| {
2113 s.resources.insert(key, value);
2114 })
2115 .await?;
2116 Ok(())
2117 }
2118
2119 // ═══════════════════════════════════════════════════════════════════════
2120 // Senior suspend / resume
2121 // ═══════════════════════════════════════════════════════════════════════
2122
2123 /// Ask a question of the Senior, mark the task `Suspended`, and
2124 /// return a `ResumeKey`. The suspended state persists until another
2125 /// task calls `resume(key, answer)`.
2126 ///
2127 /// Resume-side waiting is `Notify`-based, so a caller (typically
2128 /// MainAI) can detach, reattach from a different process, and still
2129 /// pull the answer out via `await_resume(key, timeout)` — the answer
2130 /// is stored inside `EngineState`.
2131 pub async fn query_senior(
2132 &self,
2133 token: &CapToken,
2134 task_id: &StepId,
2135 question: Value,
2136 ) -> Result<ResumeKey, EngineError> {
2137 self.verify_token(token, Verb::QuerySenior).await?;
2138 let task_id = task_id.clone();
2139 let key = ResumeKey::for_senior(&task_id);
2140 let task_notify = self
2141 .with_state("query_senior.notify_ensure", |s| {
2142 s.ensure_task_notify(&task_id)
2143 })
2144 .await?;
2145
2146 let key_clone = key.clone();
2147 let task_id_inner = task_id.clone();
2148 let question_clone = question.clone();
2149 self.with_state("query_senior.suspend", move |s| {
2150 let task = s
2151 .tasks
2152 .get_mut(&task_id_inner)
2153 .ok_or_else(|| EngineError::TaskNotFound(task_id_inner.to_string()))?;
2154 task.status = TaskStatus::Suspended;
2155 task.suspended_on = Some(key_clone.clone());
2156 task.updated_at = now_unix();
2157 s.pending_resumes
2158 .insert(key_clone.clone(), ResumePending::new());
2159 s.push_event(Event::SeniorQueried {
2160 task_id: task_id_inner.clone(),
2161 question: question_clone.clone(),
2162 });
2163 s.push_event(Event::TaskSuspended {
2164 task_id: task_id_inner.clone(),
2165 key: key_clone.clone(),
2166 });
2167 Ok::<(), EngineError>(())
2168 })
2169 .await??;
2170
2171 // Notify callers waiting for a task status change (Running → Suspended).
2172 task_notify.notify_waiters();
2173
2174 let _ = self
2175 .inner
2176 .event_tx
2177 .send(Event::SeniorQueried { task_id, question });
2178 Ok(key)
2179 }
2180
2181 /// Store the answer for a `ResumeKey` in `EngineState` and wake the
2182 /// waiting caller via `Notify`. Also flips the suspended task's
2183 /// status back to `Running` and fires the per-task notifier.
2184 pub async fn resume(&self, key: ResumeKey, answer: Value) -> Result<(), EngineError> {
2185 let answer_for_state = answer.clone();
2186 let answer_for_event = answer.clone();
2187 let key_clone = key.clone();
2188 let (notify, task_notify, task_id_opt) = self
2189 .with_state("resume.set", move |s| {
2190 let pending = s
2191 .pending_resumes
2192 .get_mut(&key_clone)
2193 .ok_or(EngineError::ResumeKeyNotFound)?;
2194 pending.answer = Some(answer_for_state);
2195 let notify = pending.notify.clone();
2196
2197 let task_id = s
2198 .tasks
2199 .iter()
2200 .find(|(_, t)| t.suspended_on.as_ref() == Some(&key_clone))
2201 .map(|(id, _)| id.clone());
2202
2203 let task_notify = task_id.as_ref().map(|tid| s.ensure_task_notify(tid));
2204
2205 if let Some(tid) = &task_id {
2206 if let Some(task) = s.tasks.get_mut(tid) {
2207 task.suspended_on = None;
2208 task.status = TaskStatus::Running;
2209 task.updated_at = now_unix();
2210 }
2211 s.push_event(Event::TaskResumed {
2212 task_id: tid.clone(),
2213 key: key_clone.clone(),
2214 });
2215 s.push_event(Event::SeniorAnswered {
2216 task_id: tid.clone(),
2217 answer: answer_for_event.clone(),
2218 });
2219 }
2220 Ok::<_, EngineError>((notify, task_notify, task_id))
2221 })
2222 .await??;
2223
2224 // Outside the lock: notify_waiters for both the ResumePending and task-status waits.
2225 notify.notify_waiters();
2226 if let Some(n) = task_notify {
2227 n.notify_waiters();
2228 }
2229
2230 if let Some(tid) = task_id_opt {
2231 let _ = self
2232 .inner
2233 .event_tx
2234 .send(Event::TaskResumed { task_id: tid, key });
2235 }
2236 Ok(())
2237 }
2238
2239 /// Wait for the resume answer. Even if the caller (an Operator)
2240 /// detached and reattached, the answer is available immediately here
2241 /// — if it was already stored, this returns without waiting on the
2242 /// notifier.
2243 ///
2244 /// `timeout = Duration::ZERO` performs an instant check without
2245 /// waiting.
2246 pub async fn await_resume(
2247 &self,
2248 key: ResumeKey,
2249 timeout: Duration,
2250 ) -> Result<Value, EngineError> {
2251 // (1) Under the lock: clone the notify handle and check for an existing answer.
2252 let key_clone = key.clone();
2253 let (notify, existing) = self
2254 .with_state("await_resume.snapshot", move |s| {
2255 let pending = s
2256 .pending_resumes
2257 .get(&key_clone)
2258 .ok_or(EngineError::ResumeKeyNotFound)?;
2259 Ok::<_, EngineError>((pending.notify.clone(), pending.answer.clone()))
2260 })
2261 .await??;
2262
2263 // (2) If an answer has already been stored, return immediately (detach / reattach pattern).
2264 if let Some(v) = existing {
2265 return Ok(v);
2266 }
2267
2268 // (3) Outside the lock: wait on the notify with a timeout.
2269 if timeout.is_zero() {
2270 return Err(EngineError::PollTimeout);
2271 }
2272 let waited = tokio::time::timeout(timeout, notify.notified()).await;
2273 if waited.is_err() {
2274 return Err(EngineError::PollTimeout);
2275 }
2276
2277 // (4) Under the lock: re-read the answer (should be present now that we were notified).
2278 let key_clone = key.clone();
2279 self.with_state("await_resume.read", move |s| {
2280 let pending = s
2281 .pending_resumes
2282 .get(&key_clone)
2283 .ok_or(EngineError::ResumeKeyNotFound)?;
2284 pending
2285 .answer
2286 .clone()
2287 .ok_or_else(|| EngineError::Internal("notified but answer missing".into()))
2288 })
2289 .await?
2290 }
2291
2292 // ═══════════════════════════════════════════════════════════════════════
2293 // poll_task — the "wait" path that waits for task-status changes (works for long-poll and regular wait).
2294 // ═══════════════════════════════════════════════════════════════════════
2295
2296 /// Wait until the task's status **transitions to terminal or
2297 /// `Suspended`**, then return the latest `TaskState`. Returns
2298 /// immediately if the task is already in a terminal state.
2299 /// Exceeding the timeout returns `EngineError::PollTimeout`.
2300 ///
2301 /// A `hold` of `Duration::from_secs(0)` returns a snapshot immediately
2302 /// (no wait). Larger holds — tens of minutes up to days — are fine;
2303 /// the wait state is kept in memory inside the engine and does not
2304 /// degrade.
2305 pub async fn poll_task(
2306 &self,
2307 token: &CapToken,
2308 task_id: &StepId,
2309 hold: Duration,
2310 ) -> Result<TaskState, EngineError> {
2311 self.verify_token_for_task(token, Verb::PollTask, task_id)
2312 .await?;
2313 let task_id_inner = task_id.clone();
2314
2315 // (1) Under the lock: take a snapshot and clone task_notify.
2316 let (state, notify) = self
2317 .with_state("poll_task.snapshot", move |s| {
2318 let task = s
2319 .tasks
2320 .get(&task_id_inner)
2321 .cloned()
2322 .ok_or_else(|| EngineError::TaskNotFound(task_id_inner.to_string()))?;
2323 let notify = s.ensure_task_notify(&task_id_inner);
2324 Ok::<_, EngineError>((task, notify))
2325 })
2326 .await??;
2327
2328 // (2) Immediate-return condition: already terminal / Suspended (nothing left to wait on).
2329 if matches!(
2330 state.status,
2331 TaskStatus::Pass | TaskStatus::Blocked | TaskStatus::Cancelled | TaskStatus::Suspended
2332 ) {
2333 return Ok(state);
2334 }
2335 if hold.is_zero() {
2336 return Ok(state);
2337 }
2338
2339 // (3) Outside the lock: wait on Notify with a timeout.
2340 let waited = tokio::time::timeout(hold, notify.notified()).await;
2341 if waited.is_err() {
2342 return Err(EngineError::PollTimeout);
2343 }
2344
2345 // (4) Under the lock: take a fresh snapshot.
2346 let task_id_inner = task_id.clone();
2347 self.with_state("poll_task.reread", move |s| {
2348 s.tasks
2349 .get(&task_id_inner)
2350 .cloned()
2351 .ok_or_else(|| EngineError::TaskNotFound(task_id_inner.to_string()))
2352 })
2353 .await?
2354 }
2355
2356 // ═══════════════════════════════════════════════════════════════════════
2357 // Background: heartbeat miss → detach loop
2358 // ═══════════════════════════════════════════════════════════════════════
2359
2360 /// Background loop that scans sessions every `heartbeat_interval` and
2361 /// flips `attached = false` on any session whose `last_seen` exceeds
2362 /// `heartbeat_miss_threshold * interval`.
2363 ///
2364 /// The tasks themselves are kept (assuming
2365 /// `keepalive_on_idle = true`), so another client can reattach with
2366 /// the same token and resume immediately. Dropping the returned
2367 /// `JoinHandle` does not stop the loop — the handle exists so callers
2368 /// who want to abort can hold onto it.
2369 pub fn start_detach_loop(&self) -> tokio::task::JoinHandle<()> {
2370 let engine = self.clone();
2371 let cfg = self.inner.cfg.long_hold.clone();
2372 let interval = cfg.heartbeat_interval;
2373 let miss_secs = cfg.heartbeat_interval.as_secs() * cfg.heartbeat_miss_threshold as u64;
2374
2375 tokio::spawn(async move {
2376 let mut ticker = tokio::time::interval(interval);
2377 ticker.tick().await; // first tick is immediate
2378 loop {
2379 ticker.tick().await;
2380 let now = now_unix();
2381 let detached = engine
2382 .with_state("detach_loop.scan", |s| {
2383 let mut detached = Vec::new();
2384 for (sid, sess) in s.sessions.iter_mut() {
2385 if !sess.attached {
2386 continue;
2387 }
2388 if now.saturating_sub(sess.last_seen) >= miss_secs {
2389 sess.attached = false;
2390 detached.push(sid.clone());
2391 }
2392 }
2393 for sid in &detached {
2394 s.push_event(Event::SessionDetached {
2395 session_id: sid.clone(),
2396 });
2397 }
2398 detached
2399 })
2400 .await
2401 .unwrap_or_default();
2402 for sid in detached {
2403 let _ = engine
2404 .inner
2405 .event_tx
2406 .send(Event::SessionDetached { session_id: sid });
2407 }
2408 }
2409 })
2410 }
2411
2412 /// Helper: wake a task whose status has changed. Called from the
2413 /// method body outside the lock.
2414 async fn wake_task(&self, task_id: &StepId) -> Result<(), EngineError> {
2415 let task_id = task_id.clone();
2416 let notify_opt = self
2417 .with_state("wake_task.get_notify", move |s| {
2418 s.task_notifies.get(&task_id).cloned()
2419 })
2420 .await?;
2421 if let Some(n) = notify_opt {
2422 n.notify_waiters();
2423 }
2424 Ok(())
2425 }
2426}
2427
2428// ─── UT: issue #14 — token store keyed by fingerprint, not nonce ────────────
2429#[cfg(test)]
2430mod token_fingerprint_store_tests {
2431 use super::*;
2432
2433 /// A token that was never attached fails verify with a `TokenNotFound`
2434 /// that carries the fingerprint — never the nonce. The error string can
2435 /// surface in HTTP error bodies, so this is the secret-hygiene contract.
2436 #[tokio::test]
2437 async fn verify_unknown_token_reports_fingerprint_not_nonce() {
2438 let engine = Engine::new(EngineCfg::default());
2439 // Signed by the engine's own signer (sig passes) but never inserted
2440 // into the store — verify must fail at step (4), the store lookup.
2441 let token = engine.signer().session(
2442 "ghost",
2443 Role::Operator,
2444 vec!["*".into()],
2445 Duration::from_secs(60),
2446 );
2447 let err = engine
2448 .verify_token(&token, Verb::ReadTaskState)
2449 .await
2450 .expect_err("token is not in the store");
2451 let msg = err.to_string();
2452 assert!(
2453 msg.contains(&token.fingerprint()),
2454 "error must carry the fingerprint: {msg}"
2455 );
2456 assert!(
2457 !msg.contains(&token.nonce),
2458 "error must not leak the nonce: {msg}"
2459 );
2460 }
2461
2462 /// attach → verify → heartbeat → detach all resolve the session /
2463 /// token record through fingerprint keys (mint/verify lifecycle
2464 /// regression guard for the issue #14 key migration).
2465 #[tokio::test]
2466 async fn attach_verify_heartbeat_detach_cycle_with_fp_keying() {
2467 let engine = Engine::new(EngineCfg::default());
2468 let token = engine
2469 .attach("op-1", Role::Operator, Duration::from_secs(60))
2470 .await
2471 .expect("attach");
2472 engine
2473 .verify_token(&token, Verb::ReadTaskState)
2474 .await
2475 .expect("verify consumes via fp key");
2476 engine
2477 .heartbeat(&token)
2478 .await
2479 .expect("heartbeat finds the session by fp");
2480 engine
2481 .detach(&token)
2482 .await
2483 .expect("detach finds the session by fp");
2484 }
2485}
2486
2487// ─── UT: `OperatorKind` "Runtime Global" tier — `Option` semantics ─────────
2488//
2489// Regression coverage for the "explicit Automate is indistinguishable from
2490// unspecified" defect: `OperatorSession.operator_kind` (and the
2491// `attach_with_ids` `kind` parameter it stores) is `Option<OperatorKind>`,
2492// so `Some(Automate)` is an explicit Runtime Global request that must
2493// outrank `bp_global`, while `None` must let `bp_global` decide. Exercises
2494// the real `resolve_operator_info` cascade path (not just
2495// `collapse_operator_kind` in isolation), attaching via `attach_with_ids`
2496// exactly as `TaskLaunchService::launch` does.
2497#[cfg(test)]
2498mod resolve_operator_info_runtime_global_tests {
2499 use super::*;
2500
2501 async fn attach_and_resolve(
2502 runtime_global: Option<OperatorKind>,
2503 bp_global: Option<OperatorKind>,
2504 ) -> OperatorInfo {
2505 let engine = Engine::new(EngineCfg::default());
2506 let token = engine
2507 .attach_with_ids(
2508 "ut-op",
2509 Role::Operator,
2510 Duration::from_secs(30),
2511 runtime_global,
2512 None,
2513 None,
2514 None,
2515 HashMap::new(),
2516 HashMap::new(),
2517 bp_global,
2518 )
2519 .await
2520 .expect("attach_with_ids ok");
2521 let session = engine
2522 .with_state("test.find_session", |s| {
2523 s.sessions
2524 .values()
2525 .find(|sess| sess.token_fp == token.fingerprint())
2526 .cloned()
2527 })
2528 .await
2529 .expect("with_state ok")
2530 .expect("session present after attach_with_ids");
2531 engine.resolve_operator_info(&session, "agent-x").await
2532 }
2533
2534 #[tokio::test]
2535 async fn explicit_some_automate_outranks_bp_global_main_ai() {
2536 // Runtime Global explicitly requests Automate; bp_global is MainAi.
2537 // The explicit `Some(Automate)` must win — this is exactly the case
2538 // the old `== OperatorKind::default()` convention got wrong (it
2539 // could not tell "explicitly Automate" from "unspecified" and would
2540 // have let `bp_global` (MainAi) take over instead).
2541 let info =
2542 attach_and_resolve(Some(OperatorKind::Automate), Some(OperatorKind::MainAi)).await;
2543 assert_eq!(
2544 info.kind,
2545 OperatorKind::Automate,
2546 "explicit Some(Automate) runtime_global must outrank bp_global MainAi"
2547 );
2548 }
2549
2550 #[tokio::test]
2551 async fn none_lets_bp_global_main_ai_win() {
2552 // Runtime Global left unspecified (`None`); bp_global is MainAi.
2553 // With nothing more specific set, `bp_global` must decide.
2554 let info = attach_and_resolve(None, Some(OperatorKind::MainAi)).await;
2555 assert_eq!(
2556 info.kind,
2557 OperatorKind::MainAi,
2558 "None runtime_global must let bp_global MainAi win"
2559 );
2560 }
2561}
2562
2563/// issue #13 run_id propagation: `dispatch_attempt_with`'s `run_id` param
2564/// must land in `Ctx.meta.runtime["run_id"]` (the same slot pattern as the
2565/// pre-existing `worker_handle`), or be omitted entirely when `None`. Same
2566/// `CtxProbe` shape as `middleware::worker_binding`'s test module — an
2567/// inner `SpawnerAdapter` that snapshots the `Ctx` it was called with and
2568/// fails the spawn (only the ctx snapshot matters here).
2569#[cfg(test)]
2570mod dispatch_attempt_with_run_id_tests {
2571 use super::*;
2572 use crate::worker::adapter::{SpawnError, SpawnerAdapter};
2573 use crate::worker::Worker;
2574 use std::sync::Mutex as StdMutex;
2575
2576 struct CtxProbe {
2577 seen: Arc<StdMutex<Option<Ctx>>>,
2578 }
2579
2580 #[async_trait::async_trait]
2581 impl SpawnerAdapter for CtxProbe {
2582 async fn spawn(
2583 &self,
2584 _engine: &Engine,
2585 ctx: &Ctx,
2586 _task_id: StepId,
2587 _attempt: u32,
2588 _token: CapToken,
2589 ) -> Result<Box<dyn Worker>, SpawnError> {
2590 *self.seen.lock().unwrap() = Some(ctx.clone());
2591 Err(SpawnError::Internal("probe stop".into()))
2592 }
2593 }
2594
2595 async fn dispatch_with_probe(run_id: Option<&RunId>) -> Ctx {
2596 let engine = Engine::new(EngineCfg::default());
2597 let token = engine
2598 .attach("ut-op", Role::Operator, Duration::from_secs(30))
2599 .await
2600 .expect("attach");
2601 let tid = engine
2602 .start_task(
2603 &token,
2604 TaskSpec {
2605 agent: "probe".into(),
2606 initial_directive: "hi".into(),
2607 step_ctx: None,
2608 },
2609 )
2610 .await
2611 .expect("start_task");
2612 let seen: Arc<StdMutex<Option<Ctx>>> = Arc::new(StdMutex::new(None));
2613 let spawner: Arc<dyn SpawnerAdapter> = Arc::new(CtxProbe { seen: seen.clone() });
2614 // The probe always errors the spawn (`SpawnError::Internal`); we
2615 // only care about the `Ctx` snapshot it captured, so the dispatch
2616 // outcome itself (`Err`) is discarded.
2617 let _ = engine
2618 .dispatch_attempt_with(&token, &tid, &spawner, run_id)
2619 .await;
2620 let captured = seen.lock().unwrap().clone();
2621 captured.expect("inner ctx captured")
2622 }
2623
2624 #[tokio::test]
2625 async fn run_id_lands_in_ctx_meta_runtime_when_some() {
2626 let run_id = RunId::new();
2627 let observed = dispatch_with_probe(Some(&run_id)).await;
2628 assert_eq!(
2629 observed.meta.runtime.get("run_id").and_then(|v| v.as_str()),
2630 Some(run_id.as_str()),
2631 "ctx.meta.runtime[\"run_id\"] must carry the run_id passed to dispatch_attempt_with"
2632 );
2633 }
2634
2635 #[tokio::test]
2636 async fn run_id_key_absent_when_none() {
2637 let observed = dispatch_with_probe(None).await;
2638 assert!(
2639 !observed.meta.runtime.contains_key("run_id"),
2640 "no run_id key must be injected when dispatch_attempt_with is called with None"
2641 );
2642 }
2643}
2644
2645/// GH #21 Phase 2: `TaskSpec.step_ctx` must land in
2646/// `Ctx.meta.runtime[STEP_CTX_KEY]` — re-read from the spec on EVERY
2647/// attempt (the prep closure re-reads `task.spec.step_ctx` every call, not
2648/// caching it once at `start_task`), so a retry (attempt 2) carries it
2649/// too. Same `CtxProbe` shape as `dispatch_attempt_with_run_id_tests`.
2650#[cfg(test)]
2651mod dispatch_attempt_with_step_ctx_tests {
2652 use super::*;
2653 use crate::worker::adapter::{SpawnError, SpawnerAdapter};
2654 use crate::worker::Worker;
2655 use std::sync::Mutex as StdMutex;
2656
2657 struct CtxProbe {
2658 seen: Arc<StdMutex<Option<Ctx>>>,
2659 }
2660
2661 #[async_trait::async_trait]
2662 impl SpawnerAdapter for CtxProbe {
2663 async fn spawn(
2664 &self,
2665 _engine: &Engine,
2666 ctx: &Ctx,
2667 _task_id: StepId,
2668 _attempt: u32,
2669 _token: CapToken,
2670 ) -> Result<Box<dyn Worker>, SpawnError> {
2671 *self.seen.lock().unwrap() = Some(ctx.clone());
2672 Err(SpawnError::Internal("probe stop".into()))
2673 }
2674 }
2675
2676 #[tokio::test]
2677 async fn step_ctx_lands_in_ctx_meta_runtime_on_attempt_1_and_2() {
2678 let engine = Engine::new(EngineCfg::default());
2679 let token = engine
2680 .attach("ut-op", Role::Operator, Duration::from_secs(30))
2681 .await
2682 .expect("attach");
2683 let tid = engine
2684 .start_task(
2685 &token,
2686 TaskSpec {
2687 agent: "probe".into(),
2688 initial_directive: "hi".into(),
2689 step_ctx: Some(serde_json::json!({ "work_dir": "/step" })),
2690 },
2691 )
2692 .await
2693 .expect("start_task");
2694 let seen: Arc<StdMutex<Option<Ctx>>> = Arc::new(StdMutex::new(None));
2695 let spawner: Arc<dyn SpawnerAdapter> = Arc::new(CtxProbe { seen: seen.clone() });
2696
2697 // The probe always errors the spawn; only the ctx snapshot matters.
2698 let _ = engine
2699 .dispatch_attempt_with(&token, &tid, &spawner, None)
2700 .await;
2701 let first = seen
2702 .lock()
2703 .unwrap()
2704 .clone()
2705 .expect("attempt 1 ctx captured");
2706 assert_eq!(
2707 first.meta.runtime.get(STEP_CTX_KEY),
2708 Some(&serde_json::json!({ "work_dir": "/step" })),
2709 "attempt 1 must carry TaskSpec.step_ctx in ctx.meta.runtime[STEP_CTX_KEY]"
2710 );
2711
2712 let _ = engine
2713 .dispatch_attempt_with(&token, &tid, &spawner, None)
2714 .await;
2715 let second = seen
2716 .lock()
2717 .unwrap()
2718 .clone()
2719 .expect("attempt 2 ctx captured");
2720 assert_eq!(
2721 second.meta.runtime.get(STEP_CTX_KEY),
2722 Some(&serde_json::json!({ "work_dir": "/step" })),
2723 "attempt 2 (retry) must ALSO carry TaskSpec.step_ctx — prep re-reads the spec every attempt"
2724 );
2725 }
2726
2727 #[tokio::test]
2728 async fn step_ctx_key_absent_when_none() {
2729 let engine = Engine::new(EngineCfg::default());
2730 let token = engine
2731 .attach("ut-op", Role::Operator, Duration::from_secs(30))
2732 .await
2733 .expect("attach");
2734 let tid = engine
2735 .start_task(
2736 &token,
2737 TaskSpec {
2738 agent: "probe".into(),
2739 initial_directive: "hi".into(),
2740 step_ctx: None,
2741 },
2742 )
2743 .await
2744 .expect("start_task");
2745 let seen: Arc<StdMutex<Option<Ctx>>> = Arc::new(StdMutex::new(None));
2746 let spawner: Arc<dyn SpawnerAdapter> = Arc::new(CtxProbe { seen: seen.clone() });
2747 let _ = engine
2748 .dispatch_attempt_with(&token, &tid, &spawner, None)
2749 .await;
2750 let observed = seen.lock().unwrap().clone().expect("ctx captured");
2751 assert!(
2752 !observed.meta.runtime.contains_key(STEP_CTX_KEY),
2753 "no step_ctx key must be injected when TaskSpec.step_ctx is None"
2754 );
2755 }
2756}
2757
2758// ─── issue #18: `TaskSpec.initial_directive` `Value` pass-through ──────────
2759#[cfg(test)]
2760mod initial_directive_value_passthrough_tests {
2761 use super::*;
2762
2763 async fn seeded_engine(initial_directive: Value) -> (Engine, CapToken, StepId) {
2764 let engine = Engine::new(EngineCfg::default());
2765 let op_token = engine
2766 .attach("ut-op", Role::Operator, Duration::from_secs(30))
2767 .await
2768 .expect("attach");
2769 let task_id = engine
2770 .start_task(
2771 &op_token,
2772 TaskSpec {
2773 agent: "planner".to_string(),
2774 initial_directive,
2775 step_ctx: None,
2776 },
2777 )
2778 .await
2779 .expect("start_task");
2780 (engine, op_token, task_id)
2781 }
2782
2783 /// Mint + register a `Role::Worker` token the same way
2784 /// `dispatch_attempt_with` does — `fetch_prompt` is worker-verb-gated.
2785 async fn mint_worker_token(engine: &Engine, task_id: &StepId) -> CapToken {
2786 let worker_token = engine.signer().session(
2787 format!("worker-of-{task_id}"),
2788 Role::Worker,
2789 vec!["*".into()],
2790 Duration::from_secs(600),
2791 );
2792 let fp = worker_token.fingerprint();
2793 let record = CapTokenRecord::from_worker_token(worker_token.clone(), task_id.clone());
2794 engine
2795 .with_state("test.mint_worker", move |s| {
2796 s.tokens.insert(fp, record);
2797 })
2798 .await
2799 .expect("mint worker token");
2800 worker_token
2801 }
2802
2803 /// `EngineDispatcher::dispatch` no longer stringifies the evaluated
2804 /// `Step.in` value before seeding `TaskSpec.initial_directive` — an
2805 /// Object seed must round-trip through `start_task` /
2806 /// `read_task_state` byte-for-byte as the same `Value::Object`, not a
2807 /// JSON-stringified `Value::String`.
2808 #[tokio::test]
2809 async fn object_seed_passes_through_task_spec_unchanged() {
2810 let seed = serde_json::json!({"key": "value"});
2811 let (engine, token, task_id) = seeded_engine(seed.clone()).await;
2812 let state = engine
2813 .read_task_state(&token, &task_id)
2814 .await
2815 .expect("read_task_state");
2816 assert_eq!(
2817 state.spec.initial_directive, seed,
2818 "TaskSpec.initial_directive must equal the raw Object seed, not a stringified copy"
2819 );
2820 }
2821
2822 /// `Engine::fetch_prompt` returns the `Value` end-to-end (issue #18):
2823 /// an Object seed stays a `Value::Object` and is not stringified in
2824 /// the engine layer. The Worker HTTP boundary
2825 /// (`fetch_worker_payload*`) is what performs the render down to a
2826 /// JSON literal `String` for `WorkerPayload.prompt`.
2827 #[tokio::test]
2828 async fn object_seed_passes_through_fetch_prompt_as_value() {
2829 let seed = serde_json::json!({"key": "value"});
2830 let (engine, _token, task_id) = seeded_engine(seed.clone()).await;
2831 let worker_token = mint_worker_token(&engine, &task_id).await;
2832 let prompt = engine
2833 .fetch_prompt(&worker_token, &task_id)
2834 .await
2835 .expect("fetch_prompt");
2836 assert_eq!(
2837 prompt, seed,
2838 "fetch_prompt must return the raw Object Value, not a stringified copy"
2839 );
2840 }
2841
2842 /// The Worker HTTP boundary is the render point: `fetch_worker_payload*`
2843 /// coerces the stored `Value` down to `WorkerPayload.prompt: String`
2844 /// (JSON-literal shape for non-strings). Verifies the boundary render
2845 /// stays intact for an Object seed.
2846 #[tokio::test]
2847 async fn object_seed_renders_as_json_literal_at_worker_payload_boundary() {
2848 let seed = serde_json::json!({"key": "value"});
2849 let (engine, _token, task_id) = seeded_engine(seed).await;
2850 let worker_token = mint_worker_token(&engine, &task_id).await;
2851 let payload = engine
2852 .fetch_worker_payload(&worker_token, &task_id)
2853 .await
2854 .expect("fetch_worker_payload");
2855 assert_eq!(
2856 payload.prompt, r#"{"key":"value"}"#,
2857 "WorkerPayload.prompt must be the JSON literal String render of the Value seed"
2858 );
2859 }
2860
2861 /// A `String` seed is unaffected — still passes through verbatim, both
2862 /// as the `TaskSpec.initial_directive` `Value` and as the Worker
2863 /// `fetch_prompt` return (issue #18 Invariant 2).
2864 #[tokio::test]
2865 async fn string_seed_passes_through_unchanged() {
2866 let (engine, token, task_id) = seeded_engine(serde_json::json!("do the thing")).await;
2867 let state = engine
2868 .read_task_state(&token, &task_id)
2869 .await
2870 .expect("read_task_state");
2871 assert_eq!(
2872 state.spec.initial_directive,
2873 serde_json::json!("do the thing")
2874 );
2875 let worker_token = mint_worker_token(&engine, &task_id).await;
2876 let prompt = engine
2877 .fetch_prompt(&worker_token, &task_id)
2878 .await
2879 .expect("fetch_prompt");
2880 assert_eq!(prompt, serde_json::json!("do the thing"));
2881 }
2882}
2883
2884/// GH #31: `fetch_worker_payload{,_trusted}`'s size-threshold branch
2885/// between inline (`WorkerPayload.system`) and by-reference
2886/// (`WorkerPayload.system_ref`) delivery, plus the `bake_worker_system_prompt`
2887/// `agent_render_sizes` bookkeeping that feeds `agent_last_rendered_size`.
2888#[cfg(test)]
2889mod system_ref_threshold_tests {
2890 use super::*;
2891
2892 async fn seeded_engine_with_cfg(cfg: EngineCfg) -> (Engine, CapToken, StepId) {
2893 let engine = Engine::new(cfg);
2894 let op_token = engine
2895 .attach("ut-op", Role::Operator, Duration::from_secs(30))
2896 .await
2897 .expect("attach");
2898 let task_id = engine
2899 .start_task(
2900 &op_token,
2901 TaskSpec {
2902 agent: "planner".to_string(),
2903 initial_directive: serde_json::json!("do the thing"),
2904 step_ctx: None,
2905 },
2906 )
2907 .await
2908 .expect("start_task");
2909 (engine, op_token, task_id)
2910 }
2911
2912 /// Same worker-token-minting fixture as
2913 /// `initial_directive_value_passthrough_tests::mint_worker_token`
2914 /// (kept local to this module — the two `mod`s do not share private
2915 /// helpers across `cfg(test)` boundaries).
2916 async fn mint_worker_token(engine: &Engine, task_id: &StepId) -> CapToken {
2917 let worker_token = engine.signer().session(
2918 format!("worker-of-{task_id}"),
2919 Role::Worker,
2920 vec!["*".into()],
2921 Duration::from_secs(600),
2922 );
2923 let fp = worker_token.fingerprint();
2924 let record = CapTokenRecord::from_worker_token(worker_token.clone(), task_id.clone());
2925 engine
2926 .with_state("test.mint_worker", move |s| {
2927 s.tokens.insert(fp, record);
2928 })
2929 .await
2930 .expect("mint worker token");
2931 worker_token
2932 }
2933
2934 /// Under-threshold: `system` stays inline, `system_ref` stays `None`.
2935 #[tokio::test]
2936 async fn under_threshold_stays_inline() {
2937 let (engine, _op_token, task_id) = seeded_engine_with_cfg(EngineCfg::default()).await;
2938 let worker_token = mint_worker_token(&engine, &task_id).await;
2939 let rendered = "a short system prompt".to_string();
2940 engine
2941 .bake_worker_system_prompt(&task_id, 1, Some(rendered.clone()))
2942 .await
2943 .expect("bake");
2944 let payload = engine
2945 .fetch_worker_payload(&worker_token, &task_id)
2946 .await
2947 .expect("fetch_worker_payload");
2948 assert_eq!(payload.system, Some(rendered));
2949 assert!(payload.system_ref.is_none());
2950 }
2951
2952 /// Over-threshold: `system` is cleared and `system_ref` is populated
2953 /// with a `sha256` matching the known input string. Exercises
2954 /// `fetch_worker_payload_trusted` (the `_trusted` sibling must be
2955 /// behaviorally identical to `fetch_worker_payload`).
2956 #[tokio::test]
2957 async fn over_threshold_switches_to_system_ref_with_matching_sha256() {
2958 let mut cfg = EngineCfg::default();
2959 cfg.system_ref.threshold_bytes = 16;
2960 cfg.system_ref.mode = crate::types::SystemRefMode::File;
2961 cfg.system_ref.store_dir =
2962 std::env::temp_dir().join(format!("mse-system-ref-test-{}", crate::types::now_unix()));
2963 let (engine, _op_token, task_id) = seeded_engine_with_cfg(cfg).await;
2964 let rendered =
2965 "this system prompt is deliberately longer than the 16 byte threshold".to_string();
2966 engine
2967 .bake_worker_system_prompt(&task_id, 1, Some(rendered.clone()))
2968 .await
2969 .expect("bake");
2970 let payload = engine
2971 .fetch_worker_payload_trusted(&task_id)
2972 .await
2973 .expect("fetch_worker_payload_trusted");
2974 assert!(
2975 payload.system.is_none(),
2976 "over-threshold response must not also inline `system`"
2977 );
2978 let system_ref = payload
2979 .system_ref
2980 .expect("over-threshold response must populate system_ref");
2981 assert_eq!(system_ref.size_bytes, rendered.len() as u64);
2982 assert_eq!(system_ref.mode, crate::types::SystemRefMode::File);
2983 use sha2::Digest;
2984 let expected_sha256 = hex::encode(sha2::Sha256::digest(rendered.as_bytes()));
2985 assert_eq!(system_ref.sha256, expected_sha256);
2986 assert!(system_ref.uri.starts_with("file://"));
2987 let written = tokio::fs::read_to_string(system_ref.uri.trim_start_matches("file://"))
2988 .await
2989 .expect("File mode must have written the referenced path");
2990 assert_eq!(written, rendered);
2991 }
2992
2993 /// `Http` mode never writes a file — `system_ref.uri` is the bare path
2994 /// the engine can construct on its own, scheme/host-free.
2995 #[tokio::test]
2996 async fn over_threshold_http_mode_constructs_path_only_uri() {
2997 let mut cfg = EngineCfg::default();
2998 cfg.system_ref.threshold_bytes = 16;
2999 cfg.system_ref.mode = crate::types::SystemRefMode::Http;
3000 let (engine, _op_token, task_id) = seeded_engine_with_cfg(cfg).await;
3001 let worker_token = mint_worker_token(&engine, &task_id).await;
3002 let rendered =
3003 "this system prompt is deliberately longer than the 16 byte threshold".to_string();
3004 engine
3005 .bake_worker_system_prompt(&task_id, 1, Some(rendered))
3006 .await
3007 .expect("bake");
3008 let payload = engine
3009 .fetch_worker_payload(&worker_token, &task_id)
3010 .await
3011 .expect("fetch_worker_payload");
3012 let system_ref = payload.system_ref.expect("system_ref must be populated");
3013 assert_eq!(system_ref.mode, crate::types::SystemRefMode::Http);
3014 assert_eq!(
3015 system_ref.uri,
3016 format!("/v1/worker/prompt/system?task_id={task_id}&attempt=1")
3017 );
3018 }
3019
3020 /// `bake_worker_system_prompt` records the render size keyed by agent
3021 /// name (last-write-wins), readable via `agent_last_rendered_size`.
3022 #[tokio::test]
3023 async fn bake_records_agent_render_size_last_write_wins() {
3024 let (engine, _op_token, task_id) = seeded_engine_with_cfg(EngineCfg::default()).await;
3025 assert_eq!(engine.agent_last_rendered_size("planner").await, None);
3026 engine
3027 .bake_worker_system_prompt(&task_id, 1, Some("a".repeat(10)))
3028 .await
3029 .expect("bake 1");
3030 assert_eq!(engine.agent_last_rendered_size("planner").await, Some(10));
3031 engine
3032 .bake_worker_system_prompt(&task_id, 2, Some("b".repeat(20)))
3033 .await
3034 .expect("bake 2");
3035 assert_eq!(
3036 engine.agent_last_rendered_size("planner").await,
3037 Some(20),
3038 "most-recently-observed size wins, not the largest"
3039 );
3040 }
3041}
3042
3043/// subtask-4 / ST2 rework: `submit_output` / `submit_worker_result_trusted`'s
3044/// submit-time projection sink (`Engine::materialize_final_submission`) —
3045/// the Data-plane `OutputStore` dual-write plus the
3046/// `FileProjectionAdapter`-backed file materialize, both fail-open. See
3047/// the subtask-4 Tests this module covers inline on each test.
3048#[cfg(test)]
3049mod submit_time_projection_sink_tests {
3050 use super::*;
3051 use crate::core::agent_context::AgentContextView;
3052 use crate::store::output::{ContentRef, InMemoryOutputStore, OutputEvent};
3053
3054 /// Starts a task under `agent`, returning `(engine, op_token, task_id,
3055 /// worker_token)` — same helper shape as the sibling test modules
3056 /// above (`initial_directive_value_passthrough_tests::seeded_engine` /
3057 /// `mint_worker_token`), duplicated locally per this file's
3058 /// established per-module convention.
3059 async fn seeded_task(agent: &str) -> (Engine, CapToken, StepId, CapToken) {
3060 let engine = Engine::new(EngineCfg::default());
3061 let op_token = engine
3062 .attach("ut-op", Role::Operator, Duration::from_secs(30))
3063 .await
3064 .expect("attach");
3065 let task_id = engine
3066 .start_task(
3067 &op_token,
3068 TaskSpec {
3069 agent: agent.to_string(),
3070 initial_directive: Value::String("go".into()),
3071 step_ctx: None,
3072 },
3073 )
3074 .await
3075 .expect("start_task");
3076 let worker_token = engine.signer().session(
3077 format!("worker-of-{task_id}"),
3078 Role::Worker,
3079 vec!["*".into()],
3080 Duration::from_secs(600),
3081 );
3082 let fp = worker_token.fingerprint();
3083 let record = CapTokenRecord::from_worker_token(worker_token.clone(), task_id.clone());
3084 engine
3085 .with_state("test.mint_worker", move |s| {
3086 s.tokens.insert(fp, record);
3087 })
3088 .await
3089 .expect("mint worker token");
3090 (engine, op_token, task_id, worker_token)
3091 }
3092
3093 /// Seeds `EngineState.agent_ctx[(task_id, attempt)].view` directly —
3094 /// the same snapshot `AgentContextMiddleware` writes at spawn time
3095 /// (see its module doc), stood up here without the full spawner
3096 /// stack so these tests can exercise `submit_output` in isolation.
3097 async fn seed_agent_context(engine: &Engine, task_id: &StepId, attempt: u32, work_dir: &str) {
3098 let task_id = task_id.clone();
3099 let work_dir = work_dir.to_string();
3100 engine
3101 .with_state("test.seed_agent_context", move |s| {
3102 s.agent_ctx.insert(
3103 (task_id, attempt),
3104 crate::core::state::AgentCtxEntry {
3105 view: AgentContextView {
3106 work_dir: Some(work_dir),
3107 ..Default::default()
3108 },
3109 policy: Default::default(),
3110 },
3111 );
3112 })
3113 .await
3114 .expect("seed agent_ctx");
3115 }
3116
3117 /// GH #27 (follow-up to #23): seeds `EngineState.agent_ctx` with an
3118 /// arbitrary `work_dir` / `project_root` pair (either may be `None`),
3119 /// unlike [`seed_agent_context`] (which only ever sets `work_dir`) —
3120 /// lets these tests exercise `ProjectionPlacement::resolve_root`'s
3121 /// fallback in both directions.
3122 async fn seed_agent_context_roots(
3123 engine: &Engine,
3124 task_id: &StepId,
3125 attempt: u32,
3126 work_dir: Option<&str>,
3127 project_root: Option<&str>,
3128 ) {
3129 let task_id = task_id.clone();
3130 let work_dir = work_dir.map(str::to_string);
3131 let project_root = project_root.map(str::to_string);
3132 engine
3133 .with_state("test.seed_agent_context_roots", move |s| {
3134 s.agent_ctx.insert(
3135 (task_id, attempt),
3136 crate::core::state::AgentCtxEntry {
3137 view: AgentContextView {
3138 work_dir,
3139 project_root,
3140 ..Default::default()
3141 },
3142 policy: Default::default(),
3143 },
3144 );
3145 })
3146 .await
3147 .expect("seed agent_ctx");
3148 }
3149
3150 /// GH #27 (follow-up to #23): seeds `EngineState.projection_placements`
3151 /// directly — the same snapshot `EngineDispatcher::dispatch` stashes
3152 /// at dispatch time (mirroring [`seed_step_naming`]'s contract) — so
3153 /// these tests can exercise a declared `ProjectionPlacement` without
3154 /// driving a real `Compiler::compile`.
3155 async fn seed_projection_placement(
3156 engine: &Engine,
3157 task_id: &StepId,
3158 placement: crate::core::projection_placement::ProjectionPlacement,
3159 ) {
3160 let task_id = task_id.clone();
3161 let placement = Arc::new(placement);
3162 engine
3163 .with_state("test.seed_projection_placement", move |s| {
3164 s.projection_placements.insert(task_id, placement);
3165 })
3166 .await
3167 .expect("seed projection_placements");
3168 }
3169
3170 /// GH #23 subtask-2: builds a fixture
3171 /// [`crate::core::step_naming::StepNaming`] table declaring `producer`
3172 /// → `canonical` (`AgentMeta.projection_name`), then seeds it into
3173 /// `EngineState.step_namings` for `task_id` — the same snapshot
3174 /// `EngineDispatcher::dispatch` stashes at dispatch time
3175 /// (`blueprint.rs`'s "construct once, read many" contract), stood up
3176 /// here without the full Blueprint-compile path so these tests can
3177 /// exercise the canonical-sink resolution in isolation.
3178 async fn seed_step_naming(engine: &Engine, task_id: &StepId, producer: &str, canonical: &str) {
3179 use crate::blueprint::{
3180 current_schema_version, AgentDef, AgentKind, AgentMeta, Blueprint, BlueprintMetadata,
3181 CompilerHints, CompilerStrategy,
3182 };
3183 use crate::core::step_naming::StepNaming;
3184 use mlua_flow_ir::{Expr, Node};
3185
3186 let flow = Node::Step {
3187 ref_: producer.to_string(),
3188 in_: Expr::Path {
3189 at: "$.in".to_string(),
3190 },
3191 out: Expr::Path {
3192 at: format!("$.{producer}_out"),
3193 },
3194 };
3195 let bp = Blueprint {
3196 schema_version: current_schema_version(),
3197 id: "sink-canonical-ut".into(),
3198 flow,
3199 agents: vec![AgentDef {
3200 name: producer.to_string(),
3201 kind: AgentKind::RustFn,
3202 spec: serde_json::json!({ "fn_id": producer }),
3203 profile: None,
3204 meta: Some(AgentMeta {
3205 projection_name: Some(canonical.to_string()),
3206 ..Default::default()
3207 }),
3208 }],
3209 operators: vec![],
3210 metas: vec![],
3211 hints: CompilerHints::default(),
3212 strategy: CompilerStrategy::default(),
3213 metadata: BlueprintMetadata::default(),
3214 spawner_hints: Default::default(),
3215 default_agent_kind: AgentKind::Operator,
3216 default_operator_kind: None,
3217 default_init_ctx: None,
3218 default_agent_ctx: None,
3219 default_context_policy: None,
3220 projection_placement: None,
3221 audits: vec![],
3222 degradation_policy: None,
3223 };
3224 let (naming, warnings) = StepNaming::from_blueprint(&bp).expect("no collision");
3225 assert!(warnings.is_empty(), "single-step fixture has no collisions");
3226 let naming = Arc::new(naming);
3227 let task_id = task_id.clone();
3228 engine
3229 .with_state("test.seed_step_naming", move |s| {
3230 s.step_namings.insert(task_id, naming);
3231 })
3232 .await
3233 .expect("seed step_namings");
3234 }
3235
3236 fn final_event(value: Value, ok: bool) -> crate::worker::output::OutputEvent {
3237 crate::worker::output::OutputEvent::Final {
3238 content: crate::worker::output::ContentRef::Inline { value },
3239 ok,
3240 }
3241 }
3242
3243 /// Subtask 4 Test #2: `submit_output`'s `Final` writes
3244 /// `<root>/workspace/tasks/<task_id>/ctx/<agent>.md`, content matching
3245 /// the submitted value.
3246 #[tokio::test]
3247 async fn submit_output_final_materializes_file_when_work_dir_resolved() {
3248 let dir = tempfile::TempDir::new().unwrap();
3249 let (engine, _op, task_id, worker_token) = seeded_task("planner").await;
3250 seed_agent_context(&engine, &task_id, 1, &dir.path().to_string_lossy()).await;
3251
3252 engine
3253 .submit_output(
3254 &worker_token,
3255 &task_id,
3256 1,
3257 final_event(serde_json::json!({"plan": "do it"}), true),
3258 )
3259 .await
3260 .expect("submit_output");
3261
3262 let expected_file = dir
3263 .path()
3264 .join("workspace/tasks")
3265 .join(task_id.as_str())
3266 .join("ctx/planner.md");
3267 assert!(
3268 expected_file.exists(),
3269 "materialized submission file missing at {expected_file:?}"
3270 );
3271 let body = std::fs::read_to_string(expected_file).unwrap();
3272 assert!(body.contains(r#""plan": "do it""#), "body: {body}");
3273 }
3274
3275 /// Subtask 4 Test #3: `work_dir` unresolved (no `agent_ctx`
3276 /// snapshot for this `(task_id, attempt)`) — submit still succeeds,
3277 /// fail-open, no file.
3278 #[tokio::test]
3279 async fn submit_output_final_skips_file_when_root_unresolved() {
3280 let (engine, _op, task_id, worker_token) = seeded_task("planner").await;
3281 // No seed_agent_context call — root is unresolved.
3282
3283 let result = engine
3284 .submit_output(
3285 &worker_token,
3286 &task_id,
3287 1,
3288 final_event(serde_json::json!("hi"), true),
3289 )
3290 .await;
3291 assert!(
3292 result.is_ok(),
3293 "submit must succeed even with no resolvable root (fail-open, Invariant 1)"
3294 );
3295 }
3296
3297 /// Subtask 4 Test #4 (file half): re-submitting under the same
3298 /// `(task_id, agent)` overwrites the materialized file with the
3299 /// latest value.
3300 #[tokio::test]
3301 async fn resubmit_overwrites_materialized_file_with_latest() {
3302 let dir = tempfile::TempDir::new().unwrap();
3303 let (engine, _op, task_id, worker_token) = seeded_task("planner").await;
3304 seed_agent_context(&engine, &task_id, 1, &dir.path().to_string_lossy()).await;
3305
3306 engine
3307 .submit_output(
3308 &worker_token,
3309 &task_id,
3310 1,
3311 final_event(serde_json::json!("first"), true),
3312 )
3313 .await
3314 .expect("first submit");
3315 engine
3316 .submit_output(
3317 &worker_token,
3318 &task_id,
3319 1,
3320 final_event(serde_json::json!("second"), true),
3321 )
3322 .await
3323 .expect("second submit");
3324
3325 let expected_file = dir
3326 .path()
3327 .join("workspace/tasks")
3328 .join(task_id.as_str())
3329 .join("ctx/planner.md");
3330 let body = std::fs::read_to_string(expected_file).unwrap();
3331 assert!(body.contains("second"), "body must reflect latest: {body}");
3332 assert!(
3333 !body.contains("first"),
3334 "body must not carry the stale value: {body}"
3335 );
3336 }
3337
3338 /// GH #27 (follow-up to #23): the byte-compat default
3339 /// `ProjectionPlacement` (`root_preference = WorkDir`) falls back to
3340 /// `project_root` when `work_dir` is absent — the same fallback
3341 /// [`crate::core::projection_placement::ProjectionPlacement::resolve_root`]
3342 /// now performs for every one of the "3 path" call sites, this one
3343 /// exercised at the submit-sink layer.
3344 #[tokio::test]
3345 async fn submit_output_final_falls_back_to_project_root_when_work_dir_absent() {
3346 let dir = tempfile::TempDir::new().unwrap();
3347 let (engine, _op, task_id, worker_token) = seeded_task("planner").await;
3348 seed_agent_context_roots(
3349 &engine,
3350 &task_id,
3351 1,
3352 None,
3353 Some(&dir.path().to_string_lossy()),
3354 )
3355 .await;
3356
3357 engine
3358 .submit_output(
3359 &worker_token,
3360 &task_id,
3361 1,
3362 final_event(serde_json::json!({"plan": "via project_root"}), true),
3363 )
3364 .await
3365 .expect("submit_output");
3366
3367 let expected_file = dir
3368 .path()
3369 .join("workspace/tasks")
3370 .join(task_id.as_str())
3371 .join("ctx/planner.md");
3372 assert!(
3373 expected_file.exists(),
3374 "materialized submission file missing at {expected_file:?} \
3375 (work_dir absent must fall back to project_root)"
3376 );
3377 }
3378
3379 /// GH #27 (follow-up to #23): a declared `ProjectionPlacement`
3380 /// (`root_preference = ProjectRoot`, custom `dir_template`) changes
3381 /// BOTH which root is preferred (project_root wins even though
3382 /// work_dir is also present) AND the target directory layout — proof
3383 /// the submit sink consults the snapshotted resolver rather than a
3384 /// hardcoded layout.
3385 #[tokio::test]
3386 async fn submit_output_final_uses_declared_projection_placement() {
3387 let work_dir = tempfile::TempDir::new().unwrap();
3388 let project_root = tempfile::TempDir::new().unwrap();
3389 let (engine, _op, task_id, worker_token) = seeded_task("planner").await;
3390 seed_agent_context_roots(
3391 &engine,
3392 &task_id,
3393 1,
3394 Some(&work_dir.path().to_string_lossy()),
3395 Some(&project_root.path().to_string_lossy()),
3396 )
3397 .await;
3398 seed_projection_placement(
3399 &engine,
3400 &task_id,
3401 crate::core::projection_placement::ProjectionPlacement {
3402 root_preference: crate::core::projection_placement::RootPreference::ProjectRoot,
3403 dir_template: "custom/{task_id}/out".to_string(),
3404 },
3405 )
3406 .await;
3407
3408 engine
3409 .submit_output(
3410 &worker_token,
3411 &task_id,
3412 1,
3413 final_event(serde_json::json!({"plan": "via custom placement"}), true),
3414 )
3415 .await
3416 .expect("submit_output");
3417
3418 let expected_file = project_root
3419 .path()
3420 .join("custom")
3421 .join(task_id.as_str())
3422 .join("out/planner.md");
3423 assert!(
3424 expected_file.exists(),
3425 "materialized submission file missing at custom placement target {expected_file:?}"
3426 );
3427 let unexpected_file = work_dir
3428 .path()
3429 .join("workspace/tasks")
3430 .join(task_id.as_str())
3431 .join("ctx/planner.md");
3432 assert!(
3433 !unexpected_file.exists(),
3434 "declared root_preference=ProjectRoot must not fall back to work_dir: {unexpected_file:?}"
3435 );
3436 }
3437
3438 /// Subtask 4 Invariant 3 / crux requirement #3: when
3439 /// [`Engine::set_output_store`] wires a Data-plane [`crate::store::output::OutputStore`],
3440 /// `submit_output`'s `Final` dual-writes into it under
3441 /// `producer_agent = TaskState.spec.agent` — the store becomes
3442 /// queryable via `get_latest_by_name`, independent of whether a root
3443 /// resolved for the file half.
3444 #[tokio::test]
3445 async fn submit_output_final_dual_writes_into_configured_output_store() {
3446 let (engine, _op, task_id, worker_token) = seeded_task("reviewer").await;
3447 let data_store: Arc<dyn crate::store::output::OutputStore> =
3448 Arc::new(InMemoryOutputStore::new());
3449 engine.set_output_store(data_store.clone());
3450
3451 engine
3452 .submit_output(
3453 &worker_token,
3454 &task_id,
3455 1,
3456 final_event(serde_json::json!({"verdict": "pass"}), true),
3457 )
3458 .await
3459 .expect("submit_output");
3460
3461 let record = data_store
3462 .get_latest_by_name("reviewer")
3463 .await
3464 .expect("dual-written record");
3465 match record.event {
3466 OutputEvent::Final { content, ok } => {
3467 assert!(ok);
3468 match content {
3469 ContentRef::Inline { value } => {
3470 assert_eq!(value, serde_json::json!({"verdict": "pass"}));
3471 }
3472 other => panic!("expected Inline content, got {other:?}"),
3473 }
3474 }
3475 other => panic!("expected Final event, got {other:?}"),
3476 }
3477 }
3478
3479 /// GH #34 subtask-3 gap fix: an `Artifact` event submitted via
3480 /// `submit_output` dual-writes into a wired Data-plane `OutputStore`
3481 /// under its OWN `name`, verbatim — mirrors
3482 /// `submit_output_final_dual_writes_into_configured_output_store`
3483 /// above, but for the `Artifact` variant.
3484 #[tokio::test]
3485 async fn submit_output_artifact_dual_writes_into_configured_output_store() {
3486 let (engine, _op, task_id, worker_token) = seeded_task("echo").await;
3487 let data_store: Arc<dyn crate::store::output::OutputStore> =
3488 Arc::new(InMemoryOutputStore::new());
3489 engine.set_output_store(data_store.clone());
3490
3491 engine
3492 .submit_output(
3493 &worker_token,
3494 &task_id,
3495 1,
3496 OutputEvent::Artifact {
3497 name: "audit:echo".to_string(),
3498 content: ContentRef::Inline {
3499 value: serde_json::json!({"finding": "clean"}),
3500 },
3501 },
3502 )
3503 .await
3504 .expect("submit_output");
3505
3506 let record = data_store
3507 .get_latest_by_name("audit:echo")
3508 .await
3509 .expect("dual-written artifact record");
3510 match record.event {
3511 OutputEvent::Artifact { name, content } => {
3512 assert_eq!(name, "audit:echo");
3513 match content {
3514 ContentRef::Inline { value } => {
3515 assert_eq!(value, serde_json::json!({"finding": "clean"}));
3516 }
3517 other => panic!("expected Inline content, got {other:?}"),
3518 }
3519 }
3520 other => panic!("expected Artifact event, got {other:?}"),
3521 }
3522 // The `Artifact` dual-write must never collide with / overwrite
3523 // the producing step's own `Final` name — `submit_output` never
3524 // materialized a `Final` here, so `"echo"` must stay unresolved.
3525 assert!(
3526 data_store.get_latest_by_name("echo").await.is_err(),
3527 "artifact write must not fabricate a record under the raw producer_agent name"
3528 );
3529 }
3530
3531 /// Invariant 1 (fail-open) for `Artifact`, mirroring
3532 /// `submit_output_final_skips_file_when_root_unresolved`'s Final-side
3533 /// coverage: no `OutputStore` wired at all — submit still succeeds.
3534 #[tokio::test]
3535 async fn submit_output_artifact_is_fail_open_when_no_output_store_configured() {
3536 let (engine, _op, task_id, worker_token) = seeded_task("echo").await;
3537
3538 let result = engine
3539 .submit_output(
3540 &worker_token,
3541 &task_id,
3542 1,
3543 OutputEvent::Artifact {
3544 name: "audit:echo".to_string(),
3545 content: ContentRef::Inline {
3546 value: serde_json::json!("finding"),
3547 },
3548 },
3549 )
3550 .await;
3551 assert!(
3552 result.is_ok(),
3553 "submit must succeed even with no OutputStore wired (fail-open, Invariant 1)"
3554 );
3555 }
3556
3557 /// `submit_worker_result_trusted` (the `/v1/worker/submit` short-handle
3558 /// path) triggers the exact same sink as `submit_output` — parity
3559 /// across both worker-submit entry points.
3560 #[tokio::test]
3561 async fn submit_worker_result_trusted_also_triggers_projection_sink() {
3562 let dir = tempfile::TempDir::new().unwrap();
3563 let (engine, _op, task_id, _worker_token) = seeded_task("planner").await;
3564 seed_agent_context(&engine, &task_id, 1, &dir.path().to_string_lossy()).await;
3565 let data_store: Arc<dyn crate::store::output::OutputStore> =
3566 Arc::new(InMemoryOutputStore::new());
3567 engine.set_output_store(data_store.clone());
3568
3569 engine
3570 .submit_worker_result_trusted(&task_id, 1, serde_json::json!("trusted-value"), true)
3571 .await
3572 .expect("submit_worker_result_trusted");
3573
3574 let expected_file = dir
3575 .path()
3576 .join("workspace/tasks")
3577 .join(task_id.as_str())
3578 .join("ctx/planner.md");
3579 assert!(expected_file.exists());
3580 let record = data_store
3581 .get_latest_by_name("planner")
3582 .await
3583 .expect("dual-written record");
3584 assert!(matches!(record.event, OutputEvent::Final { ok: true, .. }));
3585 }
3586
3587 /// GH #23 subtask-2 (canonical sink): a declared `projection_name`
3588 /// (`AgentMeta.projection_name`, surfaced via `StepNaming`) redirects
3589 /// `submit_output`'s Final canonical sink — both the Data-plane
3590 /// dual-write name and the materialized file stem resolve to the
3591 /// canonical name, not the raw `producer_agent`.
3592 #[tokio::test]
3593 async fn submit_output_final_uses_canonical_name_when_step_naming_declares_one() {
3594 let dir = tempfile::TempDir::new().unwrap();
3595 let (engine, _op, task_id, worker_token) = seeded_task("reviewer").await;
3596 seed_agent_context(&engine, &task_id, 1, &dir.path().to_string_lossy()).await;
3597 seed_step_naming(&engine, &task_id, "reviewer", "verdict-final").await;
3598 let data_store: Arc<dyn crate::store::output::OutputStore> =
3599 Arc::new(InMemoryOutputStore::new());
3600 engine.set_output_store(data_store.clone());
3601
3602 engine
3603 .submit_output(
3604 &worker_token,
3605 &task_id,
3606 1,
3607 final_event(serde_json::json!({"verdict": "pass"}), true),
3608 )
3609 .await
3610 .expect("submit_output");
3611
3612 let record = data_store
3613 .get_latest_by_name("verdict-final")
3614 .await
3615 .expect("dual-written record under canonical name");
3616 assert!(matches!(record.event, OutputEvent::Final { ok: true, .. }));
3617 assert!(
3618 data_store.get_latest_by_name("reviewer").await.is_err(),
3619 "raw producer_agent name must not be written once canonical resolves"
3620 );
3621
3622 let expected_file = dir
3623 .path()
3624 .join("workspace/tasks")
3625 .join(task_id.as_str())
3626 .join("ctx/verdict-final.md");
3627 assert!(
3628 expected_file.exists(),
3629 "materialized file stem must be canonical at {expected_file:?}"
3630 );
3631 }
3632
3633 /// GH #23 subtask-2: no `StepNaming` table snapshotted for this
3634 /// `task_id` (the pre-GH-#23 / no-`with_step_naming` path) is a
3635 /// defensive fail-open — the canonical sink falls back to the raw
3636 /// `producer_agent`, byte-identical to
3637 /// `submit_output_final_dual_writes_into_configured_output_store`
3638 /// above (which never calls `seed_step_naming`).
3639 #[tokio::test]
3640 async fn submit_output_final_falls_back_to_producer_agent_when_no_step_naming_table() {
3641 let (engine, _op, task_id, worker_token) = seeded_task("reviewer").await;
3642 let data_store: Arc<dyn crate::store::output::OutputStore> =
3643 Arc::new(InMemoryOutputStore::new());
3644 engine.set_output_store(data_store.clone());
3645
3646 engine
3647 .submit_output(
3648 &worker_token,
3649 &task_id,
3650 1,
3651 final_event(serde_json::json!({"verdict": "pass"}), true),
3652 )
3653 .await
3654 .expect("submit_output");
3655
3656 let record = data_store
3657 .get_latest_by_name("reviewer")
3658 .await
3659 .expect("fail-open dual-write under raw producer_agent name");
3660 assert!(matches!(record.event, OutputEvent::Final { ok: true, .. }));
3661 }
3662
3663 /// GH #23 subtask-2 (Layer 2): `OutputStore::get_latest_by_name_in_run`
3664 /// resolves the value `submit_output` dual-wrote for this exact
3665 /// `(task_id, attempt)` run, independent of `get_latest_by_name`'s
3666 /// cross-Run race (two Runs sharing a producer name never bleed into
3667 /// each other through the Run-scoped accessor).
3668 #[tokio::test]
3669 async fn submit_output_final_is_resolvable_via_run_scoped_lookup() {
3670 let (engine, _op, task_id, worker_token) = seeded_task("reviewer").await;
3671 let data_store: Arc<dyn crate::store::output::OutputStore> =
3672 Arc::new(InMemoryOutputStore::new());
3673 engine.set_output_store(data_store.clone());
3674
3675 engine
3676 .submit_output(
3677 &worker_token,
3678 &task_id,
3679 1,
3680 final_event(serde_json::json!({"verdict": "pass"}), true),
3681 )
3682 .await
3683 .expect("submit_output");
3684
3685 let record = data_store
3686 .get_latest_by_name_in_run(task_id.as_str(), 1, "reviewer")
3687 .await
3688 .expect("run-scoped lookup resolves the dual-written record");
3689 assert!(matches!(record.event, OutputEvent::Final { ok: true, .. }));
3690
3691 // A different attempt of the same task must not resolve — the
3692 // Run-scoped lookup does not fall back across attempts.
3693 assert!(
3694 data_store
3695 .get_latest_by_name_in_run(task_id.as_str(), 2, "reviewer")
3696 .await
3697 .is_err(),
3698 "a different attempt must not resolve the same-named record"
3699 );
3700 }
3701}
3702
3703/// GH #36 ST1: named multi-part worker output. Covers (a) the pure
3704/// `fold_final_and_parts` assembly `dispatch_attempt_with`'s Final-pull
3705/// delegates to, (b) `stage_worker_artifact_trusted`'s per-attempt
3706/// isolation on `EngineState.output_store` / `.worker_artifact_names` (the
3707/// same `HashMap<(StepId, u32), _>` key shape `submit_worker_result_trusted`
3708/// uses — a fresh attempt is a fresh key, so nothing to explicitly "clean
3709/// up"), and (c) the allowlist behavior that keeps a non-opt-in `Artifact`
3710/// producer (e.g. `AfterRunAuditMiddleware`) from being folded in.
3711#[cfg(test)]
3712mod named_multi_part_worker_output_tests {
3713 use super::*;
3714 use crate::worker::output::{ContentRef, OutputEvent};
3715
3716 fn artifact(name: &str, value: Value) -> OutputEvent {
3717 OutputEvent::Artifact {
3718 name: name.to_string(),
3719 content: ContentRef::Inline { value },
3720 }
3721 }
3722
3723 fn final_ev(value: Value, ok: bool) -> OutputEvent {
3724 OutputEvent::Final {
3725 content: ContentRef::Inline { value },
3726 ok,
3727 }
3728 }
3729
3730 fn names(list: &[&str]) -> Vec<String> {
3731 list.iter().map(|s| s.to_string()).collect()
3732 }
3733
3734 /// Two staged parts (both in `staged_names`) + a `Final` fold into
3735 /// `{"out", "parts"}`, each value carried through verbatim.
3736 #[test]
3737 fn fold_final_and_parts_assembles_out_and_parts_shape() {
3738 let tail = vec![
3739 artifact("summary", serde_json::json!("the summary")),
3740 artifact("diff", serde_json::json!({"lines": 3})),
3741 final_ev(serde_json::json!("final text"), true),
3742 ];
3743 let staged = names(&["summary", "diff"]);
3744 let (value, ok) = fold_final_and_parts(&tail, &staged).expect("Final present");
3745 assert!(ok);
3746 assert_eq!(
3747 value,
3748 serde_json::json!({
3749 "out": "final text",
3750 "parts": {
3751 "summary": "the summary",
3752 "diff": {"lines": 3},
3753 }
3754 })
3755 );
3756 }
3757
3758 /// Zero staged parts: the value is exactly the plain `Final` value — no
3759 /// `{"out", "parts"}` wrapping. This is the back-compat guarantee (GH
3760 /// #36 must not change the shape for a worker that never POSTs to
3761 /// `/v1/worker/artifact`).
3762 #[test]
3763 fn fold_final_and_parts_with_no_parts_returns_plain_final_value() {
3764 let tail = vec![final_ev(serde_json::json!("plain value"), true)];
3765 let (value, ok) = fold_final_and_parts(&tail, &[]).expect("Final present");
3766 assert!(ok);
3767 assert_eq!(value, serde_json::json!("plain value"));
3768 }
3769
3770 /// The same staged part `name` appearing twice in one attempt: the
3771 /// LATER (tail-order) value wins — `parts` is a `Map`, not an
3772 /// accumulating list.
3773 #[test]
3774 fn fold_final_and_parts_same_name_twice_last_write_wins() {
3775 let tail = vec![
3776 artifact("a", serde_json::json!("first")),
3777 artifact("a", serde_json::json!("second")),
3778 final_ev(serde_json::json!("f"), true),
3779 ];
3780 let staged = names(&["a"]);
3781 let (value, _ok) = fold_final_and_parts(&tail, &staged).expect("Final present");
3782 assert_eq!(
3783 value,
3784 serde_json::json!({"out": "f", "parts": {"a": "second"}})
3785 );
3786 }
3787
3788 /// No `Final` anywhere in the tail (only staged parts, e.g. the worker
3789 /// crashed before submitting) — `None`, the caller's pre-existing "no
3790 /// Final in output_tail" error path.
3791 #[test]
3792 fn fold_final_and_parts_returns_none_when_no_final_present() {
3793 let tail = vec![artifact("a", serde_json::json!("v"))];
3794 let staged = names(&["a"]);
3795 assert!(fold_final_and_parts(&tail, &staged).is_none());
3796 }
3797
3798 /// An `Artifact` on the tail whose name is NOT in `staged_names` (e.g.
3799 /// `AfterRunAuditMiddleware`'s `"audit:<step_ref>"` sidecar finding on
3800 /// an audited step's own tail) must NOT be folded into `"parts"` — the
3801 /// value stays the plain `Final` value, exactly the pre-GH-#36
3802 /// behavior for every producer that isn't the worker's own
3803 /// `/v1/worker/artifact` staging. This is the regression this fold was
3804 /// almost shipped without (see `dispatch_attempt_with`'s doc).
3805 #[test]
3806 fn fold_final_and_parts_ignores_artifacts_outside_the_staged_allowlist() {
3807 let tail = vec![
3808 final_ev(serde_json::json!({"echoed": "hi"}), true),
3809 artifact("audit:echo", serde_json::json!({"finding": "clean"})),
3810 ];
3811 // `staged_names` empty: the worker itself never staged anything —
3812 // the audit sidecar Artifact must be ignored.
3813 let (value, ok) = fold_final_and_parts(&tail, &[]).expect("Final present");
3814 assert!(ok);
3815 assert_eq!(value, serde_json::json!({"echoed": "hi"}));
3816 }
3817
3818 /// Mixed tail: one staged (allowlisted) part and one non-staged
3819 /// (audit-style) `Artifact` — only the staged one is folded in.
3820 #[test]
3821 fn fold_final_and_parts_folds_only_the_staged_subset_of_a_mixed_tail() {
3822 let tail = vec![
3823 artifact("summary", serde_json::json!("s")),
3824 artifact("audit:echo", serde_json::json!({"finding": "clean"})),
3825 final_ev(serde_json::json!("f"), true),
3826 ];
3827 let staged = names(&["summary"]);
3828 let (value, _ok) = fold_final_and_parts(&tail, &staged).expect("Final present");
3829 assert_eq!(
3830 value,
3831 serde_json::json!({"out": "f", "parts": {"summary": "s"}})
3832 );
3833 }
3834
3835 /// `stage_worker_artifact_trusted` writes onto the `(task_id, attempt)`
3836 /// key exactly like `submit_worker_result_trusted` does — a part staged
3837 /// under attempt N is invisible to an `output_tail` / allowlist read of
3838 /// attempt N+1 (a fresh attempt starts empty; nothing carries over).
3839 #[tokio::test]
3840 async fn stage_worker_artifact_trusted_is_isolated_per_attempt() {
3841 let engine = Engine::new(EngineCfg::default());
3842 let task_id = StepId::new();
3843
3844 engine
3845 .stage_worker_artifact_trusted(&task_id, 1, "a".to_string(), serde_json::json!("v1"))
3846 .await
3847 .expect("stage attempt 1");
3848
3849 let attempt_1_tail = engine.output_tail(&task_id, 1).await;
3850 assert_eq!(attempt_1_tail.len(), 1);
3851 assert!(matches!(
3852 &attempt_1_tail[0],
3853 OutputEvent::Artifact { name, .. } if name == "a"
3854 ));
3855 assert_eq!(
3856 engine.worker_artifact_names_for(&task_id, 1).await,
3857 vec!["a".to_string()]
3858 );
3859
3860 let attempt_2_tail = engine.output_tail(&task_id, 2).await;
3861 assert!(
3862 attempt_2_tail.is_empty(),
3863 "attempt 2 must not see attempt 1's staged part"
3864 );
3865 assert!(
3866 engine
3867 .worker_artifact_names_for(&task_id, 2)
3868 .await
3869 .is_empty(),
3870 "attempt 2's allowlist must not see attempt 1's staged name"
3871 );
3872 }
3873}