rig_core/effect/

mod.rs

//! Serializable effect requests, handler descriptions, outcomes, and recording data.
//! [`HandlerKey`] identifies a destination; [`EffectKind`] and [`Outcome`] describe
//! the exchange. Host-defined operations use [`CustomEffect`].
//!
//! ```
//! use rig_core::effect::{EffectFamily, EffectRow, model_key};
//!
//! let mut row = EffectRow::new();
//! row.insert(model_key("primary"), EffectFamily::Completion);
//! assert_eq!(row.len(), 1);
//! ```

use std::{fmt, sync::Arc};

use serde::{Deserialize, Serialize};

use crate::{
    completion::{CompletionRequest, CompletionResponse, Message, ModelRef, ProviderCapabilities},
    embeddings::EmbeddingResponse,
    error::ErrorReport,
    id::ConversationId,
    operation::RerankRequest,
    rerank::RerankResponse,
    streaming::Transcript,
    tool::ToolResult,
    vector_store::request::{Filter, VectorSearchRequest},
    wasm_compat::WasmCompatSend,
};

/// The identity of one dispatch, minted by the dispatcher.
///
/// It is the correlation key between a bus-tap [`EffectRecord`], a hook
/// observation, and a host's own bookkeeping (a Bevy driver maps it to an
/// `Entity`). Ids are unique per dispatcher and strictly increasing.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
#[serde(transparent)]
pub struct EffectId(u64);

impl EffectId {
    /// Build an id from its raw value (a replayer restoring a log does this).
    pub const fn from_raw(raw: u64) -> Self {
        Self(raw)
    }

    /// The raw value.
    pub const fn as_u64(self) -> u64 {
        self.0
    }
}

impl fmt::Display for EffectId {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "effect:{}", self.0)
    }
}

/// String identifier for a registered handler, serializable without a live handle.
#[derive(Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct HandlerKey(Arc<str>);

impl HandlerKey {
    /// Build a key from any string-like value.
    pub fn new(key: impl Into<Arc<str>>) -> Self {
        Self(key.into())
    }

    /// The key as a string slice.
    pub fn as_str(&self) -> &str {
        &self.0
    }

    /// Parses the key into [`KeyParts`] without failure. Formatting the parts
    /// reproduces the original key.
    pub fn parts(&self) -> KeyParts {
        KeyParts::parse(&self.0)
    }
}

/// A key's parts, by the grammar every key the tree mints follows:
///
/// ```text
/// [<owner>/]<name>[#<generation>]     where <name> is [<kind>:]<label>
/// ```
///
/// `golden/tool:add#2` is owner `golden`, kind `tool`, label `add`,
/// generation `2`; `host/note` is owner `host`, no kind, label `note`;
/// `model:fast` is no owner, kind `model`, label `fast`. The owner is the
/// text before the first `/`, the kind the text before the first `:` of
/// what follows, the generation the digits after the last `#`; a label
/// holds anything else, so a key whose label contains one of those
/// separators is read as its grammar says, not as its author meant. The
/// string is the wire form; this is how a reader takes it apart without
/// a hand parser of its own.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct KeyParts {
    /// The program or host the key belongs to (`golden`, `host`).
    pub owner: Option<Arc<str>>,
    /// The family-shaped prefix of the name (`model`, `tool`, `retrieve`).
    pub kind: Option<Arc<str>>,
    /// The name proper: a model label, a tool name, `memory`, `note`.
    pub label: Arc<str>,
    /// The disambiguating generation (`#2`): the same tool registered
    /// again under a fresh key. Only canonical `u64` decimal suffixes are
    /// generations; leading zeroes and overflow remain literal label text.
    pub generation: Option<u64>,
}

impl KeyParts {
    /// Read `key` by the grammar.
    pub fn parse(key: &str) -> Self {
        let (owner, rest) = match key.split_once('/') {
            Some((owner, rest)) if !owner.is_empty() && !rest.is_empty() => {
                (Some(Arc::from(owner)), rest)
            }
            _ => (None, key),
        };
        let (rest, generation) = match rest.rsplit_once('#') {
            Some((head, digits))
                if !digits.is_empty()
                    && (digits.len() == 1 || !digits.starts_with('0'))
                    && digits.bytes().all(|b| b.is_ascii_digit()) =>
            {
                match digits.parse::<u64>() {
                    Ok(generation) => (head, Some(generation)),
                    Err(_) => (rest, None),
                }
            }
            _ => (rest, None),
        };
        let (kind, label) = match rest.split_once(':') {
            Some((kind, label)) if !kind.is_empty() && !label.is_empty() => {
                (Some(Arc::from(kind)), label)
            }
            _ => (None, rest),
        };
        Self {
            owner,
            kind,
            label: Arc::from(label),
            generation,
        }
    }

