mlua-swarm 0.20.0

//! The second stage of the two-stage pipeline: `SpawnerAdapter`.
//!
//! From the engine's viewpoint there is only one trait,
//! `SpawnerAdapter`; its `spawn` returns `Box<dyn Worker>` (see
//! `crate::worker::Worker`). Worker shape is an implementation detail of
//! each spawner; the engine only touches Workers through three
//! operations — `id()` / `cancel_token()` / `join()`.
//!
//! The old `WorkerAdapter` trait and `InProcWorker` struct — which
//! assumed a three-stage `Spawner.spawn → WorkerAdapter → invoke`
//! pipeline — were removed on this turn. Nothing instantiated or
//! dispatched them (dead code), and the multi-invocation path from
//! was collapsed in the implementation anyway.
//! The interface is now consolidated into the new `trait Worker` in
//! `src/worker.rs`.

use crate::core::agent_context::AgentContextView;
use crate::core::ctx::Ctx;
use crate::core::engine::Engine;
use crate::types::{CapToken, StepId};
use crate::worker::Worker;
use async_trait::async_trait;
use serde_json::Value;
use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use thiserror::Error;

/// Errors that can occur while `SpawnerAdapter::spawn` is setting up a
/// worker, before the worker itself starts running.
#[derive(Debug, Error)]
pub enum SpawnError {
    /// No `WorkerFn` is registered for the requested agent name.
    #[error("worker not registered: {0}")]
    NotRegistered(String),
    /// A middleware layer vetoed the spawn (e.g. capability check, rate
    /// limit, policy gate).
    #[error("spawn rejected by middleware: {0}")]
    RejectedByMiddleware(String),
    /// Any other setup failure (e.g. `fetch_prompt` failed).
    #[error("internal: {0}")]
    Internal(String),
}

/// Errors surfaced once a worker is running, via `Worker::join`.
#[derive(Debug, Error)]
pub enum WorkerError {
    /// The worker fn itself returned an error.
    #[error("worker fn returned error: {0}")]
    Failed(String),
    /// The worker was cancelled through its `CancellationToken`.
    #[error("cancelled")]
    Cancelled,
}

/// The value a `WorkerFn` hands back on success, folded into an
/// `OutputEvent::Final` by the spawner.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct WorkerResult {
    /// The worker fn's output payload.
    pub value: Value,
    /// Whether the agent itself considers this a successful result
    /// (distinct from `Result::Err` — a worker fn can return `Ok(..)`
    /// with `ok: false` to signal an agent-level failure).
    pub ok: bool,
    /// Optional normalized per-attempt stats sidecar (token usage /
    /// model / num_turns / adapter-specific raw data), produced by the
    /// worker boundary that knows them (agent-block result captor,
    /// subprocess stdout normalization, …). The spawner's fold site
    /// forwards it to `Engine::record_worker_stats`; it never rides
    /// into `OutputEvent::Final` (the BP-chain value stays stats-free).
    /// `None` = no stats reported — every pre-stats worker fn is
    /// unaffected (`#[serde(default)]` keeps wire compat).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub stats: Option<crate::store::trace::WorkerStats>,
}

impl WorkerResult {
    /// Ensure `stats.worker_kind` is set — the design invariant that
    /// every dispatched step's terminal `StepEntry` carries at least
    /// its worker kind label ("rust_fn" / "lua" / "agent_block" /
    /// "subprocess" / "operator" / …), even when no LLM-shaped stats
    /// (usage / model / num_turns) apply. Idempotent: if a boundary
    /// already reported a `worker_kind` (agent_block / subprocess
    /// stats sidecar, operator ack), that value wins; otherwise the
    /// fold site's `kind` becomes the value. Called at every worker
    /// fold site (`InProcSpawner` spawn task, subprocess spawn task,
    /// `OperatorDelegateMiddleware`).
    pub fn ensure_worker_kind(mut self, kind: &str) -> Self {
        let stats = self
            .stats
            .get_or_insert_with(crate::store::trace::WorkerStats::default);
        if stats.worker_kind.is_none() {
            stats.worker_kind = Some(kind.to_string());
        }
        self
    }
}

