Skip to main content

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}