    /// The parts as the key they came from (`Display` writes the same).
    pub fn to_key(&self) -> HandlerKey {
        HandlerKey::from(self.to_string())
    }
}

impl fmt::Display for KeyParts {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        if let Some(owner) = &self.owner {
            write!(f, "{owner}/")?;
        }
        if let Some(kind) = &self.kind {
            write!(f, "{kind}:")?;
        }
        f.write_str(&self.label)?;
        if let Some(generation) = self.generation {
            write!(f, "#{generation}")?;
        }
        Ok(())
    }
}

impl fmt::Debug for HandlerKey {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "HandlerKey({:?})", &*self.0)
    }
}

impl fmt::Display for HandlerKey {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(&self.0)
    }
}

impl From<&str> for HandlerKey {
    fn from(key: &str) -> Self {
        Self(Arc::from(key))
    }
}

impl From<String> for HandlerKey {
    fn from(key: String) -> Self {
        Self(Arc::from(key))
    }
}

impl AsRef<str> for HandlerKey {
    fn as_ref(&self) -> &str {
        &self.0
    }
}

// Serialize as a string without requiring serde's rc feature.
impl Serialize for HandlerKey {
    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        serializer.serialize_str(&self.0)
    }
}

impl<'de> Deserialize<'de> for HandlerKey {
    fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
        let key = <std::borrow::Cow<'de, str>>::deserialize(deserializer)?;
        Ok(Self(Arc::from(&*key)))
    }
}

/// Serde for a bare `Arc<str>` field (the `Custom` kind label).
mod arc_str {
    use std::sync::Arc;

    use serde::{Deserialize, Deserializer, Serializer};

    pub(super) fn serialize<S: Serializer>(
        value: &Arc<str>,
        serializer: S,
    ) -> Result<S::Ok, S::Error> {
        serializer.serialize_str(value)
    }

    pub(super) fn deserialize<'de, D: Deserializer<'de>>(
        deserializer: D,
    ) -> Result<Arc<str>, D::Error> {
        let value = <std::borrow::Cow<'de, str>>::deserialize(deserializer)?;
        Ok(Arc::from(&*value))
    }
}

/// `Option<Arc<str>>` as an optional string on the wire.
mod opt_arc_str {
    use std::sync::Arc;

    use serde::{Deserialize, Deserializer, Serializer};

    pub(super) fn serialize<S: Serializer>(
        value: &Option<Arc<str>>,
        serializer: S,
    ) -> Result<S::Ok, S::Error> {
        match value {
            Some(value) => serializer.serialize_some(&**value),
            None => serializer.serialize_none(),
        }
    }

    pub(super) fn deserialize<'de, D: Deserializer<'de>>(
        deserializer: D,
    ) -> Result<Option<Arc<str>>, D::Error> {
        let value = <Option<std::borrow::Cow<'de, str>>>::deserialize(deserializer)?;
        Ok(value.map(|value| Arc::from(&*value)))
    }
}

/// Effect classification used by typed binding and recording.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum EffectFamily {
    /// A completion request to a model.
    Completion,
    /// A tool call.
    Tool,
    /// A text or image embedding request.
    Embed,
    /// A reranking request.
    Rerank,
    /// A conversation-memory operation.
    Memory,
    /// A vector-store retrieval.
    Retrieve,
    /// An out-of-tree kind.
    Custom,
}

impl EffectFamily {
    /// The family's stable label.
    pub const fn name(self) -> &'static str {
        match self {
            Self::Completion => "completion",
            Self::Tool => "tool_call",
            Self::Embed => "embed",
            Self::Rerank => "rerank",
            Self::Memory => "memory",
            Self::Retrieve => "retrieve",
            Self::Custom => "custom",
        }
    }
}

impl fmt::Display for EffectFamily {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(self.name())
    }
}

/// Sealed mapping between typed requests and answers and the wire's
/// [`EffectKind`] and [`Outcome`]. Hosts extend the protocol through
/// [`CustomEffect`] and [`family::Custom`], not new family implementations.
pub trait Family: sealed::Sealed + Clone + Copy + Send + Sync + 'static {
    /// The family this marker names.
    const FAMILY: EffectFamily;
    /// What a typed dispatch of this family takes.
    type Request: WasmCompatSend + 'static;
    /// What it resolves to.
    type Answer: WasmCompatSend + 'static;
    /// The wire form of a request, or the report for a request that has
    /// none (a [`CustomEffect`] whose `Serialize` fails). The in-tree
    /// families always have one; a typed dispatch of a request without one
    /// is pre-failed by the bus and never reaches a handler or a log.
    fn wrap(request: Self::Request) -> Result<EffectKind, ErrorReport>;
    /// The typed answer, or the report for an outcome of another family.
    fn unwrap(outcome: Outcome) -> Result<Self::Answer, ErrorReport>;
    /// The report [`Family::unwrap`] gives for an outcome of another family.
    fn mismatch(outcome: &Outcome) -> ErrorReport {
        ErrorReport::new(
            crate::error::ErrorKind::Internal,
            format!(
                "expected a {} outcome, the handler answered {}",
                Self::FAMILY,
                outcome.family()
            ),
        )
    }
}