/// First stage of the two-stage pipeline: builds a `Box<dyn Worker>` for
/// one attempt. Every concrete spawner (`InProcSpawner`, `ProcessSpawner`,
/// the Operator spawner) implements this; the engine only ever holds a
/// `Arc<dyn SpawnerAdapter>` and knows nothing about the Worker shape
/// behind it.
#[async_trait]
pub trait SpawnerAdapter: Send + Sync {
    /// Spawn one attempt as a worker. Returns `Box<dyn Worker>`.
    ///
    /// The `directive` argument was removed in design intent: prompts are
    /// pulled on demand through
    /// `engine.fetch_prompt(token, task_id, attempt)`. Spawners are free
    /// to use whatever protocol they like internally — push, pull, or a
    /// hybrid. `ProcessSpawner` runs `fetch_prompt` and pushes the
    /// result into the child's stdin; `InProcSpawner` injects a prep
    /// snapshot as `WorkerInvocation.prompt`; a child process could
    /// even re-pull with the token itself.
    async fn spawn(
        &self,
        engine: &Engine,
        ctx: &Ctx,
        task_id: StepId,
        attempt: u32,
        token: CapToken,
    ) -> Result<Box<dyn Worker>, SpawnError>;
}

// ─── InProcSpawner ────────────────────────────────────────────────────────

/// Invocation context handed to a Worker fn. Bundles `token` +
/// `task_id` + `prompt` + `sink` + `context`.
///
/// The `prompt` field was added in design intent, folding the old
/// `Fn(inv, directive)` `directive` argument into the invocation. The
/// spawner is expected to call
/// `engine.fetch_prompt(token, task_id, attempt)` in its prep step and
/// inject the snapshot into the invocation (push form). The `WorkerFn`
/// side may still re-pull if it needs to — for example to fetch the
/// prompt for a different attempt.
///
/// The `sink` field was added in design intent as the formal contract for
/// the spawner's intake surface. A worker fn can stream intermediate
/// events with things like
/// `inv.sink.emit(OutputEvent::Progress { .. })`. Child-process
/// spawners (`ProcessSpawner`, etc.) do not use `sink` — the child
/// speaks the stdout protocol; `InProcSpawner` injects one. Even
/// without `sink`, the `WorkerResult` returned by the fn is still
/// folded into a `Final` event on the spawner side, running alongside
/// the older return-value path.
/// `#[non_exhaustive]`: this struct is the in-process seam every backend
/// reads task context off, so it is expected to keep growing (GH #86 added
/// `context`). Marking it non-exhaustive makes each future field a
/// non-breaking addition. Construct it with [`WorkerInvocation::new`] plus
/// the `with_*` setters — the same shape `agent-block-core` moved
/// `BlockConfig` to, and for the same reason.
#[derive(Clone)]
#[non_exhaustive]
pub struct WorkerInvocation {
    /// Capability token authorizing this attempt.
    pub token: CapToken,
    /// The task this invocation belongs to.
    pub task_id: StepId,
    /// Attempt number within the task (used to key output events).
    pub attempt: u32,
    /// Registered agent name the `WorkerFn` was looked up under.
    pub agent: String,
    /// The prompt/prep snapshot pulled via `engine.fetch_prompt`,
    /// injected here (push form) so the worker fn does not need to call
    /// back into the engine for the common case.
    pub prompt: String,
    /// Intake: sink the worker fn uses to emit intermediate
    /// `OutputEvent`s. Injected by `InProcSpawner`. `None` means the
    /// sink path is not wired for this invocation.
    pub sink: Option<std::sync::Arc<dyn crate::worker::output::OutputSink>>,
    /// Upstream task cancel token — the clone of `cancel_inner`
    /// generated by `InProcSpawner` for `JoinHandleWorker`. Worker fns
    /// bridge this to their child futures or their SDK's
    /// `shutdown_token`, propagating external cancellation all the way
    /// down. `None` — like `sink` above — means the caller path is not
    /// carrying the cancel channel.
    pub cancel_token: Option<tokio_util::sync::CancellationToken>,
    /// The materialized, policy-applied task context for this attempt —
    /// the **in-process twin of [`crate::types::WorkerPayload::context`]**
    /// (which is how the same view reaches an out-of-process Operator over
    /// `GET /v1/worker/prompt`).
    ///
    /// This is the single seam through which task-level context reaches an
    /// in-process worker. `InProcSpawner::spawn` fills it once, from
    /// [`AgentContextView::materialized_or_from_ctx`], so a worker fn reads
    /// `inv.context` instead of hand-rolling its own `Ctx` peek — the
    /// duplication that previously left the Lua / RustFn workers with no
    /// context at all while each other backend re-derived its own subset.
    ///
    /// `None` means the caller path did not carry a `Ctx` (the same
    /// "not wired for this invocation" convention `sink` / `cancel_token`
    /// use above). It is never `None` on the `InProcSpawner` path.
    pub context: Option<AgentContextView>,
}

