Skip to main content

mlua_swarm/core/
state.rs

1//! `EngineState` — the single `Mutex`-guarded state object — plus the
2//! supporting types.
3//!
4//! `EngineState` holds every mutable piece of engine flow state (task
5//! table, session table, prompts, token records, worker handles, resume
6//! table, per-task notifiers, resources, per-attempt output events, and the
7//! event log tail). It sits on the Domain side of the Data / Domain split
8//! and is unchanged by the Data-plane (`output_store` module) refactor.
9
10use crate::types::{CapToken, Role, SessionId, StepId};
11use serde::{Deserialize, Serialize};
12use serde_json::Value;
13use std::collections::HashMap;
14use std::sync::Arc;
15use tokio::sync::{broadcast, Notify};
16
17// ─── Resume / Task ─────────────────────────────────────────────────────────
18
19/// Opaque handle identifying one `query_senior` suspend/`resume` cycle.
20/// Stored on `TaskState.suspended_on` and as the key of
21/// `EngineState.pending_resumes`.
22#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
23pub struct ResumeKey(pub String);
24
25impl ResumeKey {
26    /// Generate a random key (`RK-<12 hex bytes>`).
27    ///
28    /// The prefix moved from `R-` to `RK-` in issue #14: `R-` is reserved
29    /// for [`crate::types::RunId`], and sharing it would let a resume key
30    /// pass a run-id prefix check.
31    pub fn new() -> Self {
32        Self(format!("RK-{}", crate::types::uid_hex(12)))
33    }
34
35    /// Deterministic key for a Senior-escalation suspend on `task_id`
36    /// (`RK-senior-<task_id>`), so repeated escalations on the same task
37    /// are addressable without extra bookkeeping.
38    pub fn for_senior(task_id: &StepId) -> Self {
39        Self(format!("RK-senior-{}", task_id))
40    }
41}
42
43impl Default for ResumeKey {
44    fn default() -> Self {
45        Self::new()
46    }
47}
48
49/// Lifecycle state of a task. `Pending` is the only non-terminal,
50/// non-`Suspended` state before the first `dispatch_attempt_with`;
51/// `Pass` / `Blocked` / `Cancelled` are terminal.
52#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
53#[serde(rename_all = "snake_case")]
54pub enum TaskStatus {
55    /// Created via `start_task`, not yet dispatched.
56    Pending,
57    /// A `dispatch_attempt_with` call is in flight for this task.
58    Running,
59    /// Suspended awaiting a `query_senior`/`resume` round-trip.
60    Suspended,
61    /// The last attempt completed with `ok = true`.
62    Pass,
63    /// The last attempt completed with `ok = false` (or dispatch failed).
64    Blocked,
65    /// Cancelled via `cancel_task`.
66    Cancelled,
67}
68
69/// Static task definition supplied to `start_task`: which agent runs it
70/// and the initial prompt/directive value.
71#[derive(Debug, Clone, Serialize, Deserialize)]
72pub struct TaskSpec {
73    /// Name of the agent that should execute this task.
74    pub agent: String,
75    /// Prompt/directive value seeded for attempt 1. Passed through
76    /// verbatim from the evaluated `Step.in` (issue #18 — no premature
77    /// `Value → String` coercion at this layer). Consumers that need a
78    /// rendered `String` (the `EngineState.prompts` table feeding the
79    /// Worker HTTP path, and the WS `Spawn.directive` reminder text)
80    /// render it at their own late boundary; strings pass through
81    /// verbatim, anything else is serde-stringified.
82    pub initial_directive: Value,
83    /// GH #21 Phase 2 — the Step tier's resolved context bundle, threaded
84    /// through from `EngineDispatcher::dispatch`'s `$step_meta` envelope
85    /// resolution (`None` when the dispatched `Step.in` carried no
86    /// envelope — pre-#21-Phase-2 Blueprints unaffected). Re-read from
87    /// the spec on EVERY `Engine::dispatch_attempt_with` attempt (not
88    /// cached once), so retries and Run-rekicks all carry it; inserted
89    /// into `Ctx.meta.runtime[STEP_CTX_KEY]`
90    /// (`crate::core::agent_context::STEP_CTX_KEY`), consumed by
91    /// `crate::middleware::agent_context::AgentContextMiddleware`.
92    #[serde(default, skip_serializing_if = "Option::is_none")]
93    pub step_ctx: Option<Value>,
94    /// Per-run override for
95    /// [`crate::core::config::CheckPolicy`] — governs how submit-time
96    /// projection sinks (`Engine::materialize_final_submission` /
97    /// `Engine::materialize_artifact_submission`) react to fail-open
98    /// conditions for THIS task. `None` (the default; backward-compat
99    /// with pre-`CheckPolicy` `TaskSpec`s deserialised without the
100    /// field) falls back to `EngineCfg.check_policy` (server-wide
101    /// default). `Some(policy)` overrides for this task only — see the
102    /// `CheckPolicy` doc for the semantics of the three modes.
103    #[serde(default, skip_serializing_if = "Option::is_none")]
104    pub check_policy: Option<crate::core::config::CheckPolicy>,
105}
106
107/// The full mutable record of one task: its static `spec`, current
108/// `status`, attempt counter, and bookkeeping timestamps. Cloned out of
109/// `EngineState` on every read (e.g. by `read_task_state` / `poll_task`).
110#[derive(Debug, Clone, Serialize, Deserialize)]
111pub struct TaskState {
112    /// Unique task identifier (assigned by `start_task`).
113    pub id: StepId,
114    /// The static spec this task was created from.
115    pub spec: TaskSpec,
116    /// Current lifecycle status.
117    pub status: TaskStatus,
118    /// 1-based counter, bumped by `Engine::dispatch_attempt_with` each
119    /// time this task is dispatched.
120    pub attempt: u32,
121    /// Set while `status == Suspended`; the key needed to `resume` it.
122    pub suspended_on: Option<ResumeKey>,
123    /// Most recent result value posted via `post_result` or produced by a
124    /// completed attempt.
125    pub last_result: Option<Value>,
126    /// Unix timestamp (seconds) when the task was created.
127    pub created_at: u64,
128    /// Unix timestamp (seconds) of the last state mutation.
129    pub updated_at: u64,
130    /// Recursive swarm depth. The root (an Operator calling
131    /// `start_task`) is 0; a child spawned by a Worker calling
132    /// `start_task` is its parent's `depth + 1`. Exceeding
133    /// `EngineCfg.max_spawn_depth` raises `SpawnDepthExceeded`.
134    #[serde(default)]
135    pub spawn_depth: u32,
136}
137
138impl TaskState {
139    /// Construct a new `Pending` task with `attempt = 0` and
140    /// `spawn_depth = 0`; `created_at`/`updated_at` are set to now.
141    pub fn new(id: StepId, spec: TaskSpec) -> Self {
142        let now = crate::types::now_unix();
143        Self {
144            id,
145            spec,
146            status: TaskStatus::Pending,
147            attempt: 0,
148            suspended_on: None,
149            last_result: None,
150            created_at: now,
151            updated_at: now,
152            spawn_depth: 0,
153        }
154    }
155}
156
157/// Reserved sentinel key used by [`wrap_skip_marker`] / [`is_skip_marker`]
158/// to encode a [`DispatchOutcome::Skip`] payload as a plain
159/// `serde_json::Value` (GH #76 Skip tier: Skip tier). Documented as a reserved
160/// key on the wire: an ordinary worker payload MUST NOT contain a
161/// top-level object field named `"__mse_skip"`.
162pub const SKIP_MARKER_KEY: &str = "__mse_skip";
163
164/// Wrap `v` in the reserved `{ "__mse_skip": true, "value": v }` sentinel
165/// shape produced by a [`SubmitOutcome::Skip`] submission and consumed by
166/// [`is_skip_marker`] / [`unwrap_skip_marker`] on the read side. See
167/// [`SKIP_MARKER_KEY`] for the reserved-key contract.
168pub fn wrap_skip_marker(v: Value) -> Value {
169    let mut map = serde_json::Map::new();
170    map.insert(SKIP_MARKER_KEY.to_string(), Value::Bool(true));
171    map.insert("value".to_string(), v);
172    Value::Object(map)
173}
174
175/// Return `true` when `v` is an object with `{ "__mse_skip": true }` set —
176/// the sentinel shape [`wrap_skip_marker`] produces. Used by the
177/// dispatcher (and the flow-ir binding boundary in future subtasks) to
178/// route a Skip completion out of the ordinary Pass/Blocked value path.
179pub fn is_skip_marker(v: &Value) -> bool {
180    v.as_object()
181        .and_then(|m| m.get(SKIP_MARKER_KEY))
182        .and_then(|b| b.as_bool())
183        .unwrap_or(false)
184}
185
186/// If `v` is a skip-marker sentinel (see [`is_skip_marker`]), return the
187/// carried inner value cloned out of the `"value"` field (or `Value::Null`
188/// when the field is absent — a malformed sentinel is still a skip signal
189/// with no payload). Returns `None` otherwise, so the caller can fall back
190/// to the ordinary Pass/Blocked path.
191pub fn unwrap_skip_marker(v: &Value) -> Option<Value> {
192    if !is_skip_marker(v) {
193        return None;
194    }
195    Some(
196        v.as_object()
197            .and_then(|m| m.get("value"))
198            .cloned()
199            .unwrap_or(Value::Null),
200    )
201}
202
203/// Result of a `dispatch_attempt_with` call (or the conceptual outcome of
204/// a task attempt more broadly).
205///
206/// `#[non_exhaustive]` (GH #76 Skip tier) so future tier additions (e.g. a
207/// `Deferred` sibling of `Skip`) are additive — external crates cannot
208/// exhaustively match on this enum, and must include a `_` arm.
209#[derive(Debug, Clone, Serialize, Deserialize)]
210#[non_exhaustive]
211pub enum DispatchOutcome {
212    /// The attempt completed with `ok = true`; carries the result value.
213    Pass(Value),
214    /// The attempt completed with `ok = false`, or dispatch itself failed;
215    /// carries the result/error value.
216    Blocked(Value),
217    /// GH #76 Skip tier: sibling tier of [`Self::Pass`] / [`Self::Blocked`]. The
218    /// worker completed successfully (`ok = true`) but declared its
219    /// output is NOT applicable to the surrounding flow — the enclosing
220    /// [`crate::blueprint::EngineDispatcher::dispatch`] treats this as
221    /// flow-continuation (does NOT propagate an error to flow-ir) while
222    /// short-circuiting the write to the step's declared `out` binding.
223    /// Carries the returning agent's verdict payload for observability
224    /// (`StepEntry.status = "skipped"`); downstream `$.<step_id>`
225    /// references see whatever pre-existing value the binding held
226    /// (typically absent). Mirrors `spawn_halt`'s "ok=true + marker" wire
227    /// pattern.
228    Skip(Value),
229    /// The task suspended (e.g. via `query_senior`) before completing;
230    /// carries the key needed to `resume` it.
231    Suspended(ResumeKey),
232    /// The task was cancelled before completing.
233    Cancelled,
234    /// The attempt did not complete within the allotted time.
235    Timeout,
236}
237
238/// GH #76 Skip tier: the completion tier a caller of
239/// [`crate::core::engine::Engine::submit_worker_result_trusted`] signals
240/// alongside its `value`. Sibling to `DispatchOutcome` (the engine-side
241/// outcome enum) — this one is the CALLER's intent enum, whereas
242/// [`DispatchOutcome`] is the reduced outcome the engine derives from the
243/// completed attempt.
244///
245/// Mapping into the wire shape stored in `EngineState.output_store`'s
246/// terminal `OutputEvent::Final`:
247///
248/// | outcome  | `Final.ok` | `Final.content`                                   |
249/// |----------|------------|---------------------------------------------------|
250/// | `Pass`   | `true`     | `value` verbatim                                  |
251/// | `Blocked`| `false`    | `value` verbatim                                  |
252/// | `Skip`   | `true`     | [`wrap_skip_marker(value)`](wrap_skip_marker)     |
253///
254/// `#[non_exhaustive]` so future tier additions stay additive.
255#[derive(Debug, Clone, Copy, PartialEq, Eq)]
256#[non_exhaustive]
257pub enum SubmitOutcome {
258    /// Ordinary success — the worker's output is the step's value.
259    Pass,
260    /// Ordinary failure — the worker itself declared `ok = false`. Maps
261    /// to `DispatchOutcome::Blocked` at dispatch time; the
262    /// verdict-contract completion check is exempt (same "ok=false is
263    /// exempt" rule the pre-Skip world already applied).
264    Blocked,
265    /// Skip tier (GH #76): ok=true for flow-continuation purposes but the
266    /// value is wrapped in the skip-marker sentinel so the dispatcher can
267    /// route it to `DispatchOutcome::Skip` and downstream binding-write
268    /// paths can short-circuit. The verdict-contract completion check is
269    /// intentionally skipped — a Skip is the agent declaring "not
270    /// applicable", not a real verdict value.
271    Skip,
272}
273
274// ─── Session ───────────────────────────────────────────────────────────────
275
276/// Persisted record of one attached Operator session: identity, role,
277/// heartbeat bookkeeping, owned tasks, and the `OperatorKind` cascade
278/// inputs plus registry IDs used to rebuild `OperatorInfo` on dispatch
279/// (see `Engine::resolve_operator_info`).
280#[derive(Debug, Clone, Serialize, Deserialize)]
281pub struct OperatorSession {
282    /// Unique session identifier (distinct from the token nonce).
283    pub id: SessionId,
284    /// Caller-supplied name identifying the Operator (not necessarily
285    /// unique across sessions).
286    pub operator_id: String,
287    /// Role the session's token was minted with.
288    pub role: Role,
289    /// Unix timestamp (seconds) when the session was attached.
290    pub attached_at: u64,
291    /// Unix timestamp (seconds) of the last heartbeat/attach touch.
292    pub last_seen: u64,
293    /// Whether the session is currently considered live. Flipped to
294    /// `false` by `detach` or by `start_detach_loop` on a heartbeat miss.
295    pub attached: bool,
296    /// Task IDs started by this session (via `start_task` while this
297    /// session's token was current).
298    pub owned_task_ids: Vec<StepId>,
299    /// Fingerprint (`CapToken::fingerprint`, SHA-256 of the nonce) of the
300    /// `CapToken` this session was attached with; used to look sessions up
301    /// by token in `with_state` closures. Holds the fingerprint rather
302    /// than the nonce so the session table carries no secret material
303    /// (issue #14).
304    pub token_fp: String,
305    /// The Operator's `kind`, plus IDs of
306    /// the `SeniorBridge` / `SpawnHook` registered on the engine's
307    /// `BridgeRegistry`. Persisted (all `String`; no `Arc<dyn ...>`). At
308    /// `dispatch_attempt` time the engine looks these up in the registry
309    /// and builds an `OperatorInfo` to inject into `Ctx`.
310    ///
311    /// # 4-tier `OperatorKind` cascade — "Runtime Global" tier
312    ///
313    /// This field is the literal value passed to `Engine::attach_with_ids`'s
314    /// `kind` parameter, and is fed to `crate::core::ctx::collapse_operator_kind`
315    /// as the `runtime_global` tier verbatim: `Some(_)` is always an
316    /// explicit Runtime Global request that outranks both BP tiers — even
317    /// `Some(OperatorKind::Automate)` — and `None` means "not requested",
318    /// letting the BP-level tiers (`bp_agent_kinds` / `bp_global_kind`) take
319    /// over. `#[serde(default)]` keeps existing persisted sessions (from
320    /// before this field existed / was `Option`) deserializing as `None`.
321    /// See `crate::core::ctx::collapse_operator_kind` for the full cascade +
322    /// rationale.
323    #[serde(default)]
324    pub operator_kind: Option<crate::core::ctx::OperatorKind>,
325    /// "Runtime Agent-level" tier (highest priority) of the `OperatorKind`
326    /// cascade — per-agent override supplied at task-launch time via
327    /// `TaskLaunchInput.operator_kind_overrides` / `TaskApplicationInput
328    /// .operator_kind_overrides`. Keyed by `AgentDef.name`.
329    #[serde(default)]
330    pub runtime_agent_kinds: HashMap<String, crate::core::ctx::OperatorKind>,
331    /// "BP Agent-level" tier of the `OperatorKind` cascade — baked at
332    /// `TaskLaunchService::launch` time from `Blueprint.operators[].kind`,
333    /// resolved per-agent via `AgentDef.spec.operator_ref`. Keyed by
334    /// `AgentDef.name` (not `OperatorDef.name`).
335    #[serde(default)]
336    pub bp_agent_kinds: HashMap<String, crate::core::ctx::OperatorKind>,
337    /// "BP Global" tier of the `OperatorKind` cascade — baked at
338    /// `TaskLaunchService::launch` time from `Blueprint.default_operator_kind`.
339    #[serde(default)]
340    pub bp_global_kind: Option<crate::core::ctx::OperatorKind>,
341    /// ID of the `Arc<dyn SeniorBridge>` registered on the engine's
342    /// `BridgeRegistry`, if any; resolved back into `OperatorInfo.senior_bridge`.
343    #[serde(default)]
344    pub bridge_id: Option<String>,
345    /// ID of the `Arc<dyn SpawnHook>` registered on the engine's
346    /// `BridgeRegistry`, if any; resolved back into `OperatorInfo.spawn_hook`.
347    #[serde(default)]
348    pub hook_id: Option<String>,
349    /// ID of the `Arc<dyn Operator>` registered on the `OperatorRegistry`.
350    /// Used by `OperatorDelegateMiddleware` when `kind = MainAi` /
351    /// `Composite` and `operator_id` is `Some`: it delegates the entire
352    /// spawn to `operator.execute`.
353    #[serde(default)]
354    pub operator_backend_id: Option<String>,
355}
356
357// ─── Token record (= server-side counter holder) ──────────────────────────
358
359/// Server-side counter/state holder paired 1:1 with a minted `CapToken`
360/// (keyed by nonce in `EngineState.tokens`). Tracks remaining uses,
361/// revocation, and — for Worker tokens — the task the token is bound to.
362#[derive(Debug, Clone, Serialize, Deserialize)]
363pub struct CapTokenRecord {
364    /// The token this record backs.
365    pub token: CapToken,
366    /// Remaining number of verb-consuming calls. `None` means unlimited
367    /// (session-style tokens); `Some(0)` makes `consume` fail.
368    pub uses_left: Option<u32>, // None = unlimited (session)
369    /// When `true`, `consume` always fails regardless of `uses_left`.
370    pub revoked: bool,
371    /// The task this Worker token is bound to (set when minted via
372    /// `dispatch_attempt`). Used on two axes:
373    ///   1. **Depth tracking.** When a Worker calls `start_task` to spawn a
374    ///      child, the child receives this task's `spawn_depth + 1`.
375    ///   2. **Ownership gate.** When a Worker calls a state-touch verb
376    ///      (`fetch_prompt` / `post_result` / `read_task_state` /
377    ///      `cancel_task` / `poll_task`), the argument's `task_id` must
378    ///      match this value. `start_task`
379    ///      and `dispatch_attempt` are exempt — recursive swarming must
380    ///      stay open, and depth is capped by `max_spawn_depth`.
381    ///
382    ///      Operator tokens (minted at attach time) leave this `None`, so
383    ///      they can touch any task.
384    #[serde(default)]
385    pub task_id: Option<StepId>,
386}
387
388impl CapTokenRecord {
389    /// Wrap a freshly minted `CapToken` with no bound task (`task_id =
390    /// None`) — the shape used for Operator/session tokens.
391    pub fn from_token(token: CapToken) -> Self {
392        Self {
393            uses_left: token.max_uses,
394            token,
395            revoked: false,
396            task_id: None,
397        }
398    }
399
400    /// Convenience constructor used when minting a Worker token — binds
401    /// the record to the target task.
402    pub fn from_worker_token(token: CapToken, task_id: StepId) -> Self {
403        Self {
404            uses_left: token.max_uses,
405            token,
406            revoked: false,
407            task_id: Some(task_id),
408        }
409    }
410
411    /// Consume one use. `None` (session token) always returns `Ok`;
412    /// `Some(0)` returns `Err`.
413    pub fn consume(&mut self) -> Result<(), CapTokenConsumeError> {
414        if self.revoked {
415            return Err(CapTokenConsumeError::Revoked);
416        }
417        match self.uses_left.as_mut() {
418            None => Ok(()),
419            Some(0) => Err(CapTokenConsumeError::Exhausted),
420            Some(n) => {
421                *n -= 1;
422                Ok(())
423            }
424        }
425    }
426}
427
428/// Why [`CapTokenRecord::consume`] refused to spend a use.
429#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
430pub enum CapTokenConsumeError {
431    /// The record was explicitly revoked (`revoked = true`); revocation
432    /// is permanent and independent of `uses_left`.
433    #[error("token revoked")]
434    Revoked,
435    /// The record's `uses_left` budget (`Some(0)`) is spent.
436    #[error("token uses exhausted")]
437    Exhausted,
438}
439
440// ─── Event ─────────────────────────────────────────────────────────────────
441
442/// Engine lifecycle event. Every event is both appended to
443/// `EngineState.event_log_tail` (in-process ring buffer) and broadcast on
444/// `Engine::event_tx` for live subscribers.
445#[derive(Debug, Clone, Serialize, Deserialize)]
446#[serde(tag = "kind", rename_all = "snake_case")]
447pub enum Event {
448    /// A session was attached (`attach` / `attach_with` / `attach_with_ids`).
449    SessionAttached {
450        /// The newly attached session.
451        session_id: SessionId,
452        /// Role its token was minted with.
453        role: Role,
454    },
455    /// A session was detached (`detach`, or a heartbeat-miss timeout).
456    SessionDetached {
457        /// The session that was detached.
458        session_id: SessionId,
459    },
460    /// A new task was created via `start_task`.
461    TaskCreated {
462        /// The newly created task.
463        task_id: StepId,
464    },
465    /// An attempt began dispatching (not currently emitted by
466    /// `dispatch_attempt_with`; reserved for future use).
467    TaskAttemptStarted {
468        /// The task being dispatched.
469        task_id: StepId,
470        /// The attempt number.
471        attempt: u32,
472    },
473    /// An attempt finished, Pass or Blocked, with the resulting value.
474    TaskAttemptCompleted {
475        /// The task whose attempt completed.
476        task_id: StepId,
477        /// The attempt number that completed.
478        attempt: u32,
479        /// The result value produced by the attempt.
480        result: Value,
481    },
482    /// The task attempt completed with `ok = true`.
483    TaskPass {
484        /// The task that passed.
485        task_id: StepId,
486        /// The result value.
487        result: Value,
488    },
489    /// The task attempt completed with `ok = false`.
490    TaskBlocked {
491        /// The task that was blocked.
492        task_id: StepId,
493        /// The result/error value.
494        result: Value,
495    },
496    /// A worker appended an `OutputEvent` via `submit_output`.
497    WorkerOutput {
498        /// The task the output belongs to.
499        task_id: StepId,
500        /// The attempt the output belongs to.
501        attempt: u32,
502        /// The appended output event.
503        event: crate::worker::output::OutputEvent,
504    },
505    /// The task suspended pending a `resume` for `key`.
506    TaskSuspended {
507        /// The suspended task.
508        task_id: StepId,
509        /// The key needed to `resume` it.
510        key: ResumeKey,
511    },
512    /// The task resumed after `resume(key, ..)` was called.
513    TaskResumed {
514        /// The resumed task.
515        task_id: StepId,
516        /// The key that was resumed.
517        key: ResumeKey,
518    },
519    /// The task was cancelled via `cancel_task`.
520    TaskCancelled {
521        /// The cancelled task.
522        task_id: StepId,
523    },
524    /// `query_senior` was called, asking `question` on behalf of `task_id`.
525    SeniorQueried {
526        /// The task that triggered the query.
527        task_id: StepId,
528        /// The question posed to the Senior.
529        question: Value,
530    },
531    /// A Senior's `answer` was stored via `resume`.
532    SeniorAnswered {
533        /// The task the answer applies to.
534        task_id: StepId,
535        /// The Senior's answer.
536        answer: Value,
537    },
538}
539
540/// Receiver half of the engine-wide `Event` broadcast channel, obtained
541/// via `Engine::subscribe`.
542pub type EventStream = broadcast::Receiver<Event>;
543
544// ─── Resume pending (= Notify-based wait + stored answer) ─────────────────
545
546/// Entry for a task suspended via `query_senior`, waiting to be resumed.
547///
548/// The `Notify` + `answer: Option<Value>` form (rather than a oneshot
549/// channel) is deliberate: the answer stays inside `EngineState` even if
550/// the caller (an Operator) **detaches and reattaches**, so it can pull
551/// the answer out via `await_resume` after reattach.
552#[derive(Debug, Clone)]
553pub struct ResumePending {
554    /// Wakes any `await_resume` waiter once `answer` is set.
555    pub notify: Arc<Notify>,
556    /// The stored answer, once `resume` has been called for this key.
557    pub answer: Option<Value>,
558}
559
560impl ResumePending {
561    /// Create an unanswered pending entry (fresh `Notify`, `answer = None`).
562    pub fn new() -> Self {
563        Self {
564            notify: Arc::new(Notify::new()),
565            answer: None,
566        }
567    }
568}
569
570impl Default for ResumePending {
571    fn default() -> Self {
572        Self::new()
573    }
574}
575
576// ─── EngineState (= the locked thing) ──────────────────────────────────────
577
578/// One `(task_id, attempt)` entry of [`EngineState::agent_ctx`] — the
579/// materialized [`crate::core::agent_context::AgentContextView`] (Contract
580/// C, GH #20) and the effective [`mlua_swarm_schema::ContextPolicy`]
581/// `AgentContextMiddleware` already applied to it (`projection-adapter`
582/// ST5), folded into one struct (GH #23) so the two values — written
583/// together at the same single insert site — can no longer drift apart.
584#[derive(Debug, Clone, Default)]
585pub struct AgentCtxEntry {
586    /// The materialized view; the Worker axis's read-back source
587    /// (`Engine::fetch_worker_payload{,_trusted}` / `Engine::agent_context_for`).
588    pub view: crate::core::agent_context::AgentContextView,
589    /// The resolved policy already applied to `view`; read back by
590    /// `Engine::context_policy_for` so a consumer can filter a *different*
591    /// step's pointer list against this key's policy without re-deriving it
592    /// from the Blueprint.
593    pub policy: mlua_swarm_schema::ContextPolicy,
594}
595
596/// The single `Mutex`-guarded blob of engine flow state, accessed only
597/// through `Engine::with_state` (see the R1-R4 discipline documented
598/// there).
599#[derive(Debug)]
600pub struct EngineState {
601    /// All known tasks, keyed by `StepId`.
602    pub tasks: HashMap<StepId, TaskState>,
603    /// All attached/detached sessions, keyed by `SessionId`.
604    pub sessions: HashMap<SessionId, OperatorSession>,
605    /// Per-`(task_id, attempt)` prompt/directive value, seeded from
606    /// `TaskSpec.initial_directive` and fetched via `fetch_prompt`. Held
607    /// as `serde_json::Value` end-to-end (issue #18): the render down to
608    /// `String` (strings verbatim, anything else serde-stringified)
609    /// happens only at the two consumer boundaries — the Worker HTTP
610    /// path (`Engine::fetch_worker_payload*` → `WorkerPayload.prompt:
611    /// String`) and the WS Spawn frame text render
612    /// (`operator_ws::session::default_spawn_directive_with_task_directive`).
613    /// Engine-internal storage and `Engine::fetch_prompt` keep the
614    /// `Value` end-to-end, so `Step.in` `Object` / `Array` seeds are not
615    /// prematurely flattened.
616    pub prompts: HashMap<(StepId, u32), Value>,
617    /// Per-attempt `system_prompt`: `AgentDef.profile.system_prompt` is
618    /// baked at compile time, rendered inside `OperatorSpawner::spawn`,
619    /// and stashed here for the SubAgent to fetch alongside its prompt via
620    /// `HTTP /v1/worker/prompt`. The value is `Option<String>` so a missing
621    /// profile can be distinguished: an absent key means "not yet baked",
622    /// while `Some(None)` means "baked and profile is explicitly absent".
623    pub systems: HashMap<(StepId, u32), Option<String>>,
624    /// GH #31: per-agent-name "most-recently-baked `system_prompt` size"
625    /// bookkeeping, updated by `Engine::bake_worker_system_prompt`
626    /// whenever `system.is_some()` (keyed by `TaskState.spec.agent`, the
627    /// same lookup `fetch_worker_payload` uses). Last-write-wins is the
628    /// deliberate semantics — "most-recently-observed", not "largest" or
629    /// "per-attempt" — since this exists only to let `bp_doctor`-style
630    /// tooling ask "how big does this agent's rendered system prompt
631    /// currently run" without plumbing the live `Engine` into
632    /// `BlueprintsState`. Read via `Engine::agent_last_rendered_size`.
633    pub agent_render_sizes: HashMap<String, usize>,
634    /// Per-attempt materialized [`crate::core::agent_context::AgentContextView`]
635    /// (Contract C, GH #20) paired with the effective
636    /// [`mlua_swarm_schema::ContextPolicy`] `AgentContextMiddleware` already
637    /// applied to that same view — the Worker axis's read-back source.
638    /// Written by `crate::middleware::agent_context::AgentContextMiddleware`
639    /// (innermost spawner layer) at dispatch time via one insert (GH #23:
640    /// folded from two separately-keyed maps — see below); read by
641    /// `Engine::fetch_worker_payload{,_trusted}` / `Engine::agent_context_for`
642    /// (the [`AgentCtxEntry::view`] half, threaded into `WorkerPayload.context`)
643    /// and `Engine::context_policy_for` (the [`AgentCtxEntry::policy`] half,
644    /// read by `crates/mlua-swarm-server/src/worker.rs`'s `GET
645    /// /v1/worker/prompt` handler to filter `WorkerPayload.context.steps` via
646    /// `ContextPolicy::allows_step` without re-deriving the policy from the
647    /// Blueprint at fetch time). Keyed the same way as `prompts` / `systems`
648    /// — `Ctx` itself is not stored, so the entry has to be snapshotted here
649    /// to still be servable at fetch time. An absent entry (pre-ST5 spawns,
650    /// or a spawner stack that never layered `AgentContextMiddleware`) means
651    /// no materialized view and a pass-all policy
652    /// (`mlua_swarm_schema::ContextPolicy::default()`) for that key.
653    ///
654    /// GH #23: this map used to be two separately-keyed maps
655    /// (`agent_contexts: HashMap<_, AgentContextView>` /
656    /// `context_policies: HashMap<_, ContextPolicy>`) sharing a key and kept
657    /// in sync only by convention (the single insert site in
658    /// `AgentContextMiddleware` writing both in the same `with_state` call).
659    /// Folding both values into one map of an [`AgentCtxEntry`] struct makes
660    /// that pairing structural — the two values can no longer drift apart
661    /// key-by-key. TODO(GH #23): this map (nor `prompts` / `systems`) still
662    /// has no removal path — entries accumulate for the process lifetime; a
663    /// long-running server needs a task-completion sweep (tracked
664    /// separately from this fold).
665    pub agent_ctx: HashMap<(StepId, u32), AgentCtxEntry>,
666    /// GH #23: per-dispatch snapshot of the Blueprint-wide
667    /// [`crate::core::step_naming::StepNaming`] addressing-space table —
668    /// built once by `blueprint::compiler::Compiler::compile` and stashed
669    /// here, per `StepId`, by `crate::blueprint::EngineDispatcher::dispatch`
670    /// (its single insert site; the same `Arc` is shared across every
671    /// Step dispatched from the same Blueprint launch). `Engine::step_naming_for`
672    /// reads it back so later consumers (GH #23 subtask-2/3 —
673    /// `ContextPolicy.allows_step`, `StepPointer`/`StepSummary` assembly,
674    /// the REST `:step` resolver, `FileProjectionAdapter`) do not have to
675    /// re-derive the table from the Blueprint at read time. An absent
676    /// entry means the dispatcher was never given a `StepNaming` (e.g. a
677    /// direct `EngineDispatcher::with_spawner` caller that skipped
678    /// `with_step_naming`) — callers fall back to the pre-GH-#23 runtime
679    /// union rule in that case.
680    pub step_namings: HashMap<StepId, Arc<crate::core::step_naming::StepNaming>>,
681    /// GH #27 (follow-up to #23): per-dispatch snapshot of the
682    /// Blueprint-wide [`crate::core::projection_placement::ProjectionPlacement`]
683    /// resolver — built once by `blueprint::compiler::Compiler::compile`
684    /// and stashed here, per `StepId`, by
685    /// `crate::blueprint::EngineDispatcher::dispatch` (the same insert
686    /// site, and the same "construct once, read many" contract, as
687    /// [`Self::step_namings`]). `Engine::projection_placement_for` reads
688    /// it back so every one of the 3 materialize call sites (submit-time
689    /// sink, server read-back, spawn-time pointer) resolves the SAME root
690    /// preference / directory template. An absent entry means the
691    /// dispatcher was never given a `ProjectionPlacement` (e.g. a direct
692    /// `EngineDispatcher::with_spawner` caller that skipped
693    /// `with_projection_placement`) — callers fall back to
694    /// `ProjectionPlacement::default()` (byte-compat with the pre-#27
695    /// hardcoded layout) in that case.
696    pub projection_placements:
697        HashMap<StepId, Arc<crate::core::projection_placement::ProjectionPlacement>>,
698    /// All minted `CapToken` records, keyed by token fingerprint
699    /// (`CapToken::fingerprint` = SHA-256 of the nonce; issue #14 — the
700    /// key is loggable, the nonce is not).
701    pub tokens: HashMap<String, CapTokenRecord>, // key = token fingerprint
702    /// Short worker handle (`wh-XXXXXXXX`, 12 chars) → token-fingerprint
703    /// lookup map. Resolves the `worker_handle` field a SubAgent receives
704    /// with its prompt. There is no signature verification: `task_id` is
705    /// resolved by a plain `HashMap` lookup — deliberately thin for the
706    /// local running over WebSocket, and adopted specifically to remove
707    /// the base64 copy-paste failure mode.
708    pub worker_handles: HashMap<String, String>,
709    /// Outstanding `query_senior` suspensions awaiting `resume`.
710    pub pending_resumes: HashMap<ResumeKey, ResumePending>,
711    /// Per-task notifier — `notify_waiters` fires on every task-status
712    /// change. Used by `poll_task` on the caller side, and by callers that
713    /// need to `await` again after detach/reattach.
714    pub task_notifies: HashMap<StepId, Arc<Notify>>,
715    /// Arbitrary named resources set via `set_resource` and read via
716    /// `fetch_data`.
717    pub resources: HashMap<String, Value>,
718    /// Per-attempt output-event log. The `SpawnerAdapter` appends via
719    /// `submit_output`; the dispatch path pulls the terminal
720    /// `OutputEvent::Final` off the tail and decides Pass / Blocked.
721    pub output_store: HashMap<(StepId, u32), Vec<crate::worker::output::OutputEvent>>,
722    /// GH #36 ST1 (named multi-part worker output): the set of `Artifact`
723    /// names a WORKER staged for `(task_id, attempt)`, in staging order.
724    ///
725    /// Two population paths, one per lane, both meaning "the worker itself
726    /// staged this part": `Engine::stage_worker_artifact_trusted` (= `POST
727    /// /v1/worker/artifact`, the out-of-process lane) and
728    /// `crate::worker::output::EngineSink::emit` (the in-process lane's
729    /// `WorkerInvocation.sink`, which `InProcSpawner::spawn` is the sole
730    /// constructor of — so an `Artifact` arriving through it is by
731    /// construction the worker's own).
732    ///
733    /// `output_store` is a SHARED per-attempt tail — besides a worker's own
734    /// staged parts, other producers append `OutputEvent::Artifact` events
735    /// to the SAME tail too (e.g. `AfterRunAuditMiddleware`'s
736    /// `"audit:<step_ref>"` sidecar finding, an intentionally
737    /// BP-chain-invisible observation — see `Engine::submit_output`'s doc,
738    /// "`Artifact` dual-write" section). `Engine::dispatch_attempt_with`'s
739    /// Final-pull assembly must fold ONLY a worker's own named parts into
740    /// `"parts"`, not every `Artifact` that happens to land on the tail —
741    /// this set is that distinguishing signal, so a step under audit keeps
742    /// its BP-chain value byte-identical to pre-GH-#36 unless the WORKER
743    /// itself opted in.
744    pub worker_artifact_names: HashMap<(StepId, u32), Vec<String>>,
745    /// Bounded in-process tail of recent `Event`s (most recent last),
746    /// trimmed to `event_log_max` by `push_event`.
747    pub event_log_tail: Vec<Event>,
748    /// Maximum length of `event_log_tail` before older entries are
749    /// dropped.
750    pub event_log_max: usize,
751    /// Per-attempt normalized worker stats reported by worker
752    /// boundaries (`Engine::record_worker_stats`) and drained by the
753    /// dispatcher's outcome fold (`Engine::take_worker_stats`) into the
754    /// terminal `StepEntry`. Entries a dispatcher never drains (direct
755    /// `dispatch_attempt_with` callers without an `EngineDispatcher`)
756    /// share the process-lifetime accumulation caveat of `prompts` /
757    /// `agent_ctx` (GH #23 sweep TODO).
758    pub worker_stats: HashMap<(StepId, u32), crate::store::trace::WorkerStats>,
759    /// Per-dispatch [`crate::store::trace::TraceHandle`] registry —
760    /// inserted by `EngineDispatcher::dispatch` before spawning a step
761    /// (when its `RunContext` carries a trace handle), removed after
762    /// the outcome fold. Middlewares and other engine-adjacent writers
763    /// read it via `Engine::trace_handle` to append their own trace
764    /// kinds without any plumbing through `Ctx`. Known limitation
765    /// (holistic review, LOW): a dispatch future dropped between insert
766    /// and fold (sync-launch timeout race, caller abort) strands its
767    /// entry for the process lifetime — the same accumulation caveat as
768    /// `worker_stats` / `prompts` (GH #23 sweep TODO); each stranded
769    /// entry is one map slot + an `Arc` clone, not a per-run buffer.
770    pub trace_handles: HashMap<StepId, crate::store::trace::TraceHandle>,
771}
772
773impl EngineState {
774    /// Append `name` to the worker's own staged-part allowlist for
775    /// `(task_id, attempt)` — see [`Self::worker_artifact_names`].
776    ///
777    /// One statement, but shared by both lanes' population paths so the
778    /// "which map, keyed how" decision lives in one place: a lane that
779    /// records into the wrong shape would silently drop its parts out of
780    /// the `{out, parts}` fold rather than fail.
781    pub(crate) fn record_worker_artifact_name(
782        &mut self,
783        task_id: StepId,
784        attempt: u32,
785        name: String,
786    ) {
787        self.worker_artifact_names
788            .entry((task_id, attempt))
789            .or_default()
790            .push(name);
791    }
792
793    /// Construct an empty `EngineState` with `event_log_max = 1024`.
794    pub fn new() -> Self {
795        Self {
796            tasks: HashMap::new(),
797            sessions: HashMap::new(),
798            prompts: HashMap::new(),
799            systems: HashMap::new(),
800            agent_render_sizes: HashMap::new(),
801            agent_ctx: HashMap::new(),
802            step_namings: HashMap::new(),
803            projection_placements: HashMap::new(),
804            tokens: HashMap::new(),
805            worker_handles: HashMap::new(),
806            pending_resumes: HashMap::new(),
807            task_notifies: HashMap::new(),
808            resources: HashMap::new(),
809            output_store: HashMap::new(),
810            worker_artifact_names: HashMap::new(),
811            event_log_tail: Vec::new(),
812            event_log_max: 1024,
813            worker_stats: HashMap::new(),
814            trace_handles: HashMap::new(),
815        }
816    }
817
818    /// Ensure a per-task `Notify` exists; return the existing one if any.
819    pub fn ensure_task_notify(&mut self, task_id: &StepId) -> Arc<Notify> {
820        self.task_notifies
821            .entry(task_id.clone())
822            .or_insert_with(|| Arc::new(Notify::new()))
823            .clone()
824    }
825
826    /// Append `ev` to `event_log_tail`, trimming the oldest entries once
827    /// `event_log_max` is exceeded.
828    pub fn push_event(&mut self, ev: Event) {
829        self.event_log_tail.push(ev);
830        if self.event_log_tail.len() > self.event_log_max {
831            let overflow = self.event_log_tail.len() - self.event_log_max;
832            self.event_log_tail.drain(..overflow);
833        }
834    }
835}
836
837impl Default for EngineState {
838    fn default() -> Self {
839        Self::new()
840    }
841}