mod sealed {
    pub trait Sealed {}
}

/// Sealed handler-family declaration. Typed families expose `Some(F::FAMILY)`;
/// [`family::Dynamic`] exposes `None` and requires runtime family checks.
pub trait Served: sealed::Sealed + 'static {
    /// The family, when the handler has one.
    const SERVED: Option<EffectFamily>;
}

// `Served` is sealed; a handler's family is a `Family` marker or `Dynamic`,
// and an error about it should say so rather than suggest this impl.
#[diagnostic::do_not_recommend]
impl<F: Family> Served for F {
    const SERVED: Option<EffectFamily> = Some(F::FAMILY);
}

/// A tool call as a typed request: the name and the raw JSON arguments.
/// The context the tool runs with is not part of the request: it travels
/// beside the dispatch, attached by the driver to the handler's sink
/// (`Dispatch::scope::<ToolContext>()`), and what the tool publishes
/// comes back the same way ([`crate::tool::PublishedContext`]).
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ToolCallRequest {
    /// The tool's name (the name the model calls it by).
    pub name: String,
    /// The arguments as a JSON string.
    pub args: String,
}

/// A host's own effect, typed the way a [`ToolContext`](crate::tool::ToolContext) value is: a
/// declared kind label and a declared answer type, both serde. The wire
/// form is [`EffectKind::Custom`] / [`Outcome::Custom`]; the type never
/// crosses it.
pub trait CustomEffect: Serialize + serde::de::DeserializeOwned + WasmCompatSend + 'static {
    /// The kind label this effect dispatches under; a handler's
    /// [`FamilyDescriptor::Custom`] must name the same label.
    const KIND: &'static str;
    /// What the handler answers.
    type Answer: Serialize + serde::de::DeserializeOwned + WasmCompatSend + 'static;
}

mod key;

pub use key::Key;

/// A key for a model handler: `model:<label>`.
pub fn model_key(label: &str) -> HandlerKey {
    HandlerKey::from(format!("model:{label}"))
}

/// A key for a tool handler: `tool:<name>`.
pub fn tool_key(name: &str) -> HandlerKey {
    HandlerKey::from(format!("tool:{name}"))
}

/// A key for an embedding handler: `embed:<label>`.
pub fn embed_key(label: &str) -> HandlerKey {
    HandlerKey::from(format!("embed:{label}"))
}

/// A key for a rerank handler: `rerank:<label>`.
pub fn rerank_key(label: &str) -> HandlerKey {
    HandlerKey::from(format!("rerank:{label}"))
}

/// A key for a labelled memory handler: `memory:<label>`. An agent's own
/// memory is registered under the bare `memory` key.
pub fn memory_key(label: &str) -> HandlerKey {
    HandlerKey::from(format!("memory:{label}"))
}

/// A key for a labelled retrieval handler: `retrieve:<label>`. An agent's
/// own index is registered under the bare `retrieve` key.
pub fn retrieve_key(label: &str) -> HandlerKey {
    HandlerKey::from(format!("retrieve:{label}"))
}

/// Type-level markers for built-in and host-defined effect families.
///
/// ```
/// use rig_core::effect::{family, Family, EffectFamily};
/// assert_eq!(family::Completion::FAMILY, EffectFamily::Completion);
/// ```
pub mod family {
    use std::marker::PhantomData;

    use super::{
        CustomEffect, EffectFamily, EffectKind, EmbedInputs, EmbedOutputs, Family, MemoryOp,
        MemoryOutcome, Outcome, RetrieveQuery, RetrievedDocuments, ToolCallRequest, sealed::Sealed,
    };
    use crate::{
        completion::{CompletionRequest, CompletionResponse},
        error::{ErrorKind, ErrorReport},
        operation::RerankRequest,
        rerank::RerankResponse,
        tool::ToolResult,
    };