impl WorkerInvocation {
    /// The five fields every invocation must carry. The optional rails
    /// (`sink` / `cancel_token` / `context`) default to `None` and are
    /// added with the `with_*` setters below — `InProcSpawner::spawn`
    /// wires all three.
    pub fn new(
        token: CapToken,
        task_id: StepId,
        attempt: u32,
        agent: impl Into<String>,
        prompt: impl Into<String>,
    ) -> Self {
        Self {
            token,
            task_id,
            attempt,
            agent: agent.into(),
            prompt: prompt.into(),
            sink: None,
            cancel_token: None,
            context: None,
        }
    }

    /// Attach the intake sink (see the [`Self::sink`] field doc).
    pub fn with_sink(
        mut self,
        sink: std::sync::Arc<dyn crate::worker::output::OutputSink>,
    ) -> Self {
        self.sink = Some(sink);
        self
    }

    /// Attach the upstream cancel token (see the [`Self::cancel_token`]
    /// field doc).
    pub fn with_cancel_token(mut self, token: tokio_util::sync::CancellationToken) -> Self {
        self.cancel_token = Some(token);
        self
    }

    /// Attach the materialized task context (see the [`Self::context`]
    /// field doc).
    pub fn with_context(mut self, context: AgentContextView) -> Self {
        self.context = Some(context);
        self
    }
}

impl std::fmt::Debug for WorkerInvocation {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("WorkerInvocation")
            .field("token", &self.token)
            .field("task_id", &self.task_id)
            .field("attempt", &self.attempt)
            .field("agent", &self.agent)
            .field("prompt", &self.prompt)
            .field("sink", &self.sink.as_ref().map(|_| "<OutputSink>"))
            .field(
                "cancel_token",
                &self.cancel_token.as_ref().map(|_| "<CancellationToken>"),
            )
            .finish()
    }
}

/// A registered agent implementation: takes a `WorkerInvocation` and
/// resolves to a `WorkerResult` (or a `WorkerError`). Boxed as a
/// type-erased `Future` so heterogeneous agent implementations (async
/// fns, closures capturing state, etc.) can share one registry entry
/// type.
pub type WorkerFn = Arc<
    dyn Fn(
            WorkerInvocation,
        ) -> Pin<Box<dyn Future<Output = Result<WorkerResult, WorkerError>> + Send>>
        + Send
        + Sync,
>;

/// `agent`-string → `WorkerFn` registry. The generic parameter `W` pins
/// the per-kind Worker concrete type at the type level, so AgentBlock /
/// Lua / RustFn each produce their own Worker type through
/// `InProcSpawner<W>` and the type binding is preserved right up until
/// `SpawnerAdapter::spawn()` erases the return as `Box<dyn Worker>`.
/// `W` must be constructible from `WorkerJoinHandler` via `From` — i.e.
/// a newtype that embeds the async-signal handle.
pub struct InProcSpawner<W = crate::worker::MiddlewareWorker> {
    /// Agent name → implementation lookup table.
    pub registry: HashMap<String, WorkerFn>,
    _phantom: std::marker::PhantomData<W>,
}

