mlua_swarm/middleware/

worker_binding.rs

//! `WorkerBindingMiddleware` — a `SpawnerLayer` that propagates the
//! Blueprint-baked per-agent [`WorkerBinding`] through `Ctx.meta.runtime`.
//!
//! `service::task_launch` builds an `agent name → WorkerBinding` map from
//! `Blueprint.agents[].profile.worker_binding` at launch time and places
//! this layer on the stack (outermost, next to `ProjectNameAliasMiddleware`).
//! Just before spawn it looks the map up by `ctx.agent` and, on a hit,
//! inserts the serialized binding into `Ctx.meta.runtime` under the
//! `worker_binding` key.
//!
//! Agents with no declared binding get no entry; the WS thin-path
//! `requires_worker_binding` fail-loud stays the safety net.
//!
//! # What reads the key: nothing in this repository
//!
//! Said plainly, because a reader deciding whether to keep this layer on
//! their stack needs it before the rationale. This layer was introduced
//! *for* the delegate axis (`b5eb000`, "wire worker binding through the
//! delegate axis"): `OperatorDelegateMiddleware` had no per-agent
//! `OperatorSpawner` to carry a compile-time-baked binding, so it read the
//! binding back out of `Ctx.meta.runtime` at spawn. That axis is gone, and
//! the AgentSpec axis needs no such round trip — `OperatorSpawner` holds
//! the baked `WorkerBinding` and hands it to `Operator::execute` as an
//! argument.
//!
//! With that reader deleted, a repo-wide search for [`WORKER_BINDING_KEY`]
//! and for `ctx.meta.runtime`'s `"worker_binding"` entry finds this
//! module's own tests and nothing else. In particular
//! `AgentContextView::from_ctx` does not read it (it reads the project /
//! work-dir / metadata / run-id / alias keys), and `mlua-swarm-server`'s
//! WS session takes the binding from its `execute` argument.
//!
//! Two things the layer still does, and they are the whole of why it stays
//! wired:
//!
//! - It is a **published extension point**. This type and the key are
//!   `pub` on a library crate, and "downstream Operator / Spawner code
//!   reads the injected keys back via `ctx.meta.runtime.get(...)`" is this
//!   crate's stated convention for out-of-tree implementations —
//!   [`crate::middleware::task_input`]'s module doc names `worker_binding`
//!   as one of those keys. Having no in-tree reader is not the same as
//!   having no reader.
//! - The value lands in the `Ctx` snapshot `ReplayStore` persists for each
//!   passed step, so a stored run records the binding every spawn
//!   declared.
//!
//! Neither is load-bearing, and this doc does not pretend otherwise. If
//! the extension point is not wanted, delete the layer, its conditional
//! wiring in `service::task_launch::TaskLaunchService::launch`, and this
//! key — rather than leaving a claim of a live reader standing. That is a
//! breaking change (both items are `pub`), and it also needs the prose
//! references in [`crate::middleware::task_input`],
//! `crate::middleware::agent_context` and `crate::core::agent_context`
//! updated, one of which is an intra-doc link that would otherwise fail
//! the doc build.
//!
//! Same shape as `ProjectNameAliasMiddleware` / `CompiledAgentTable`: a
//! compile/launch-time table keyed by agent name, looked up via
//! `ctx.agent` at spawn time, no engine state touched.

use crate::core::ctx::Ctx;
use crate::core::engine::Engine;
use crate::middleware::SpawnerLayer;
use crate::operator::WorkerBinding;
use crate::types::{CapToken, StepId};
use crate::worker::adapter::{SpawnError, SpawnerAdapter};
use crate::worker::Worker;
use async_trait::async_trait;
use std::collections::HashMap;
use std::sync::Arc;

/// Key under `ctx.meta.runtime` this layer writes the serialized
/// [`WorkerBinding`] to.
///
/// `pub` because reading it back *is* the extension point. The delegate
/// axis used to be the in-tree reader; see the module doc for what is left
/// after its removal (nothing in this repository).
pub const WORKER_BINDING_KEY: &str = "worker_binding";

/// `SpawnerLayer` that drops the per-agent binding into `ctx` just before
/// spawn.
pub struct WorkerBindingMiddleware {
    bindings: Arc<HashMap<String, WorkerBinding>>,
}

impl WorkerBindingMiddleware {
    /// Wraps an `agent name → WorkerBinding` map to inject on every spawn
    /// whose `ctx.agent` has an entry.
    pub fn new(bindings: HashMap<String, WorkerBinding>) -> Self {
        Self {
            bindings: Arc::new(bindings),
        }
    }
}

impl SpawnerLayer for WorkerBindingMiddleware {
    fn wrap(&self, inner: Arc<dyn SpawnerAdapter>) -> Arc<dyn SpawnerAdapter> {
        Arc::new(WorkerBindingWrapped {
            inner,
            bindings: self.bindings.clone(),
        })
    }
}

struct WorkerBindingWrapped {
    inner: Arc<dyn SpawnerAdapter>,
    bindings: Arc<HashMap<String, WorkerBinding>>,
}

