Skip to main content

mlua_swarm/middleware/
task_input.rs

1//! `TaskInputMiddleware` — a `SpawnerLayer` that propagates task-level
2//! execution context (`project_root` / `work_dir` / `task_metadata`) into
3//! `Ctx.meta.runtime`.
4//!
5//! Issue #19 ST2: these three fields are canonical Task-level data, kept
6//! independent of `init_ctx` (the flow.ir initial `ctx` value, a pure eval
7//! seed). [`crate::service::task_launch::TaskLaunchService::launch`] reads
8//! them off [`crate::service::task_launch::TaskLaunchInput::task_input`]
9//! (a [`crate::service::task_launch::TaskInputSpec`]) and builds this layer
10//! via [`Self::new_from_fields`], layering it onto the spawner stack only
11//! when at least one field is present — mirroring the
12//! [`crate::middleware::project_name_alias::ProjectNameAliasMiddleware`] /
13//! [`crate::middleware::worker_binding::WorkerBindingMiddleware`]
14//! conditional-layering convention. Resolving `task_input` itself — sibling
15//! body field first, falling back to the legacy shape where these three
16//! keys were nested directly inside `init_ctx` — happens once at the wire
17//! boundary (`mlua-swarm-server`'s `run_flow_form`), not here; see
18//! [`Self::from_init_ctx`] for the now-deprecated constructor that used to
19//! do that extraction inline.
20//!
21//! Downstream Operator / Spawner code (for example `mlua-swarm-server`'s
22//! `Operator::execute`) reads the injected keys back via
23//! `ctx.meta.runtime.get(...)` the same way it reads `worker_binding` /
24//! `project_name_alias` — splicing them into the SubAgent's Spawn
25//! directive prompt is a downstream concern, out of scope here.
26//!
27//! # Pattern
28//!
29//! Same shape as `ProjectNameAliasMiddleware`: a task-wide (not per-agent)
30//! value set once at launch time and inserted into every spawn's `ctx`,
31//! unconditionally (no `ctx.agent` lookup — unlike `WorkerBindingMiddleware`,
32//! which is keyed by agent name).
33
34use crate::core::ctx::Ctx;
35use crate::core::engine::Engine;
36use crate::middleware::SpawnerLayer;
37use crate::types::{CapToken, StepId};
38use crate::worker::adapter::{SpawnError, SpawnerAdapter};
39use crate::worker::Worker;
40use async_trait::async_trait;
41use serde_json::Value;
42use std::sync::Arc;
43
44/// Key under `ctx.meta.runtime` that carries the project root path.
45///
46/// GH #20: canonical definition moved to
47/// [`crate::core::agent_context::TASK_PROJECT_ROOT_KEY`]; re-exported here
48/// so existing import paths (`crate::middleware::task_input::TASK_PROJECT_ROOT_KEY`)
49/// stay valid.
50pub use crate::core::agent_context::TASK_PROJECT_ROOT_KEY;
51
52/// Key under `ctx.meta.runtime` that carries the work dir path.
53///
54/// GH #20: canonical definition moved to
55/// [`crate::core::agent_context::TASK_WORK_DIR_KEY`]; re-exported here so
56/// existing import paths stay valid.
57pub use crate::core::agent_context::TASK_WORK_DIR_KEY;
58
59/// Key under `ctx.meta.runtime` that carries the free-form task metadata
60/// object.
61///
62/// GH #20: canonical definition moved to
63/// [`crate::core::agent_context::TASK_METADATA_KEY`]; re-exported here so
64/// existing import paths stay valid.
65pub use crate::core::agent_context::TASK_METADATA_KEY;
66
67/// `SpawnerLayer` that drops task-level execution context (`project_root` /
68/// `work_dir` / `task_metadata`) into `ctx` just before spawn.
69///
70/// Each field is independent: any subset may be `Some` (see
71/// [`Self::from_init_ctx`] for how they are extracted from `init_ctx`).
72/// Absent fields insert no key at all — no empty-string / `Value::Null`
73/// placeholder — so downstream `.get(...)` misses cleanly instead of
74/// observing a hollow value.
75pub struct TaskInputMiddleware {
76    project_root: Option<String>,
77    work_dir: Option<String>,
78    task_metadata: Option<Value>,
79}
80
81impl TaskInputMiddleware {
82    /// Builds a layer from already-resolved field values. Prefer
83    /// [`Self::new_from_fields`] (or [`Self::from_init_ctx`] when the
84    /// source is a raw `init_ctx` body) when the "all three absent → no
85    /// layer" contract matters to the caller — this constructor always
86    /// returns a (possibly no-op) `Self`, never `None`.
87    pub fn new(
88        project_root: Option<String>,
89        work_dir: Option<String>,
90        task_metadata: Option<Value>,
91    ) -> Self {
92        Self {
93            project_root,
94            work_dir,
95            task_metadata,
96        }
97    }
98
99    /// Issue #19 ST2 preferred constructor: builds from already-resolved
100    /// canonical Task-level field values (the wire boundary — e.g.
101    /// `mlua-swarm-server`'s `run_flow_form` — resolves sibling body
102    /// fields, falling back to the legacy `init_ctx`-nested shape only
103    /// there) rather than pulling them back out of `init_ctx` itself.
104    /// `init_ctx` stays a pure flow-ir eval seed with no promoted keys
105    /// folded back in.
106    ///
107    /// Returns `None` when all three fields are absent — same no-op
108    /// contract as [`Self::from_init_ctx`], so callers can use
109    /// `Option::map`/`match` uniformly regardless of which constructor
110    /// produced the layer.
111    pub fn new_from_fields(
112        project_root: Option<String>,
113        work_dir: Option<String>,
114        task_metadata: Option<Value>,
115    ) -> Option<Self> {
116        if project_root.is_none() && work_dir.is_none() && task_metadata.is_none() {
117            return None;
118        }
119        Some(Self::new(project_root, work_dir, task_metadata))
120    }
121
122    /// Extracts `project_root` (string) / `work_dir` (string) /
123    /// `task_metadata` (object) from a top-level `init_ctx` object, each
124    /// independently optional.
125    ///
126    /// Returns `None` when `init_ctx` is not a JSON object, or is an object
127    /// with none of the three keys present in the expected shape (a
128    /// present-but-wrong-typed value, e.g. `"work_dir": 42`, is treated the
129    /// same as absent — this is a best-effort task-level convenience
130    /// injection, not a request-body validator; malformed request bodies
131    /// are the request layer's concern). Callers only layer the returned
132    /// middleware onto the spawner stack when this is `Some`, keeping the
133    /// no-op path identical to today's behavior.
134    ///
135    /// Deprecated (issue #19 ST2): `TaskLaunchService::launch` no longer
136    /// calls this — it now reads [`crate::service::task_launch::TaskLaunchInput::task_input`]
137    /// (built via [`Self::new_from_fields`]) instead of extracting from
138    /// `init_ctx`. Kept for callers still relying on the nested-in-`init_ctx`
139    /// shape being pulled directly out of an arbitrary `Value` (not
140    /// currently exercised inside this crate, but a public, documented
141    /// constructor — removing it outright would be a breaking API change
142    /// for downstream users of this middleware in isolation).
143    #[deprecated(
144        note = "issue #19 ST2: prefer new_from_fields (or resolve the wire's legacy init_ctx-nested shape at the wire boundary and pass the result through TaskLaunchInput.task_input instead)"
145    )]
146    pub fn from_init_ctx(init_ctx: &Value) -> Option<Self> {
147        let obj = init_ctx.as_object()?;
148        let project_root = obj
149            .get(TASK_PROJECT_ROOT_KEY)
150            .and_then(Value::as_str)
151            .map(str::to_string);
152        let work_dir = obj
153            .get(TASK_WORK_DIR_KEY)
154            .and_then(Value::as_str)
155            .map(str::to_string);
156        let task_metadata = obj
157            .get(TASK_METADATA_KEY)
158            .filter(|v| v.is_object())
159            .cloned();
160        if project_root.is_none() && work_dir.is_none() && task_metadata.is_none() {
161            return None;
162        }
163        Some(Self::new(project_root, work_dir, task_metadata))
164    }
165}
166
167impl SpawnerLayer for TaskInputMiddleware {
168    fn wrap(&self, inner: Arc<dyn SpawnerAdapter>) -> Arc<dyn SpawnerAdapter> {
169        Arc::new(TaskInputWrapped {
170            inner,
171            project_root: self.project_root.clone(),
172            work_dir: self.work_dir.clone(),
173            task_metadata: self.task_metadata.clone(),
174        })
175    }
176}
177
178struct TaskInputWrapped {
179    inner: Arc<dyn SpawnerAdapter>,
180    project_root: Option<String>,
181    work_dir: Option<String>,
182    task_metadata: Option<Value>,
183}
184
185#[async_trait]
186impl SpawnerAdapter for TaskInputWrapped {
187    async fn spawn(
188        &self,
189        engine: &Engine,
190        ctx: &Ctx,
191        task_id: StepId,
192        attempt: u32,
193        token: CapToken,
194    ) -> Result<Box<dyn Worker>, SpawnError> {
195        let mut new_ctx = ctx.clone();
196        if let Some(project_root) = &self.project_root {
197            new_ctx.meta.runtime.insert(
198                TASK_PROJECT_ROOT_KEY.to_string(),
199                Value::String(project_root.clone()),
200            );
201        }
202        if let Some(work_dir) = &self.work_dir {
203            new_ctx.meta.runtime.insert(
204                TASK_WORK_DIR_KEY.to_string(),
205                Value::String(work_dir.clone()),
206            );
207        }
208        if let Some(task_metadata) = &self.task_metadata {
209            new_ctx
210                .meta
211                .runtime
212                .insert(TASK_METADATA_KEY.to_string(), task_metadata.clone());
213        }
214        self.inner
215            .spawn(engine, &new_ctx, task_id, attempt, token)
216            .await
217    }
218}
219
220#[cfg(test)]
221mod tests {
222    use super::*;
223    use crate::core::config::EngineCfg;
224    use crate::types::Role;
225    use serde_json::json;
226    use std::sync::Mutex;
227    use std::time::Duration;
228
229    /// Inner spawner stub that records the `Ctx` it was called with and
230    /// fails the spawn (we only care about the ctx snapshot).
231    struct CtxProbe {
232        seen: Arc<Mutex<Option<Ctx>>>,
233    }
234
235    #[async_trait]
236    impl SpawnerAdapter for CtxProbe {
237        async fn spawn(
238            &self,
239            _engine: &Engine,
240            ctx: &Ctx,
241            _task_id: StepId,
242            _attempt: u32,
243            _token: CapToken,
244        ) -> Result<Box<dyn Worker>, SpawnError> {
245            *self.seen.lock().unwrap() = Some(ctx.clone());
246            Err(SpawnError::Internal("probe stop".into()))
247        }
248    }
249
250    fn probe_stack(
251        layer: TaskInputMiddleware,
252    ) -> (Arc<dyn SpawnerAdapter>, Arc<Mutex<Option<Ctx>>>) {
253        let seen = Arc::new(Mutex::new(None));
254        let inner = Arc::new(CtxProbe { seen: seen.clone() });
255        let wrapped = layer.wrap(inner);
256        (wrapped, seen)
257    }
258
259    #[tokio::test]
260    async fn injects_all_three_fields_into_ctx_meta_runtime() {
261        let layer = TaskInputMiddleware::new(
262            Some("/repo".to_string()),
263            Some("/repo/work".to_string()),
264            Some(json!({ "issue": 17 })),
265        );
266        let (stack, seen) = probe_stack(layer);
267        let engine = Engine::new(EngineCfg::default());
268        let task_id = StepId::parse("ST-1").unwrap();
269        let ctx = Ctx::new(task_id.clone(), 1, "planner");
270        let token = engine
271            .attach("ut-op", Role::Operator, Duration::from_secs(30))
272            .await
273            .expect("attach");
274        let _ = stack.spawn(&engine, &ctx, task_id, 1, token).await;
275
276        let observed = seen.lock().unwrap().clone().expect("inner ctx captured");
277        assert_eq!(
278            observed.meta.runtime.get(TASK_PROJECT_ROOT_KEY),
279            Some(&Value::String("/repo".to_string()))
280        );
281        assert_eq!(
282            observed.meta.runtime.get(TASK_WORK_DIR_KEY),
283            Some(&Value::String("/repo/work".to_string()))
284        );
285        assert_eq!(
286            observed.meta.runtime.get(TASK_METADATA_KEY),
287            Some(&json!({ "issue": 17 }))
288        );
289    }
290
291    #[tokio::test]
292    async fn partial_fields_only_insert_present_keys() {
293        let layer = TaskInputMiddleware::new(Some("/repo".to_string()), None, None);
294        let (stack, seen) = probe_stack(layer);
295        let engine = Engine::new(EngineCfg::default());
296        let task_id = StepId::parse("ST-2").unwrap();
297        let ctx = Ctx::new(task_id.clone(), 1, "planner");
298        let token = engine
299            .attach("ut-op", Role::Operator, Duration::from_secs(30))
300            .await
301            .expect("attach");
302        let _ = stack.spawn(&engine, &ctx, task_id, 1, token).await;
303
304        let observed = seen.lock().unwrap().clone().expect("inner ctx captured");
305        assert_eq!(
306            observed.meta.runtime.get(TASK_PROJECT_ROOT_KEY),
307            Some(&Value::String("/repo".to_string()))
308        );
309        assert!(
310            !observed.meta.runtime.contains_key(TASK_WORK_DIR_KEY),
311            "work_dir absent must not insert a key"
312        );
313        assert!(
314            !observed.meta.runtime.contains_key(TASK_METADATA_KEY),
315            "task_metadata absent must not insert a key"
316        );
317    }
318
319    #[tokio::test]
320    async fn empty_fields_layer_is_a_no_op() {
321        let layer = TaskInputMiddleware::new(None, None, None);
322        let (stack, seen) = probe_stack(layer);
323        let engine = Engine::new(EngineCfg::default());
324        let task_id = StepId::parse("ST-3").unwrap();
325        let ctx = Ctx::new(task_id.clone(), 1, "planner");
326        let token = engine
327            .attach("ut-op", Role::Operator, Duration::from_secs(30))
328            .await
329            .expect("attach");
330        let _ = stack.spawn(&engine, &ctx, task_id, 1, token).await;
331
332        let observed = seen.lock().unwrap().clone().expect("inner ctx captured");
333        assert!(observed.meta.runtime.is_empty());
334    }
335
336    #[test]
337    #[allow(deprecated)]
338    fn from_init_ctx_extracts_all_three_fields() {
339        let init_ctx = json!({
340            "project_root": "/repo",
341            "work_dir": "/repo/work",
342            "task_metadata": { "issue": 17 },
343            "other": "kept-as-is-by-flow-eval-not-this-layer",
344        });
345        let layer = TaskInputMiddleware::from_init_ctx(&init_ctx).expect("some fields present");
346        assert_eq!(layer.project_root.as_deref(), Some("/repo"));
347        assert_eq!(layer.work_dir.as_deref(), Some("/repo/work"));
348        assert_eq!(layer.task_metadata, Some(json!({ "issue": 17 })));
349    }
350
351    #[test]
352    #[allow(deprecated)]
353    fn from_init_ctx_partial_fields_present() {
354        let init_ctx = json!({ "project_root": "/repo" });
355        let layer = TaskInputMiddleware::from_init_ctx(&init_ctx).expect("one field present");
356        assert_eq!(layer.project_root.as_deref(), Some("/repo"));
357        assert!(layer.work_dir.is_none());
358        assert!(layer.task_metadata.is_none());
359    }
360
361    #[test]
362    #[allow(deprecated)]
363    fn from_init_ctx_no_recognized_fields_returns_none() {
364        let init_ctx = json!({ "input": "hi" });
365        assert!(TaskInputMiddleware::from_init_ctx(&init_ctx).is_none());
366    }
367
368    #[test]
369    #[allow(deprecated)]
370    fn from_init_ctx_non_object_init_ctx_returns_none() {
371        assert!(TaskInputMiddleware::from_init_ctx(&json!("just a string")).is_none());
372        assert!(TaskInputMiddleware::from_init_ctx(&json!(null)).is_none());
373    }
374
375    #[test]
376    #[allow(deprecated)]
377    fn from_init_ctx_wrong_typed_field_is_treated_as_absent() {
378        // work_dir as a number is not the documented `Object` /
379        // `String` shape — treated the same as absent rather than
380        // erroring (best-effort convenience injection, not a validator).
381        let init_ctx = json!({ "work_dir": 42, "task_metadata": "not-an-object" });
382        assert!(TaskInputMiddleware::from_init_ctx(&init_ctx).is_none());
383    }
384
385    // ──────────────────────────────────────────────────────────────────
386    // issue #19 ST2: `new_from_fields` (already-resolved fields, no
387    // `init_ctx` extraction — the direct-sibling-read replacement for
388    // `from_init_ctx`)
389    // ──────────────────────────────────────────────────────────────────
390
391    #[test]
392    fn new_from_fields_all_present_returns_some() {
393        let layer = TaskInputMiddleware::new_from_fields(
394            Some("/repo".to_string()),
395            Some("/repo/work".to_string()),
396            Some(json!({ "issue": 19 })),
397        )
398        .expect("all three present");
399        assert_eq!(layer.project_root.as_deref(), Some("/repo"));
400        assert_eq!(layer.work_dir.as_deref(), Some("/repo/work"));
401        assert_eq!(layer.task_metadata, Some(json!({ "issue": 19 })));
402    }
403
404    #[test]
405    fn new_from_fields_partial_present_returns_some() {
406        let layer = TaskInputMiddleware::new_from_fields(Some("/repo".to_string()), None, None)
407            .expect("one field present");
408        assert_eq!(layer.project_root.as_deref(), Some("/repo"));
409        assert!(layer.work_dir.is_none());
410        assert!(layer.task_metadata.is_none());
411    }
412
413    #[test]
414    fn new_from_fields_none_present_returns_none() {
415        assert!(TaskInputMiddleware::new_from_fields(None, None, None).is_none());
416    }
417}