// Inherent impl for the default W = MiddlewareWorker (so `InProcSpawner::new()`
// in existing tests picks this default).
impl InProcSpawner {
    /// Creates an empty registry, defaulting the Worker type to
    /// `MiddlewareWorker` (used by existing call sites and tests).
    pub fn new() -> Self {
        Self {
            registry: HashMap::new(),
            _phantom: std::marker::PhantomData,
        }
    }

    /// Registers a `WorkerFn`-shaped async closure under `agent`,
    /// overwriting any previous registration for the same name. Returns
    /// `&mut Self` for chained registration calls.
    pub fn register<F, Fut>(&mut self, agent: impl Into<String>, f: F) -> &mut Self
    where
        F: Fn(WorkerInvocation) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = Result<WorkerResult, WorkerError>> + Send + 'static,
    {
        let f = Arc::new(f);
        let wrapped: WorkerFn = Arc::new(move |inv| {
            let f = f.clone();
            Box::pin(f(inv))
        });
        self.registry.insert(agent.into(), wrapped);
        self
    }
}

// Generic typed impl (the factory.build path that constructs a per-kind Worker).
impl<W> InProcSpawner<W>
where
    W: Worker + From<crate::worker::WorkerJoinHandler> + Send + Sync + 'static,
{
    /// Creates an empty registry pinned to Worker type `W` (the
    /// `factory.build` path uses this to get a per-kind Worker out of
    /// `spawn()` instead of the default `MiddlewareWorker`).
    pub fn typed() -> Self {
        Self {
            registry: HashMap::new(),
            _phantom: std::marker::PhantomData,
        }
    }
}