#[async_trait]
impl SpawnerAdapter for WorkerBindingWrapped {
    async fn spawn(
        &self,
        engine: &Engine,
        ctx: &Ctx,
        task_id: StepId,
        attempt: u32,
        token: CapToken,
    ) -> Result<Box<dyn Worker>, SpawnError> {
        let Some(binding) = self.bindings.get(&ctx.agent) else {
            // No declared binding for this agent — pass through untouched;
            // binding-requiring backends fail loud downstream.
            return self.inner.spawn(engine, ctx, task_id, attempt, token).await;
        };
        let value = serde_json::to_value(binding).map_err(|e| {
            SpawnError::Internal(format!(
                "worker_binding for agent '{}' failed to serialize: {e}",
                ctx.agent
            ))
        })?;
        let mut new_ctx = ctx.clone();
        new_ctx
            .meta
            .runtime
            .insert(WORKER_BINDING_KEY.to_string(), value);
        self.inner
            .spawn(engine, &new_ctx, task_id, attempt, token)
            .await
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::core::config::EngineCfg;
    use crate::types::Role;
    use std::sync::Mutex;
    use std::time::Duration;

    /// Inner spawner stub that records the `Ctx` it was called with and
    /// fails the spawn (we only care about the ctx snapshot).
    struct CtxProbe {
        seen: Arc<Mutex<Option<Ctx>>>,
    }

    #[async_trait]
    impl SpawnerAdapter for CtxProbe {
        async fn spawn(
            &self,
            _engine: &Engine,
            ctx: &Ctx,
            _task_id: StepId,
            _attempt: u32,
            _token: CapToken,
        ) -> Result<Box<dyn Worker>, SpawnError> {
            *self.seen.lock().unwrap() = Some(ctx.clone());
            Err(SpawnError::Internal("probe stop".into()))
        }
    }

    fn probe_stack(
        bindings: HashMap<String, WorkerBinding>,
    ) -> (Arc<dyn SpawnerAdapter>, Arc<Mutex<Option<Ctx>>>) {
        let seen = Arc::new(Mutex::new(None));
        let inner = Arc::new(CtxProbe { seen: seen.clone() });
        let wrapped = WorkerBindingMiddleware::new(bindings).wrap(inner);
        (wrapped, seen)
    }

    #[tokio::test]
    async fn injects_binding_into_ctx_meta_runtime_on_hit() {
        let mut map = HashMap::new();
        map.insert(
            "planner".to_string(),
            WorkerBinding {
                variant: "knowledge-worker".to_string(),
                tools: vec!["Read".to_string()],
                request_digest: Some(
                    "sha256:1111111111111111111111111111111111111111111111111111111111111111"
                        .parse()
                        .unwrap(),
                ),
                requested_model: Some("claude-sonnet".to_string()),
            },
        );
        let (stack, seen) = probe_stack(map);
        let engine = Engine::new(EngineCfg::default());
        let task_id = StepId::parse("ST-1").unwrap();
        let ctx = Ctx::new(task_id.clone(), 1, "planner");
        let token = engine
            .attach("ut-op", Role::Operator, Duration::from_secs(30))
            .await
            .expect("attach");
        let _ = stack.spawn(&engine, &ctx, task_id, 1, token).await;

        let observed = seen.lock().unwrap().clone().expect("inner ctx captured");
        let v = observed
            .meta
            .runtime
            .get(WORKER_BINDING_KEY)
            .expect("worker_binding key present");
        let wb: WorkerBinding = serde_json::from_value(v.clone()).expect("round-trip");
        assert_eq!(wb.variant, "knowledge-worker");
        assert_eq!(wb.tools, vec!["Read".to_string()]);
        // The requesting side's self-check inputs survive the ctx.meta.runtime
        // round-trip so the Operator can correlate and compare its environment.
        assert!(wb
            .request_digest
            .as_ref()
            .expect("request_digest present")
            .as_str()
            .starts_with("sha256:"));
        assert_eq!(wb.requested_model.as_deref(), Some("claude-sonnet"));
    }

    /// Wire compatibility: a binding with both self-check fields `None` omits
    /// the keys entirely, so the serialized frame matches the pre-C2 shape.
    #[test]
    fn omits_self_check_fields_when_none() {
        let json = serde_json::to_value(WorkerBinding {
            variant: "v".to_string(),
            tools: vec![],
            request_digest: None,
            requested_model: None,
        })
        .expect("serialize");
        let obj = json.as_object().expect("object");
        assert!(!obj.contains_key("request_digest"));
        assert!(!obj.contains_key("requested_model"));
    }

    #[tokio::test]
    async fn passes_through_untouched_on_miss() {
        let (stack, seen) = probe_stack(HashMap::new());
        let engine = Engine::new(EngineCfg::default());
        let task_id = StepId::parse("ST-2").unwrap();
        let ctx = Ctx::new(task_id.clone(), 1, "unbound-agent");
        let token = engine
            .attach("ut-op", Role::Operator, Duration::from_secs(30))
            .await
            .expect("attach");
        let _ = stack.spawn(&engine, &ctx, task_id, 1, token).await;

        let observed = seen.lock().unwrap().clone().expect("inner ctx captured");
        assert!(
            !observed.meta.runtime.contains_key(WORKER_BINDING_KEY),
            "no binding entry must be injected on miss"
        );
    }
}