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}