mlua_swarm/middleware/input_inject.rs
1//! `InputInjectMiddleware` — the `SpawnerLayer` for the Data plane's multi-in
2//! prompt injection.
3//!
4//! # Role
5//!
6//! A `SpawnerLayer` that, just before spawn, drops the list of already-
7//! registered [`OutputRef`]s into `Ctx.meta.runtime`. Downstream Operator /
8//! Spawner code (for example `mlua-swarm-server`'s `Operator::execute`) looks this
9//! key up and splices a line into the SubAgent's Spawn directive prompt
10//! along the lines of "`$IN_REFS = [out_id_1, out_id_2, ...]`, fetch these
11//! from the Store".
12//!
13//! MainAI only carries `OutputRef`s (small ids); the big bodies stay with the
14//! store owner. That keeps MainAI context tight even when ten SubAgents each
15//! stack up four-kilotoken bodies — MainAI only needs to hold the id list.
16//!
17//! This layer stays out of the Domain path (the verdict flow). See the
18//! [`crate::store::output`] module doc for the canonical narrative.
19//!
20//! # Pattern
21//!
22//! Same shape as `AgentResolver`, `ProjectNameAliasMiddleware`, and
23//! `SinkMiddleware`: edit `ctx`, call the inner spawner, done. Engine state
24//! is not touched.
25//!
26//! # Implementation status
27//!
28//! - **Current (scaffold):** `SpawnerLayer` trait impl plus injection of the
29//! `IN_REFS` list into `Ctx.meta.runtime`. Turning that into a real prompt
30//! line (literal expansion inside the Operator's directive) is done on the
31//! Operator side and is still a carry.
32//! - **Carry:** literal expansion on the Operator-directive side; a fetch path
33//! from the store for actual bodies (today we only inject the id refs — the
34//! SubAgent tool is responsible for pulling bodies down); end-to-end
35//! wire-through.
36
37use crate::core::ctx::Ctx;
38use crate::core::engine::Engine;
39use crate::middleware::SpawnerLayer;
40use crate::store::output::OutputRef;
41use crate::types::{CapToken, TaskId};
42use crate::worker::adapter::{SpawnError, SpawnerAdapter};
43use crate::worker::Worker;
44use async_trait::async_trait;
45use serde_json::Value;
46use std::sync::Arc;
47
48/// Key under `ctx.meta.runtime` that carries the `IN_REFS` list.
49///
50/// Downstream Operator / Spawner code is expected to look this key up and
51/// splice a literal line into the SubAgent's Spawn directive prompt body
52/// telling it to fetch `$IN_REFS = [<out_id>, ...]` from the store.
53pub const INPUT_REFS_KEY: &str = "input_refs";
54
55/// Multi-in prompt injection `SpawnerLayer`. Config: the list of
56/// `OutputRef`s to inject into the next spawn.
57///
58/// Per-spawn lists are built in the Blueprint (γ scope) or the Application
59/// layer, and are frozen at the moment this layer is placed in the stack.
60/// If you need to rewrite them dynamically mid-flight, do it in a different
61/// middleware or resolve on the Blueprint side.
62pub struct InputInjectMiddleware {
63 refs: Vec<OutputRef>,
64}
65
66impl InputInjectMiddleware {
67 /// Build a new layer. `refs` is the `OutputRef` list to inject into the
68 /// spawn; an empty list is fine (the initial agent).
69 pub fn new(refs: Vec<OutputRef>) -> Self {
70 Self { refs }
71 }
72
73 /// Borrow the inner refs list (tests / observers).
74 pub fn refs(&self) -> &[OutputRef] {
75 &self.refs
76 }
77}
78
79impl SpawnerLayer for InputInjectMiddleware {
80 fn wrap(&self, inner: Arc<dyn SpawnerAdapter>) -> Arc<dyn SpawnerAdapter> {
81 Arc::new(InputInjectWrapped {
82 inner,
83 refs: self.refs.clone(),
84 })
85 }
86}
87
88struct InputInjectWrapped {
89 inner: Arc<dyn SpawnerAdapter>,
90 refs: Vec<OutputRef>,
91}
92
93#[async_trait]
94impl SpawnerAdapter for InputInjectWrapped {
95 async fn spawn(
96 &self,
97 engine: &Engine,
98 ctx: &Ctx,
99 task_id: TaskId,
100 attempt: u32,
101 token: CapToken,
102 ) -> Result<Box<dyn Worker>, SpawnError> {
103 let mut new_ctx = ctx.clone();
104 let refs_json: Vec<Value> = self
105 .refs
106 .iter()
107 .map(|r| Value::String(r.0.clone()))
108 .collect();
109 new_ctx
110 .meta
111 .runtime
112 .insert(INPUT_REFS_KEY.to_string(), Value::Array(refs_json));
113 self.inner
114 .spawn(engine, &new_ctx, task_id, attempt, token)
115 .await
116 }
117}
118
119#[cfg(test)]
120mod tests {
121 use super::*;
122
123 #[test]
124 fn new_layer_holds_refs() {
125 let r1 = OutputRef::new();
126 let r2 = OutputRef::new();
127 let layer = InputInjectMiddleware::new(vec![r1.clone(), r2.clone()]);
128 assert_eq!(layer.refs(), &[r1, r2]);
129 }
130
131 #[test]
132 fn empty_refs_are_valid() {
133 let layer = InputInjectMiddleware::new(vec![]);
134 assert!(layer.refs().is_empty());
135 }
136}