    macro_rules! marker {
        ($($(#[$doc:meta])* $name:ident => $family:ident, $request:ty, $answer:ty, $wrap:expr, $unwrap:expr;)+) => {$(
            $(#[$doc])*
            #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
            pub struct $name;

            impl Sealed for $name {}

            impl Family for $name {
                const FAMILY: EffectFamily = EffectFamily::$family;
                type Request = $request;
                type Answer = $answer;

                fn wrap(request: Self::Request) -> Result<EffectKind, ErrorReport> {
                    let wrap: fn(Self::Request) -> EffectKind = $wrap;
                    Ok(wrap(request))
                }

                fn unwrap(outcome: Outcome) -> Result<Self::Answer, ErrorReport> {
                    let unwrap: fn(Outcome) -> Result<Self::Answer, ErrorReport> = $unwrap;
                    unwrap(outcome)
                }
            }
        )+};
    }

    marker! {
        /// The completion family: a unary completion (a streaming dispatch is
        /// `ModelHandle::stream`, not a typed request).
        Completion => Completion, CompletionRequest, CompletionResponse,
            |request| EffectKind::Completion { request, stream: false },
            |outcome| match outcome {
                Outcome::Completion(response) => Ok(response),
                other => Err(Completion::mismatch(&other)),
            };
        /// The tool family.
        Tool => Tool, ToolCallRequest, ToolResult,
            |request| EffectKind::ToolCall { name: request.name, args: request.args },
            |outcome| match outcome {
                Outcome::ToolResult { result } => Ok(result),
                other => Err(Tool::mismatch(&other)),
            };
        /// The embedding family.
        Embed => Embed, EmbedInputs, EmbedOutputs,
            |inputs| EffectKind::Embed { inputs },
            |outcome| match outcome {
                Outcome::Embeddings(outputs) => Ok(outputs),
                other => Err(Embed::mismatch(&other)),
            };
        /// The reranking family.
        Rerank => Rerank, RerankRequest, RerankResponse,
            |request| EffectKind::Rerank { request },
            |outcome| match outcome {
                Outcome::Reranked(response) => Ok(response),
                other => Err(Rerank::mismatch(&other)),
            };
        /// The conversation-memory family.
        Memory => Memory, MemoryOp, MemoryOutcome,
            |op| EffectKind::Memory { op },
            |outcome| match outcome {
                Outcome::Memory(answer) => Ok(answer),
                other => Err(Memory::mismatch(&other)),
            };
        /// The retrieval family.
        Retrieve => Retrieve, RetrieveQuery, RetrievedDocuments,
            |query| EffectKind::Retrieve { query },
            |outcome| match outcome {
                Outcome::Documents(documents) => Ok(documents),
                other => Err(Retrieve::mismatch(&other)),
            };
    }

    /// A handler with no one family: a replayer answering whatever its log
    /// holds, or an erased handler forwarding to whatever it wraps.
    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
    pub struct Dynamic;

    impl Sealed for Dynamic {}

    impl super::Served for Dynamic {
        const SERVED: Option<EffectFamily> = None;
    }

    /// The family of one host-defined effect `E`: dispatches
    /// [`EffectKind::Custom`] under `E::KIND` and answers `E::Answer`.
    pub struct Custom<E: CustomEffect>(PhantomData<fn() -> E>);

    impl<E: CustomEffect> Custom<E> {
        /// The marker.
        pub const fn new() -> Self {
            Self(PhantomData)
        }
    }

    // Written by hand: a derive would demand `E: Clone` (and friends), and
    // the marker must be `Copy` for every `E`.
    impl<E: CustomEffect> Clone for Custom<E> {
        fn clone(&self) -> Self {
            *self
        }
    }
    impl<E: CustomEffect> Copy for Custom<E> {}
    impl<E: CustomEffect> Default for Custom<E> {
        fn default() -> Self {
            Self::new()
        }
    }
    impl<E: CustomEffect> std::fmt::Debug for Custom<E> {
        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            write!(f, "Custom<{}>", E::KIND)
        }
    }
    impl<E: CustomEffect> PartialEq for Custom<E> {
        fn eq(&self, _: &Self) -> bool {
            true
        }
    }
    impl<E: CustomEffect> Eq for Custom<E> {}
    impl<E: CustomEffect> std::hash::Hash for Custom<E> {
        fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
            E::KIND.hash(state);
        }
    }

    impl<E: CustomEffect> Sealed for Custom<E> {}

    impl<E: CustomEffect> Family for Custom<E> {
        const FAMILY: EffectFamily = EffectFamily::Custom;
        type Request = E;
        type Answer = E::Answer;

        fn wrap(request: E) -> Result<EffectKind, ErrorReport> {
            // Reject unencodable requests before dispatch so no handler or log
            // receives an effect without a wire representation.
            serde_json::to_value(&request)
                .map(|payload| EffectKind::Custom {
                    kind: std::sync::Arc::from(E::KIND),
                    payload,
                })
                .map_err(|error| {
                    ErrorReport::new(
                        ErrorKind::Request,
                        format!("the `{}` effect did not serialize: {error}", E::KIND),
                    )
                })
        }

        fn unwrap(outcome: Outcome) -> Result<E::Answer, ErrorReport> {
            match outcome {
                Outcome::Custom { payload: value } => {
                    serde_json::from_value(value).map_err(|error| {
                        ErrorReport::new(
                            ErrorKind::Internal,
                            format!(
                                "the answer to the `{}` effect did not deserialize: {error}",
                                E::KIND
                            ),
                        )
                    })
                }
                other => Err(Self::mismatch(&other)),
            }
        }
    }
}

