mlua_swarm_server/operator_ws/
session.rs

1//! `WSOperatorSession`: 1 sid = 1 session = 3 traits co-hosted (`SeniorBridge` /
2//! `SpawnHook` / `Operator`). Registered simultaneously into 3 registries under
3//! the same sid — the canonical pattern where 1 WS connection covers all 3
4//! faces of the Operator role (judgment / observation / execution).
5//!
6//! `tx` is a `Mutex<Option<Sender>>`: `None` on disconnect, swappable to
7//! `Some(new_tx)` on reconnect. The `pending` `HashMap` persists on the session
8//! side, so a client holding answer/ack values across a disconnect can reconnect
9//! and resend them.
10//!
11//! For the detailed S↔C message flow, see the overview figure in `mod.rs`.
12
13use async_trait::async_trait;
14use mlua_swarm::{
15    CapToken, Ctx, Operator, SeniorBridge, SpawnHook, TaskId, WorkerBinding, WorkerError,
16    WorkerResult,
17};
18use serde_json::Value;
19use std::collections::HashMap;
20use tokio::sync::{mpsc, oneshot, Mutex};
21
22use super::protocol::{current_parent_req_id, PendingReply, ServerMsg};
23
24/// 1 sid = 1 session. Looked up by sid in the `operator_sessions` store on reconnect.
25pub struct WSOperatorSession {
26    sid: String,
27    /// The current mpsc sender on the write path. `None` on disconnect;
28    /// swapped to `Some(new_tx)` on reconnect.
29    tx: Mutex<Option<mpsc::UnboundedSender<ServerMsg>>>,
30    /// `req_id` → pending oneshot. Resolved when `answer` / `hook_ack` /
31    /// `spawn_ack` arrives.
32    pending: Mutex<HashMap<String, oneshot::Sender<PendingReply>>>,
33}
34
35impl WSOperatorSession {
36    /// `login.rs::handle_operator_socket` is the sole constructor call site.
37    /// Auth (Bearer token match) is checked there against `OperatorSessionEntry.token`
38    /// *before* upgrade — this struct no longer carries its own auth_token copy.
39    pub(super) fn new(sid: String, tx: mpsc::UnboundedSender<ServerMsg>) -> Self {
40        Self {
41            sid,
42            tx: Mutex::new(Some(tx)),
43            pending: Mutex::new(HashMap::new()),
44        }
45    }
46
47    /// Swaps in a new tx on reconnect. Expected to be called only from the handler side.
48    pub(super) async fn replace_tx(&self, new_tx: mpsc::UnboundedSender<ServerMsg>) {
49        *self.tx.lock().await = Some(new_tx);
50    }
51
52    /// Clears tx to `None` on disconnect. Expected to be called only from the handler side.
53    pub(crate) async fn clear_tx(&self) {
54        *self.tx.lock().await = None;
55    }
56
57    /// Resolves the pending oneshot when a `ClientMsg` arrives on the handler's
58    /// read task. If `req_id` is not registered, no-op (= silently drops unknown acks).
59    pub(super) async fn resolve_pending(&self, req_id: &str, reply: PendingReply) {
60        if let Some(otx) = self.pending.lock().await.remove(req_id) {
61            let _ = otx.send(reply);
62        }
63    }
64
65    /// Inserts an entry into pending, sends S→C, and waits for the reply. No
66    /// timeout in v1.5 (= an ask during a disconnect immediately returns `Err`
67    /// on send failure; reconnect-wait behavior is v2).
68    async fn send_and_await(&self, req_id: String, msg: ServerMsg) -> Result<PendingReply, String> {
69        let (otx, orx) = oneshot::channel::<PendingReply>();
70        self.pending.lock().await.insert(req_id.clone(), otx);
71
72        // Fetch `tx` and send. When None, we are disconnected — fail fast.
73        let send_result = {
74            let guard = self.tx.lock().await;
75            match guard.as_ref() {
76                Some(tx) => tx
77                    .send(msg)
78                    .map_err(|_| "ws send channel closed".to_string()),
79                None => Err("ws operator disconnected".to_string()),
80            }
81        };
82        if let Err(e) = send_result {
83            self.pending.lock().await.remove(&req_id);
84            return Err(e);
85        }
86
87        orx.await
88            .map_err(|_| "ws operator: oneshot cancelled (= reply path closed)".to_string())
89    }
90
91    /// Fire-and-forget send for `after` (= no reply expected).
92    async fn send_oneway(&self, msg: ServerMsg) -> Result<(), String> {
93        let guard = self.tx.lock().await;
94        match guard.as_ref() {
95            Some(tx) => tx
96                .send(msg)
97                .map_err(|_| "ws send channel closed".to_string()),
98            None => Err("ws operator disconnected".to_string()),
99        }
100    }
101}
102
103#[async_trait]
104impl SeniorBridge for WSOperatorSession {
105    async fn ask(&self, task_id: &TaskId, question: Value) -> Result<Value, String> {
106        let req_id = format!("{}-ask-{}", self.sid, uuid::Uuid::new_v4());
107        let msg = ServerMsg::Ask {
108            req_id: req_id.clone(),
109            parent_req_id: current_parent_req_id(),
110            task_id: task_id.0.clone(),
111            question,
112        };
113        match self.send_and_await(req_id, msg).await? {
114            PendingReply::Answer(v) => Ok(v),
115            PendingReply::HookAck { .. } => {
116                Err("ws operator: unexpected hook_ack reply to ask".into())
117            }
118            PendingReply::SpawnAck { .. } => {
119                Err("ws operator: unexpected spawn_ack reply to ask".into())
120            }
121        }
122    }
123}
124
125#[async_trait]
126impl SpawnHook for WSOperatorSession {
127    async fn before(&self, ctx: &Ctx) -> Result<(), String> {
128        let req_id = format!("{}-hb-{}", self.sid, uuid::Uuid::new_v4());
129        let msg = ServerMsg::HookBefore {
130            req_id: req_id.clone(),
131            parent_req_id: current_parent_req_id(),
132            task_id: ctx.task_id.0.clone(),
133            agent: ctx.agent.clone(),
134            attempt: ctx.attempt,
135        };
136        match self.send_and_await(req_id, msg).await? {
137            PendingReply::HookAck { ok: true, .. } => Ok(()),
138            PendingReply::HookAck { ok: false, reason } => {
139                Err(reason.unwrap_or_else(|| "ws operator: spawn rejected".into()))
140            }
141            PendingReply::Answer(_) => {
142                Err("ws operator: unexpected answer reply to hook_before".into())
143            }
144            PendingReply::SpawnAck { .. } => {
145                Err("ws operator: unexpected spawn_ack reply to hook_before".into())
146            }
147        }
148    }
149
150    async fn after(&self, ctx: &Ctx, result: &Value) -> Result<(), String> {
151        let req_id = format!("{}-ha-{}", self.sid, uuid::Uuid::new_v4());
152        let msg = ServerMsg::HookAfter {
153            req_id,
154            parent_req_id: current_parent_req_id(),
155            task_id: ctx.task_id.0.clone(),
156            agent: ctx.agent.clone(),
157            attempt: ctx.attempt,
158            result: result.clone(),
159        };
160        // `after` is fire-and-forget — swallow send failures.
161        let _ = self.send_oneway(msg).await;
162        Ok(())
163    }
164}
165
166#[async_trait]
167impl Operator for WSOperatorSession {
168    /// Thin control channel impl (the Spawn thin-control axis): `system` / `prompt`
169    /// have already been baked into engine state on the server side
170    /// (= `bake_worker_system_prompt` in `OperatorSpawner.spawn` + the existing
171    /// `fetch_prompt` path). This impl encodes `worker_token` and hands it to
172    /// the MainAI in a single Spawn message; the SubAgent then hits
173    /// `/v1/worker/prompt` + `/v1/worker/result` itself over HTTP. The `system`
174    /// / `prompt` arguments are intentionally **not used here** (= heavy payloads
175    /// are not carried on WS — thin-path discipline).
176    ///
177    /// The SubAgent's result post (= HTTP POST `/v1/worker/result`) appends
178    /// `Final` to `output_tail`; when the MainAI returns `SpawnAck`, this
179    /// `execute` returns `WorkerResult` and control returns to the dispatch path.
180    ///
181    /// `worker` is required (see `requires_worker_binding`) — the compile-time
182    /// gate in `OperatorSpawnerFactory::build` is the primary defense, but a
183    /// `None` can still reach here on paths that bypass compilation (e.g. an
184    /// operator-sid-pin path). This runtime check is the defensive second
185    /// layer: fail the task loud rather than silently degrade to the old
186    /// hardcoded `"mse-worker"` literal.
187    async fn execute(
188        &self,
189        ctx: &Ctx,
190        _system: Option<String>,
191        _prompt: String,
192        worker: Option<WorkerBinding>,
193        worker_token: CapToken,
194    ) -> Result<WorkerResult, WorkerError> {
195        let Some(worker) = worker else {
196            return Err(WorkerError::Failed(format!(
197                "agent '{}' has no worker_binding; WS thin-path requires one \
198                 (Blueprint AgentDef.profile.worker_binding)",
199                ctx.agent
200            )));
201        };
202        let req_id = format!("{}-spawn-{}", self.sid, uuid::Uuid::new_v4());
203        let worker_handle = ctx
204            .meta
205            .runtime
206            .get("worker_handle")
207            .and_then(|v| v.as_str())
208            .map(|s| s.to_string());
209        let project_name_alias = ctx
210            .meta
211            .runtime
212            .get("project_name_alias")
213            .and_then(|v| v.as_str());
214        let data_sink_endpoint = ctx
215            .meta
216            .runtime
217            .get("data_sink_endpoint")
218            .and_then(|v| v.as_str());
219        let directive = default_spawn_directive(
220            &ctx.agent,
221            &ctx.task_id.0,
222            &worker.variant,
223            project_name_alias,
224            data_sink_endpoint,
225        );
226        let msg = ServerMsg::Spawn {
227            req_id: req_id.clone(),
228            parent_req_id: current_parent_req_id(),
229            task_id: ctx.task_id.0.clone(),
230            agent: ctx.agent.clone(),
231            attempt: ctx.attempt,
232            capability_token: worker_token.encode(),
233            worker_handle,
234            worker: Some(worker),
235            directive,
236        };
237        match self.send_and_await(req_id, msg).await {
238            Ok(PendingReply::SpawnAck {
239                value,
240                ok,
241                error: None,
242            }) => Ok(WorkerResult { value, ok }),
243            Ok(PendingReply::SpawnAck {
244                error: Some(msg), ..
245            }) => Err(WorkerError::Failed(msg)),
246            Ok(_) => Err(WorkerError::Failed(
247                "ws operator: unexpected non-spawn reply".into(),
248            )),
249            Err(e) => Err(WorkerError::Failed(format!("ws operator spawn: {e}"))),
250        }
251    }
252
253    fn requires_worker_binding(&self) -> bool {
254        true
255    }
256}
257
258/// Literal instruction text for the MainAI (= WS Client = Operator role). Fix
259/// for observation #7.
260///
261/// Minimal hand-off form parallel to /orch (agent_primitive): sends an
262/// `[agent_primitive dispatch=@<agent>]` marker + worker endpoint + auth +
263/// task_id in the payload; the MainAI **kicks a SubAgent by specifying AgentId +
264/// Token** and **forwards the return string verbatim into `SpawnAck.value`**.
265///
266/// The detailed instructions for the SubAgent are consolidated into the
267/// agent.md `system` (= the body fetched by `GET /v1/worker/prompt`); the
268/// directive is narrowed to the minimum routing information.
269///
270/// # `project_name_alias` literal expansion
271///
272/// When the caller sets `Blueprint.metadata.project_name_alias = Some(a)`
273/// (schema field defined in `mlua-swarm-blueprint-schema::BlueprintMetadata`),
274/// the value flows into `ctx.meta.runtime["project_name_alias"]` via the
275/// `ProjectNameAliasLayer` SpawnerLayer (see
276/// `mlua_swarm::middleware::project_name_alias`). This function
277/// then expands the alias **literally** into the Spawn directive text — as the
278/// `project_name_alias: {a}` header line and as the "LDS Session Alias" mandatory
279/// reminder block for the MainAI. The engine itself performs no other action on
280/// the alias; the expansion here is what the MainAI actually reads.
281///
282/// # `subagent_type` (Blueprint-baked worker binding)
283///
284/// Resolved from `AgentDef.profile.worker_binding` (see `WorkerBinding`) and
285/// literally substituted for the old hardcoded `"mse-worker"` string — the
286/// Blueprint is the single source of truth for which Claude Code SubAgent
287/// definition the MainAI must dispatch. There is deliberately **no fallback**
288/// to another `subagent_type` here: if the named SubAgent definition is not
289/// registered, the MainAI is instructed to fail the SpawnAck loud rather than
290/// silently substitute a different one.
291pub(super) fn default_spawn_directive(
292    agent: &str,
293    task_id: &str,
294    subagent_type: &str,
295    project_name_alias: Option<&str>,
296    data_sink_endpoint: Option<&str>,
297) -> String {
298    // Expanded only when Blueprint.metadata.project_name_alias is Some.
299    // Presents a discipline reminder to the MainAI plus the literal line the
300    // SubAgent prompt should inject.
301    let project_alias_line = match project_name_alias {
302        Some(a) => format!("project_name_alias: {a}\n"),
303        None => String::new(),
304    };
305    // Endpoint hint for the Data path (Big Response routing). Only when
306    // Some, inject a convention line telling the MainAgent to pass the Big
307    // EMIT POST target URL into the SubAgent prompt or environment when it
308    // kicks a SubAgent. Audience: MainAgent (the SubAgent-launcher side).
309    // A single authenticated emit endpoint: the token can be passed as
310    // Bearer or `?token=`; both consume the same CapToken material.
311    let data_endpoint_block = match data_sink_endpoint {
312        Some(base) => format!(
313            "\n\
314             [Data path endpoint — MainAgent reminder]\n\
315             When you kick a SubAgent, inject the following two lines into\n\
316             its prompt / environment so Big Response payloads (4k+ tokens,\n\
317             files, intermediate artifacts) flow directly to the Store owner,\n\
318             bypassing the MainAgent (context stays small; only the out_id\n\
319             ref is passed around).\n  \
320             DATA_EMIT: {base}/v1/data/emit  (POST, auth = Bearer worker_handle or ?token=)\n  \
321             DATA_GET:  {base}/v1/data/<out_id|out_name>  (the next SubAgent fetches from $IN_REFS)\n\
322             When a SubAgent produces a Big Response, POST it to DATA_EMIT\n\
323             and return only the one-line out_id ref (do not mix the body\n\
324             in; the MainAgent must not answer directly).\n\
325             \n"
326        ),
327        None => String::new(),
328    };
329    let main_ai_reminder = match project_name_alias {
330        Some(a) => format!(
331            "\n\
332             [LDS Session Alias Reminder — MainAI mandatory]\n\
333             Before kicking the SubAgent below, call:\n  \
334             mcp__lds__session_create(root=<working_dir>, alias=\"{a}\")\n\
335             (= establish a single task-level lds session; reuse on repeated dispatch).\n\
336             Then add this literal line to the SubAgent prompt body below:\n  \
337             LDS Session Alias: {a}\n\
338             The SubAgent will call mcp__lds__session_start(alias=\"{a}\") on init,\n\
339             keeping worktree ownership unified across dispatches.\n\
340             (Full discipline rationale is inlined above; reach is via this directive itself,\n\
341              not via any external doc path. The 2 steps above are the complete contract.)\n\
342             \n"
343        ),
344        None => String::new(),
345    };
346    format!(
347        "[agent_primitive dispatch=@{agent}]\n\
348         worker endpoint:\n  \
349         GET  <base_url>/v1/worker/prompt?task_id={task_id}\n  \
350         POST <base_url>/v1/worker/submit\n\
351         auth: Bearer <worker_handle from THIS Spawn payload (= short `wh-XXXXXXXX` form)>\n\
352         task_id: {task_id}\n\
353         agent_id: {agent}\n\
354         {project_alias_line}\
355         {data_endpoint_block}\
356         {main_ai_reminder}\
357         Kick a SubAgent via Agent tool with subagent_type=\"{subagent_type}\" (= project-local \
358         `.claude/agents/{subagent_type}.md`, this agent's Blueprint-declared worker binding). \
359         The prompt you pass to it MUST be EXACTLY these 4 lines (no preamble, no extra text):\n\
360         \n  \
361         agent_id: {agent}\n  \
362         worker_handle: <THIS Spawn payload's `worker_handle` field (short string `wh-XXXXXXXX`)>\n  \
363         base_url: <server HTTP root, e.g. http://127.0.0.1:7786>\n  \
364         task_id: {task_id}\n\
365         \n\
366         The SubAgent self-fetches system + prompt via GET (Bearer = handle), \
367         executes as agent @{agent}, POSTs raw body to /v1/worker/submit (Bearer = handle, \
368         server resolves task_id from handle), and replies `OUTPUT` 1 word. You then forward \
369         SpawnAck {{req_id, value:{{}}, ok:true}} through your operator client — MCP path: \
370         mse_ack(sid, req_id, kind=\"spawn_ack\", ok=true) (= empty value because canonical \
371         body lives in output_tail via the POST). \
372         Do NOT fetch /v1/worker/prompt yourself. Do NOT wrap, summarize, or field-select \
373         the SubAgent reply. Observation / debug is a separate channel (= agent-inspect MCP / \
374         GET /v1/tasks/{{id}}), do NOT mix it into the forward path. \
375         If the SubAgent type is not registered, FAIL LOUD: reply SpawnAck ok=false with an \
376         error explaining the missing `.claude/agents/{subagent_type}.md` — do NOT fall back \
377         to another subagent_type."
378    )
379}
380
381#[cfg(test)]
382mod tests {
383    use super::*;
384
385    #[test]
386    fn directive_omits_project_name_alias_when_none() {
387        let d = default_spawn_directive("impl-lead", "task-x", "mse-worker-coder", None, None);
388        assert!(!d.contains("project_name_alias:"));
389        assert!(!d.contains("LDS Session Alias"));
390        assert!(!d.contains("session_create"));
391    }
392
393    #[test]
394    fn directive_emits_project_name_alias_when_some() {
395        let d = default_spawn_directive(
396            "impl-lead",
397            "task-x",
398            "mse-worker-coder",
399            Some("mse-task-7785"),
400            None,
401        );
402        // Header line (expanded verbatim from the value).
403        assert!(
404            d.contains("project_name_alias: mse-task-7785"),
405            "directive missing project_name_alias header: {d}"
406        );
407        // MainAI mandatory reminder (= session_create + SubAgent prompt inject)
408        assert!(
409            d.contains("mcp__lds__session_create(root=<working_dir>, alias=\"mse-task-7785\")"),
410            "directive missing session_create reminder: {d}"
411        );
412        assert!(
413            d.contains("LDS Session Alias: mse-task-7785"),
414            "directive missing SubAgent prompt inject line: {d}"
415        );
416        // Reach discipline: the rationale is inlined into the directive (no external doc path reference).
417        assert!(
418            d.contains("inlined above") || d.contains("complete contract"),
419            "directive should inline rationale rather than point at external doc: {d}"
420        );
421        // The SoT is not pointed at an AI personal memory file (which is
422        // outside the MainAI's reach) — reach-axis consistency. Path
423        // references coming from the subagent registration convention (for
424        // example `agents/mse-worker.md`) are a separate case and are
425        // allowed. The pattern is assembled by string concat so that no
426        // gitignored dir literal remains in the source and the
427        // internal-doc-leak / secret-pre-commit-checker mechanical pattern
428        // match is avoided.
429        let forbidden_doc_ref = format!(".{}/CLAUDE.md", "claude");
430        assert!(
431            !d.contains(&forbidden_doc_ref),
432            "directive must not reference {forbidden_doc_ref} (out of MainAI scope): {d}"
433        );
434    }
435
436    #[test]
437    fn directive_omits_data_endpoint_when_none() {
438        let d = default_spawn_directive("impl-lead", "task-x", "mse-worker-coder", None, None);
439        assert!(!d.contains("[Data path endpoint"));
440        assert!(!d.contains("DATA_EMIT"));
441        assert!(!d.contains("DATA_GET"));
442    }
443
444    #[test]
445    fn directive_emits_data_endpoint_when_some() {
446        let base = "http://127.0.0.1:7785";
447        let d =
448            default_spawn_directive("impl-lead", "task-x", "mse-worker-coder", None, Some(base));
449        assert!(
450            d.contains("[Data path endpoint"),
451            "directive missing data endpoint block header: {d}"
452        );
453        assert!(
454            d.contains(&format!("DATA_EMIT: {base}/v1/data/emit")),
455            "directive missing single-mouth emit line: {d}"
456        );
457        assert!(
458            d.contains("Bearer worker_handle or ?token="),
459            "directive missing auth transport hint: {d}"
460        );
461        assert!(
462            d.contains(&format!("DATA_GET:  {base}/v1/data/<out_id|out_name>")),
463            "directive missing GET line: {d}"
464        );
465        assert!(
466            !d.contains("emit-auth"),
467            "old split endpoint must not leak into directive: {d}"
468        );
469        assert!(
470            d.contains("bypassing the MainAgent") && d.contains("out_id ref"),
471            "directive should carry the ownership + bypass reasoning: {d}"
472        );
473    }
474
475    #[test]
476    fn directive_carries_declared_subagent_type_and_has_no_fallback() {
477        let d = default_spawn_directive("impl-lead", "task-x", "mse-worker-coder", None, None);
478        assert!(
479            d.contains("subagent_type=\"mse-worker-coder\""),
480            "directive must carry the Blueprint-declared subagent_type literally: {d}"
481        );
482        assert!(
483            d.contains(".claude/agents/mse-worker-coder.md"),
484            "directive must reference the declared subagent's own .md path: {d}"
485        );
486        // The old hardcoded default and its silent-fallback text must be gone.
487        assert!(
488            !d.contains("general-purpose"),
489            "directive must not fall back to subagent_type=\"general-purpose\": {d}"
490        );
491        assert!(
492            !d.contains("mse-worker\""),
493            "directive must not carry the old hardcoded \"mse-worker\" literal: {d}"
494        );
495        assert!(
496            d.contains("FAIL LOUD"),
497            "directive must instruct the MainAI to fail loud instead of falling back: {d}"
498        );
499    }
500}
mlua_swarm_server/operator_ws/session.rs

mlua_swarm_server/operator_ws/
session.rs