impl Default for InProcSpawner {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl<W: Worker + From<crate::worker::WorkerJoinHandler> + Send + Sync + 'static> SpawnerAdapter
    for InProcSpawner<W>
{
    async fn spawn(
        &self,
        engine: &Engine,
        ctx: &Ctx,
        task_id: StepId,
        attempt: u32,
        token: CapToken,
    ) -> Result<Box<dyn Worker>, SpawnError> {
        let f = self
            .registry
            .get(&ctx.agent)
            .cloned()
            .ok_or_else(|| SpawnError::NotRegistered(ctx.agent.clone()))?;

        // design intent: prompts are pulled via engine.fetch_prompt (the directive argument is retired)
        let prompt = engine
            .fetch_prompt(&token, &task_id)
            .await
            .map_err(|e| SpawnError::Internal(format!("fetch_prompt: {e}")))?;
        // In-process WorkerInvocation consumes `prompt` as `String` (issue #18
        // boundary render): Value flows end-to-end through the engine, and is
        // stringified here for the RustFn / Lua worker.
        let prompt = crate::core::engine::render_directive_to_string(&prompt);

        let (tx, rx) = tokio::sync::oneshot::channel();
        let cancel = tokio_util::sync::CancellationToken::new();
        let cancel_inner = cancel.clone();
        let worker_id = crate::types::WorkerId::new();
        // issue #11: surface the minted WorkerId in the trace log.
        tracing::debug!(worker_id = %worker_id, step_id = %task_id, "worker spawned (rustfn)");
        // design intent: hand `engine` / `token` to the spawn task so it can emit
        // OutputEvent::Final via submit_output (side-by-side with the
        // WorkerResult oneshot path).
        let engine_for_emit = engine.clone();
        let token_for_emit = token.clone();
        let task_id_for_emit = task_id.clone();
        // Wire the receiving end by injecting an EngineSink into WorkerInvocation.sink.
        let sink = std::sync::Arc::new(crate::worker::output::EngineSink::new(
            engine.clone(),
            token.clone(),
            task_id.clone(),
            attempt,
        )) as std::sync::Arc<dyn crate::worker::output::OutputSink>;
        let inv = WorkerInvocation::new(token, task_id, attempt, ctx.agent.clone(), prompt)
            .with_sink(sink)
            .with_cancel_token(cancel_inner.clone())
            // The one place task-level context enters the in-process lane
            // (mirrors how `Engine::fetch_worker_payload` fills
            // `WorkerPayload.context` for the out-of-process lane). Reads
            // the policy-applied view `AgentContextMiddleware` stashed, and
            // degrades to the raw `Ctx` projection when that layer is not
            // on this spawner stack.
            .with_context(AgentContextView::materialized_or_from_ctx(ctx));

        tokio::spawn(async move {
            let result = tokio::select! {
                r = f(inv) => r,
                _ = cancel_inner.cancelled() => Err(WorkerError::Cancelled),
            };
            // Fold WorkerResult into OutputEvent::Final. Contract: one Final per attempt.
            //
            // `submit_output` can REJECT this write — the GH #51
            // completion-time verdict-contract check is embedded in it and
            // runs BEFORE the `output_tail` append (see
            // `Engine::verdict_contract_completion_check`). This used to be
            // a discarded `let _ = …`, which made a rejection completely
            // invisible on this lane: the worker still signalled success, so
            // `dispatch_attempt_with`'s Final-pull reported the bare
            // `no Final in output_tail` with no cause anywhere — not even a
            // log line, while the WS Operator lane has always logged one
            // (`crate::operator`'s fallback emit). Carry the rejection into
            // the completion signal instead, so the reason travels all the
            // way to `EngineError::DispatchFailed` and the author sees which
            // contract they violated rather than a missing-Final symptom.
            //
            // Only the PRE-write rejections escalate to a failed attempt.
            // `submit_output` also returns `Err` for post-write side
            // effects (a strict-`CheckPolicy` materialize failure, say),
            // and there the `Final` IS on the tail — the attempt has a
            // value, the dispatcher can complete it, and failing it here
            // would turn a fail-open projection miss into a dead step.
            // Those still log, so the discarded-error hole is closed for
            // both, but only the contract gate changes the outcome.
            let mut emit_rejection: Option<String> = None;
            if let Ok(wr) = &result {
                // Stats sidecar: forward boundary-reported stats to the
                // engine (drained by the dispatcher's outcome fold into
                // the terminal StepEntry). This single fold site covers
                // every InProc worker kind (RustFn / Lua / AgentBlock).
                if let Some(stats) = wr.stats.clone() {
                    engine_for_emit
                        .record_worker_stats(&task_id_for_emit, attempt, stats)
                        .await;
                }
                let ev = crate::worker::output::OutputEvent::Final {
                    content: crate::worker::output::ContentRef::Inline {
                        value: wr.value.clone(),
                    },
                    ok: wr.ok,
                };
                if let Err(e) = engine_for_emit
                    .submit_output(&token_for_emit, &task_id_for_emit, attempt, ev)
                    .await
                {
                    let blocks_the_final = matches!(
                        e,
                        crate::EngineError::VerdictValueRejected { .. }
                            | crate::EngineError::VerdictPartMissing { .. }
                    );
                    tracing::warn!(
                        step_id = %task_id_for_emit,
                        attempt,
                        error = %e,
                        blocks_the_final,
                        "in-process worker's Final submission returned an error"
                    );
                    if blocks_the_final {
                        emit_rejection = Some(e.to_string());
                    }
                }
            }
            let signal: Result<(), WorkerError> = match emit_rejection {
                Some(reason) => Err(WorkerError::Failed(format!(
                    "Final rejected before output_tail: {reason}"
                ))),
                None => result.map(|_| ()),
            };
            let _ = tx.send(signal);
        });

        let handler = crate::worker::WorkerJoinHandler {
            worker_id,
            cancel,
            completion: rx,
        };
        Ok(Box::new(W::from(handler)))
    }
}