/// What a registered handler is: its key and its family-specific description.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct HandlerDescriptor {
    /// The key the handler is registered under.
    pub key: HandlerKey,
    /// The family and its advertised metadata.
    pub family: FamilyDescriptor,
    /// Interceptor names in outermost-first order, used to validate the layer
    /// stack during replay. Empty for a handler without layers.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub layers: Vec<String>,
}

/// The family-keyed description of a handler. The variant *is* the family:
/// binding a typed view compares [`Family::FAMILY`] against
/// [`FamilyDescriptor::family`].
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "family", rename_all = "snake_case")]
pub enum FamilyDescriptor {
    /// A completion model.
    Completion {
        /// The model's label.
        model: ModelRef,
        /// The capability snapshot a runtime prepares requests against.
        capabilities: ProviderCapabilities,
    },
    /// A tool.
    Tool {
        /// The tool's name (the name the model calls it by).
        name: String,
        /// The tool's description.
        description: String,
        /// The JSON schema of the tool's arguments.
        parameters: serde_json::Value,
        /// Present when the tool is retrievable by embedding.
        #[serde(default, skip_serializing_if = "Option::is_none")]
        embedding: Option<ToolEmbeddingDescriptor>,
    },
    /// An embedding model.
    Embed {
        /// The model's label.
        model: String,
        /// The vector dimension, when the model declares one.
        #[serde(default, skip_serializing_if = "Option::is_none")]
        dims: Option<usize>,
        /// The largest batch the model accepts.
        max_documents: usize,
        /// Whether the model embeds text or images.
        modality: EmbedModality,
    },
    /// A reranking model.
    Rerank {
        /// The model's label.
        model: String,
        /// The largest batch the model accepts.
        max_documents: usize,
    },
    /// A conversation-memory backend.
    Memory {},
    /// A vector-store index.
    Retrieve {},
    /// A handler for an out-of-tree kind: the label it serves.
    Custom {
        /// The [`EffectKind::Custom`] kind label this handler answers.
        kind: String,
    },
}

impl FamilyDescriptor {
    /// The family this descriptor belongs to.
    pub const fn family(&self) -> EffectFamily {
        match self {
            Self::Completion { .. } => EffectFamily::Completion,
            Self::Tool { .. } => EffectFamily::Tool,
            Self::Embed { .. } => EffectFamily::Embed,
            Self::Rerank { .. } => EffectFamily::Rerank,
            Self::Memory {} => EffectFamily::Memory,
            Self::Retrieve {} => EffectFamily::Retrieve,
            Self::Custom { .. } => EffectFamily::Custom,
        }
    }
}

/// The embedding context of a retrievable tool: what
/// [`ToolEmbedding`](crate::tool::ToolEmbedding) advertises, as data.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ToolEmbeddingDescriptor {
    /// The tool's serialized context.
    pub context: serde_json::Value,
    /// The documents the tool is retrieved by.
    pub embedding_docs: Vec<String>,
}

/// Which modality an embedding handler serves.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum EmbedModality {
    /// Text documents.
    Text,
    /// Image bytes.
    Image,
}

/// One effect: what a handler is asked to do.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "effect", rename_all = "snake_case")]
pub enum EffectKind {
    /// A completion request.
    Completion {
        /// The prepared request.
        request: CompletionRequest,
        /// Whether the response streams (`dispatch_stream`) or is unary.
        stream: bool,
    },
    /// A tool call. The context the tool runs with is not on the wire:
    /// see [`ToolCallRequest`].
    ToolCall {
        /// The tool's name.
        name: String,
        /// The raw JSON arguments as the model produced them.
        args: String,
    },
    /// An embedding request.
    Embed {
        /// The inputs to embed.
        inputs: EmbedInputs,
    },
    /// A reranking request.
    Rerank {
        /// The query and the documents.
        request: RerankRequest,
    },
    /// A conversation-memory operation.
    Memory {
        /// The operation.
        op: MemoryOp,
    },
    /// A retrieval.
    Retrieve {
        /// The query.
        query: RetrieveQuery,
    },
    /// An out-of-tree effect.
    Custom {
        /// The host-defined kind label.
        #[serde(with = "arc_str")]
        kind: Arc<str>,
        /// The host-defined payload.
        payload: serde_json::Value,
    },
}

impl EffectKind {
    /// The family of this effect.
    pub const fn family(&self) -> EffectFamily {
        match self {
            Self::Completion { .. } => EffectFamily::Completion,
            Self::ToolCall { .. } => EffectFamily::Tool,
            Self::Embed { .. } => EffectFamily::Embed,
            Self::Rerank { .. } => EffectFamily::Rerank,
            Self::Memory { .. } => EffectFamily::Memory,
            Self::Retrieve { .. } => EffectFamily::Retrieve,
            Self::Custom { .. } => EffectFamily::Custom,
        }
    }

    /// Classification label without payload data. Custom effects return their
    /// host-defined kind label.
    pub fn name(&self) -> &str {
        match self {
            Self::Custom { kind, .. } => kind,
            Self::Completion { .. }
            | Self::ToolCall { .. }
            | Self::Embed { .. }
            | Self::Rerank { .. }
            | Self::Memory { .. }
            | Self::Retrieve { .. } => self.family().name(),
        }
    }

    /// Whether this effect answers as a stream. Only a streaming completion
    /// does; every other kind is unary.
    pub const fn streams(&self) -> bool {
        match self {
            Self::Completion { stream, .. } => *stream,
            Self::ToolCall { .. }
            | Self::Embed { .. }
            | Self::Rerank { .. }
            | Self::Memory { .. }
            | Self::Retrieve { .. }
            | Self::Custom { .. } => false,
        }
    }
}

/// The inputs of an embedding request: the texts or images an embedding
/// model's call embeds.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "modality", content = "inputs", rename_all = "snake_case")]
pub enum EmbedInputs {
    /// Text documents.
    Texts(Vec<String>),
    /// Image bytes.
    Images(Vec<Vec<u8>>),
}

/// A conversation-memory operation: the transcription of
/// [`ConversationMemory`](crate::memory::ConversationMemory).
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "op", rename_all = "snake_case")]
pub enum MemoryOp {
    /// Load a conversation's history.
    Load {
        /// The conversation.
        conversation: ConversationId,
    },
    /// Append messages to a conversation.
    Append {
        /// The conversation.
        conversation: ConversationId,
        /// The messages, in order.
        messages: Vec<Message>,
    },
    /// Clear a conversation.
    Clear {
        /// The conversation.
        conversation: ConversationId,
    },
}

/// A retrieval: the transcription of
/// [`VectorStoreIndex`](crate::vector_store::VectorStoreIndex) over the
/// dynamic filter. The typed result parameter of `top_n<T>` stays on the
/// client side.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "query", rename_all = "snake_case")]
pub enum RetrieveQuery {
    /// Scored documents.
    TopN {
        /// The search.
        req: VectorSearchRequest<Filter<serde_json::Value>>,
    },
    /// Scored ids only.
    TopNIds {
        /// The search.
        req: VectorSearchRequest<Filter<serde_json::Value>>,
    },
}

/// What a handler answered.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "outcome", rename_all = "snake_case")]
pub enum Outcome {
    /// A unary completion.
    Completion(CompletionResponse),
    /// A tool call's result. The values the tool published into its
    /// context are not on the wire: they come back beside the sink
    /// ([`crate::tool::PublishedContext`]).
    ToolResult {
        /// The result.
        result: ToolResult,
    },
    /// Embeddings.
    Embeddings(EmbedOutputs),
    /// A reranking.
    Reranked(RerankResponse),
    /// A memory operation's answer.
    Memory(MemoryOutcome),
    /// Retrieved documents.
    Documents(RetrievedDocuments),
    /// An out-of-tree answer.
    Custom {
        /// The answer, nested so every JSON shape is representable.
        payload: serde_json::Value,
    },
}

impl Outcome {
    /// The family this outcome answers.
    pub const fn family(&self) -> EffectFamily {
        match self {
            Self::Completion(_) => EffectFamily::Completion,
            Self::ToolResult { .. } => EffectFamily::Tool,
            Self::Embeddings(_) => EffectFamily::Embed,
            Self::Reranked(_) => EffectFamily::Rerank,
            Self::Memory(_) => EffectFamily::Memory,
            Self::Documents(_) => EffectFamily::Retrieve,
            Self::Custom { .. } => EffectFamily::Custom,
        }
    }
}

/// The answer to an [`EmbedInputs`].
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "modality", content = "response", rename_all = "snake_case")]
pub enum EmbedOutputs {
    /// Text embeddings.
    Texts(EmbeddingResponse),
    /// Image embeddings.
    Images(EmbeddingResponse),
}

/// The answer to a [`MemoryOp`].
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "memory", rename_all = "snake_case")]
pub enum MemoryOutcome {
    /// The loaded history.
    Loaded {
        /// The messages.
        messages: Vec<Message>,
    },
    /// The append succeeded.
    Appended,
    /// The clear succeeded.
    Cleared,
}

/// The answer to a [`RetrieveQuery`].
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "retrieved", content = "results", rename_all = "snake_case")]
pub enum RetrievedDocuments {
    /// Scored documents: `(score, id, document)`.
    Scored(Vec<(f64, String, serde_json::Value)>),
    /// Scored ids: `(score, id)`.
    Ids(Vec<(f64, String)>),
}

/// One delivery observed by a scheduled consumer. Entries are ordered by
/// visibility, independently of dispatch ids and handler completion order.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct Delivery {
    /// The schedule pass in which the consumer observed this delivery.
    /// Pass numbers group deliveries; they are not wall-clock timestamps.
    pub batch: u64,
    /// The effect whose state became visible.
    pub id: EffectId,
    /// The visible transition.
    pub kind: DeliveryKind,
}

/// What a scheduled consumer collected at one observation boundary.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "delivery", rename_all = "snake_case")]
pub enum DeliveryKind {
    /// The effect's outcome became visible.
    Outcome,
    /// The effect was cancelled after its original handler answer was recorded,
    /// without delivering an outcome. Kept events may include undelivered items.
    Cancelled,
    /// A batch of stream items became visible together. Includes error items;
    /// successful event bytes remain in the effect record's event sequence.
    Stream {
        /// Items collected since the preceding delivery for this effect.
        items: usize,
    },
}

/// One recorded exchange: the effect, who served it, and the answer.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EffectRecord {
    /// Explicitly published tool-result metadata, including on errors. Never
    /// inbound authentication values or runtime scopes. `None` means the tool
    /// did not publish a context before resolving. The wire field is required:
    /// `null` explicitly records no publication; omission cannot establish
    /// whether an older recorder lost output and is rejected when decoding.
    #[serde(deserialize_with = "Option::deserialize")]
    pub tool_output: Option<crate::tool::ToolResultContext>,
    /// The dispatch's id.
    pub id: EffectId,
    /// The handler the effect was routed to.
    pub key: HandlerKey,
    /// The effect.
    pub kind: EffectKind,
    /// The answer.
    pub outcome: Result<Outcome, ErrorReport>,
    /// A streamed dispatch's items, verbatim, when the recorder was asked
    /// to keep them (`EffectLogRecorder::keeping_stream_events`); `None`
    /// otherwise, and the answer is the fold. A replayer re-emits these
    /// when present, so a replayed consumer sees the original fragment
    /// boundaries. They read back through [`Transcript::parse_prefix`].
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub events: Option<Transcript>,
    /// Who a streamed dispatch's reply was from, as its relay stated it
    /// before the first item, when the items are kept. A replayer sends it
    /// first, so a replayed stream cut short knows its origin.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub stream_origin: Option<crate::message::Origin>,
    /// Parent dispatch for a nested handler call, or `None` for a root dispatch.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub parent: Option<EffectId>,
    /// Stable identifier of the dispatching program or run, not a runtime handle.
    /// `None` when the dispatcher supplied no scope.
    #[serde(default, skip_serializing_if = "Option::is_none", with = "opt_arc_str")]
    pub scope: Option<std::sync::Arc<str>>,
}

/// Required handler keys and their effect families, ordered by key.
/// [`Self::is_subset_of`] validates the requirements against handler descriptors.
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(transparent)]
pub struct EffectRow(std::collections::BTreeMap<HandlerKey, EffectFamily>);

/// The first entry of a row a set of handlers does not serve: the key,
/// the family the row needs, and what serves the key instead, if anything.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RowGap {
    /// The key the row names.
    pub key: HandlerKey,
    /// The family the row needs it as.
    pub needed: EffectFamily,
    /// The family a handler serves it as, or `None` when nothing serves it.
    pub served: Option<EffectFamily>,
}

impl std::fmt::Display for RowGap {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self.served {
            Some(served) => write!(
                f,
                "`{}` is needed as {} but served as {served}",
                self.key, self.needed
            ),
            None => write!(f, "`{}` ({}) is not served", self.key, self.needed),
        }
    }
}

/// One difference between two rows.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum RowDiff {
    /// The other row lacks this entry.
    Missing {
        /// The key this row names.
        key: HandlerKey,
        /// The family this row names it as.
        family: EffectFamily,
    },
    /// The other row has an entry this one lacks.
    Extra {
        /// The key the other row names.
        key: HandlerKey,
        /// The family the other row names it as.
        family: EffectFamily,
    },
    /// Both rows name the key, as different families.
    Family {
        /// The key both name.
        key: HandlerKey,
        /// This row's family for it.
        this: EffectFamily,
        other: EffectFamily,
    },
}

impl std::fmt::Display for RowDiff {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Missing { key, family } => write!(f, "`{key}` ({family}) is missing"),
            Self::Extra { key, family } => write!(f, "`{key}` ({family}) is extra"),
            Self::Family { key, this, other } => {
                write!(f, "`{key}` is {this} here and {other} there")
            }
        }
    }
}

impl EffectRow {
    /// An empty row.
    pub fn new() -> Self {
        Self::default()
    }

    /// Name `key` as `family`; a key already named is re-named.
    pub fn insert(&mut self, key: HandlerKey, family: EffectFamily) -> Option<EffectFamily> {
        self.0.insert(key, family)
    }

    /// Name `key` as `family` unless the row already names it.
    pub fn insert_if_absent(&mut self, key: HandlerKey, family: EffectFamily) {
        self.0.entry(key).or_insert(family);
    }

    /// Remove `key` from the row; the family it named, if it was there.
    pub fn remove(&mut self, key: &HandlerKey) -> Option<EffectFamily> {
        self.0.remove(key)
    }

    /// The family the row names `key` as.
    pub fn get(&self, key: &HandlerKey) -> Option<&EffectFamily> {
        self.0.get(key)
    }

    /// Whether the row names `key`.
    pub fn contains_key(&self, key: &HandlerKey) -> bool {
        self.0.contains_key(key)
    }

    /// The keys, in order.
    pub fn keys(&self) -> impl Iterator<Item = &HandlerKey> {
        self.0.keys()
    }

    /// The entries, in key order.
    pub fn iter(&self) -> impl Iterator<Item = (&HandlerKey, &EffectFamily)> {
        self.0.iter()
    }

    /// How many keys the row names.
    pub fn len(&self) -> usize {
        self.0.len()
    }

    /// Whether the row names no key.
    pub fn is_empty(&self) -> bool {
        self.0.is_empty()
    }

    /// Validates that every key is served with the required family, returning
    /// the first gap in key order. Uses the first descriptor matching each key.
    pub fn is_subset_of(&self, handlers: &[HandlerDescriptor]) -> Result<(), RowGap> {
        for (key, needed) in &self.0 {
            let served = handlers
                .iter()
                .find(|descriptor| &descriptor.key == key)
                .map(|descriptor| descriptor.family.family());
            match served {
                Some(served) if served == *needed => {}
                served => {
                    return Err(RowGap {
                        key: key.clone(),
                        needed: *needed,
                        served,
                    });
                }
            }
        }
        Ok(())
    }

    /// Returns missing and mismatched entries in this row's key order, followed
    /// by extra entries in the other row's key order.
    pub fn diff(&self, other: &EffectRow) -> Vec<RowDiff> {
        let mut diffs = Vec::new();
        for (key, family) in &self.0 {
            match other.0.get(key) {
                None => diffs.push(RowDiff::Missing {
                    key: key.clone(),
                    family: *family,
                }),
                Some(theirs) if theirs != family => diffs.push(RowDiff::Family {
                    key: key.clone(),
                    this: *family,
                    other: *theirs,
                }),
                Some(_) => {}
            }
        }
        for (key, family) in &other.0 {
            if !self.0.contains_key(key) {
                diffs.push(RowDiff::Extra {
                    key: key.clone(),
                    family: *family,
                });
            }
        }
        diffs
    }
}

impl FromIterator<(HandlerKey, EffectFamily)> for EffectRow {
    fn from_iter<I: IntoIterator<Item = (HandlerKey, EffectFamily)>>(iter: I) -> Self {
        Self(iter.into_iter().collect())
    }
}

impl<'a> IntoIterator for &'a EffectRow {
    type Item = (&'a HandlerKey, &'a EffectFamily);
    type IntoIter = std::collections::btree_map::Iter<'a, HandlerKey, EffectFamily>;

    fn into_iter(self) -> Self::IntoIter {
        self.0.iter()
    }
}

impl IntoIterator for EffectRow {
    type Item = (HandlerKey, EffectFamily);
    type IntoIter = std::collections::btree_map::IntoIter<HandlerKey, EffectFamily>;

    fn into_iter(self) -> Self::IntoIter {
        self.0.into_iter()
    }
}

// The protocol crosses threads and serializes on every target.
const _: fn() = || {
    fn assert_wire<T: Clone + Send + Sync + 'static + Serialize + serde::de::DeserializeOwned>() {}
    assert_wire::<EffectId>();
    assert_wire::<HandlerKey>();
    assert_wire::<EffectFamily>();
    assert_wire::<HandlerDescriptor>();
    assert_wire::<FamilyDescriptor>();
    assert_wire::<ToolEmbeddingDescriptor>();
    assert_wire::<EmbedModality>();
    assert_wire::<EffectKind>();
    assert_wire::<EmbedInputs>();
    assert_wire::<RerankRequest>();
    assert_wire::<MemoryOp>();
    assert_wire::<RetrieveQuery>();
    assert_wire::<Outcome>();
    assert_wire::<EmbedOutputs>();
    assert_wire::<MemoryOutcome>();
    assert_wire::<RetrievedDocuments>();
    assert_wire::<EffectRecord>();
    assert_wire::<EffectRow>();
    assert_wire::<Transcript>();
};

#[cfg(test)]
mod tests;