mobius/middleware/

context.rs

use std::borrow::Cow;
use std::collections::{BTreeMap, BTreeSet};
use std::sync::Arc;

use serde_json::Value;

use super::MiddlewareStack;
use super::tools::Catalog;
use super::tools::ToolResult;
use super::{approximate_item_tokens, approximate_tokens, serialized_len};
use crate::agent::{AgentRole, WeakAgentSender};
use crate::backend::checkpoint::{
    Checkpoint, CheckpointStore, ContextRewriteReason, ExecutionOutcome, MAX_QUEUED_MESSAGES,
    QueuedMessage as DurableQueuedMessage, QueuedMessageBoundary,
};
use crate::backend::model::{ModelRouter, message_input};
use crate::backend::sandbox::ApprovalPolicy;
use crate::protocol::{
    EventMsg, FrontendEvent, MAX_CAPABILITY_INPUT_BYTES, MessageAuthor, MessageEvent,
    MessageSubmission, MessageTarget, ReviewDecision, SessionContext, SessionFileReference,
    TokenUsage, ToolCall, message_metadata,
};
use crate::{Error, Result};

/// Sends middleware-owned UI updates without depending on a concrete frontend.
pub type FrontendEventSink = Arc<dyn Fn(FrontendEvent) -> Result<()> + Send + Sync>;

/// Read-only queued message owned by the middleware receiving it.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct QueuedMessageView<'a> {
    item: &'a DurableQueuedMessage,
}

impl<'a> QueuedMessageView<'a> {
    /// Returns the identity token required by a conditional queue mutation.
    #[must_use]
    pub fn id(&self) -> &'a str {
        self.item.id()
    }

    /// Returns the prepared presentation event.
    #[must_use]
    pub fn event(&self) -> MessageEvent {
        self.item.event()
    }
}

/// Read-only startup snapshot containing only one middleware's queued messages.
#[derive(Clone, Default)]
pub struct QueuedMessageSnapshot {
    items: Vec<DurableQueuedMessage>,
}

impl QueuedMessageSnapshot {
    /// Returns every queued item owned by this middleware, oldest first.
    pub fn views(&self) -> impl Iterator<Item = QueuedMessageView<'_>> {
        self.items.iter().map(|item| QueuedMessageView { item })
    }

    pub(super) fn for_owner(owner: &str, items: &[DurableQueuedMessage]) -> Self {
        Self {
            items: items
                .iter()
                .filter(|item| item.owner() == owner)
                .cloned()
                .collect(),
        }
    }
}

/// Mutable scoped view of messages retained until their delivery boundary.
pub struct MessageQueue<'a> {
    items: &'a mut Vec<DurableQueuedMessage>,
    owner: Option<&'static str>,
}

impl<'a> MessageQueue<'a> {
    pub(crate) fn new(items: &'a mut Vec<DurableQueuedMessage>) -> Self {
        Self { items, owner: None }
    }

    pub(super) fn scope(&mut self, owner: &'static str) {
        self.owner = Some(owner);
    }

    fn owner(&self) -> Result<&'static str> {
        self.owner
            .ok_or_else(|| Error::Config("message queue is not scoped to a middleware".into()))
    }

    /// Returns the number of queued items owned by this middleware.
    #[must_use]
    pub fn count(&self) -> usize {
        let Some(owner) = self.owner else {
            return 0;
        };
        self.items
            .iter()
            .filter(|item| item.owner() == owner)
            .count()
    }

    /// Returns the newest message available to this context.
    #[must_use]
    pub fn latest(&self) -> Option<QueuedMessageView<'_>> {
        let owner = self.owner?;
        self.items
            .iter()
            .rev()
            .find(|item| item.owner() == owner)
            .map(|item| QueuedMessageView { item })
    }

    /// Returns one owned item by its revision identity.
    #[must_use]
    pub fn find(&self, id: &str) -> Option<QueuedMessageView<'_>> {
        let owner = self.owner?;
        self.items
            .iter()
            .find(|item| item.owner() == owner && item.id() == id)
            .map(|item| QueuedMessageView { item })
    }

    /// Appends one prepared message, or returns `false` when it is full or duplicated.
    /// # Errors
    ///
    /// Returns an error if validation or an operation required by this function fails.
    pub fn enqueue(
        &mut self,
        id: &str,
        boundary: QueuedMessageBoundary,
        event: MessageEvent,
    ) -> Result<bool> {
        let owner = self.owner()?;
        let item = DurableQueuedMessage::new(owner, id, boundary, event)?;
        if self.items.len() >= MAX_QUEUED_MESSAGES {
            return Ok(false);
        }
        if self
            .items
            .iter()
            .any(|item| item.owner() == owner && item.id() == id)
        {
            return Ok(false);
        }
        self.items.push(item);
        Ok(true)
    }

    /// Atomically replaces one owned item while preserving its queue position.
    /// # Errors
    ///
    /// Returns an error if validation or an operation required by this function fails.
    pub fn replace(&mut self, id: &str, replacement_id: &str, event: MessageEvent) -> Result<bool> {
        let owner = self.owner()?;
        let Some(index) = self
            .items
            .iter()
            .position(|item| item.owner() == owner && item.id() == id)
        else {
            return Ok(false);
        };
        if self.items.iter().enumerate().any(|(candidate, item)| {
            candidate != index && item.owner() == owner && item.id() == replacement_id
        }) {
            return Ok(false);
        }
        self.items[index].replace(replacement_id, event)?;
        Ok(true)
    }

    pub(crate) fn stage_model_messages(&mut self, turn_id: &str) -> Result<Vec<PreparedMessage>> {
        let Some(owner) = self.owner else {
            return Ok(Vec::new());
        };
        self.items
            .extract_if(.., |item| {
                item.owner() == owner
                    && matches!(
                        item.boundary(),
                        QueuedMessageBoundary::Steer { turn_id: target }
                            if target == turn_id
                    )
            })
            .map(PreparedMessage::try_from)
            .collect()
    }

    pub(crate) fn next_turn(&self) -> Result<Option<PreparedMessage>> {
        let owner = self.owner()?;
        self.items
            .iter()
            .find(|item| item.owner() == owner && item.boundary().starts_turn())
            .cloned()
            .map(PreparedMessage::try_from)
            .transpose()
    }

    pub(crate) fn consume_next_turn(&mut self, id: &str) -> Result<()> {
        let owner = self.owner()?;
        let index = self
            .items
            .iter()
            .position(|item| {
                item.owner() == owner && item.id() == id && item.boundary().starts_turn()
            })
            .ok_or_else(|| Error::Checkpoint("prepared message is no longer queued".into()))?;
        self.items.remove(index);
        Ok(())
    }

    pub(crate) fn promote_failed_turn(&mut self, turn_id: &str) -> Result<()> {
        let owner = self.owner()?;
        for item in self.items.iter_mut().filter(|item| {
            item.owner() == owner
                && matches!(
                    item.boundary(),
                    QueuedMessageBoundary::Steer { turn_id: target }
                        if target == turn_id
                )
        }) {
            item.promote_to_next_turn()?;
        }
        Ok(())
    }
}

/// One queued message prepared for its model boundary.
pub(crate) struct PreparedMessage {
    pub(crate) submission_id: String,
    pub(crate) input: Value,
    pub(crate) event: EventMsg,
    pub(crate) title_seed: Option<String>,
    pub(crate) boundary_events: Vec<EventMsg>,
}

impl TryFrom<DurableQueuedMessage> for PreparedMessage {
    type Error = Error;

    fn try_from(message: DurableQueuedMessage) -> Result<Self> {
        let (submission_id, event) = message.into_parts();
        let input = message_input(&event)?;
        let title_seed = matches!(
            event.author,
            MessageAuthor::User | MessageAuthor::Source { .. }
        )
        .then(|| event.text.trim().to_string())
        .filter(|title| !title.is_empty());
        Ok(Self {
            submission_id,
            input,
            event: EventMsg::Message(event),
            title_seed,
            boundary_events: Vec::new(),
        })
    }
}

/// Durable runtime identity exposed while middleware starts a session.
#[derive(Clone)]
pub struct RuntimeContext {
    /// The sender.
    pub sender: WeakAgentSender,
    /// The checkpoints.
    pub checkpoints: Arc<dyn CheckpointStore>,
    /// The session identifier.
    pub session_id: String,
    /// The model route.
    pub model_route: String,
    /// The model.
    pub model: String,
    /// The approval policy.
    pub approval_policy: ApprovalPolicy,
    /// The session context.
    pub session_context: SessionContext,
    /// The metadata.
    pub metadata: BTreeMap<String, Value>,
    /// The role.
    pub role: AgentRole,
    /// The frontend.
    pub frontend: FrontendEventSink,
}

impl RuntimeContext {
    pub(crate) fn turn_identity<'a>(
        &'a self,
        turn_id: &'a str,
        author: &'a MessageAuthor,
    ) -> TurnIdentity<'a> {
        TurnIdentity {
            session_id: &self.session_id,
            turn_id,
            model: &self.model,
            approval_policy: self.approval_policy,
            author,
        }
    }
}

/// Stable facts shared by hooks that run within one active turn.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct TurnIdentity<'a> {
    /// Trusted provenance of the message that initiated the turn.
    pub author: &'a MessageAuthor,
    /// The session identifier.
    pub session_id: &'a str,
    /// The turn identifier.
    pub turn_id: &'a str,
    /// The model.
    pub model: &'a str,
    /// The approval policy.
    pub approval_policy: ApprovalPolicy,
}

/// Why [`Middleware::session_start`](super::Middleware::session_start) is running.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SessionStartSource {
    /// Selects the startup case.
    Startup,
    /// Selects the resume case.
    Resume,
    /// Selects the compact case.
    Compact,
}

/// Mutable state shared by the declaration-ordered `SessionStart` hooks.
pub struct SessionStartContext<'a> {
    /// The runtime.
    pub runtime: &'a RuntimeContext,
    pub(crate) source: SessionStartSource,
    pub(crate) queued_messages: QueuedMessageSnapshot,
    pub(crate) input: &'a mut Vec<Value>,
    pub(crate) input_changed: bool,
    pub(crate) stop_reason: Option<String>,
}

impl SessionStartContext<'_> {
    #[must_use]
    /// Returns the session-start source.
    pub fn source(&self) -> SessionStartSource {
        self.source
    }

    #[must_use]
    /// Returns the queued messages.
    pub fn queued_messages(&self) -> &QueuedMessageSnapshot {
        &self.queued_messages
    }

    /// Appends hidden provider context produced while the session starts.
    pub fn push_input(&mut self, item: Value) {
        self.input.push(item);
        self.input_changed = true;
    }

    pub(crate) fn retain_input(&mut self, mut keep: impl FnMut(&Value) -> bool) {
        let input_len = self.input.len();
        self.input.retain(&mut keep);
        self.input_changed |= self.input.len() != input_len;
    }

    /// Stops the active turn after session-start processing completes.
    /// # Errors
    ///
    /// Returns an error if validation or an operation required by this function fails.
    pub fn stop(&mut self, reason: impl Into<String>) -> Result<()> {
        set_stop_reason(&mut self.stop_reason, "session-start stop", reason)
    }

    /// Returns the first stop requested by the ordered middleware chain.
    #[must_use]
    pub fn stop_reason(&self) -> Option<&str> {
        self.stop_reason.as_deref()
    }
}

/// Mutable state exposed before a prepared next-turn message enters durable context.
pub struct MessageSubmitContext<'a> {
    /// The turn.
    pub turn: TurnIdentity<'a>,
    /// The author.
    pub author: &'a MessageAuthor,
    /// The message.
    pub message: &'a str,
    /// The attachments.
    pub attachments: &'a [SessionFileReference],
    /// The events.
    pub events: &'a mut Vec<EventMsg>,
    pub(crate) input: Vec<Value>,
    pub(crate) rejection: Option<String>,
}

impl MessageSubmitContext<'_> {
    /// Adds provider-neutral context immediately before the submitted message.
    pub fn push_input(&mut self, item: Value) {
        self.input.push(item);
    }

    /// Rejects the submission without treating the policy decision as a hook failure.
    /// # Errors
    ///
    /// Returns an error if validation or an operation required by this function fails.
    pub fn reject(&mut self, reason: impl Into<String>) -> Result<()> {
        let reason = hook_message("prompt rejection", reason)?;
        if self.rejection.is_none() {
            self.rejection = Some(reason);
        }
        Ok(())
    }
}

pub(crate) struct MessageSubmitResult {
    pub(crate) input: Vec<Value>,
    pub(crate) rejection: Option<String>,
}

/// Mutable state exposed immediately before a model request.
pub struct ModelContext<'a> {
    /// Trusted provenance of the message that initiated the active turn.
    pub author: &'a MessageAuthor,
    /// The model.
    pub model: &'a ModelRouter,
    /// The provider.
    pub provider: &'a str,
    /// The session identifier.
    pub session_id: &'a str,
    /// The session context.
    pub session_context: &'a SessionContext,
    /// The metadata.
    pub metadata: &'a BTreeMap<String, Value>,
    /// The turn identifier.
    pub turn_id: &'a str,
    /// The model step.
    pub model_step: usize,
    /// The context window.
    pub context_window: i64,
    /// The instructions.
    pub instructions: &'a str,
    pub(crate) checkpoint_sequence: u64,
    pub(crate) available_tools: &'a mut BTreeSet<String>,
    pub(crate) allow_hosted_tools: &'a mut bool,
    pub(crate) durable_input: &'a mut Vec<Value>,
    pub(crate) transcript_delta: &'a mut Vec<Value>,
    pub(crate) context_epoch: &'a mut u64,
    pub(crate) compaction_count: &'a mut u64,
    pub(crate) rewrite_reasons: &'a mut Vec<ContextRewriteReason>,
    pub(crate) turn_stop: &'a mut Option<String>,
    pub(crate) queued_messages: Vec<DurableQueuedMessage>,
    /// The last usage.
    pub last_usage: Option<&'a TokenUsage>,
    /// The tools.
    pub tools: &'a Catalog,
    /// The events.
    pub events: &'a mut Vec<EventMsg>,
    /// The usage.
    pub usage: &'a mut Vec<TokenUsage>,
    /// Set when this hook changes durable checkpoint state.
    pub(crate) checkpoint_changed: &'a mut bool,
    pub(crate) runtime: &'a RuntimeContext,
    pub(crate) hooks: &'a MiddlewareStack,
}

/// Live capability state used to hide registered tools at a model boundary.
pub struct ToolExposureContext<'a> {
    /// The session identifier.
    pub session_id: &'a str,
    pub(crate) supports_tool_image_input: bool,
    pub(crate) input: &'a [Value],
    pub(crate) available: &'a mut BTreeSet<String>,
}

impl ToolExposureContext<'_> {
    /// Reports whether the active model accepts image input.
    #[must_use]
    pub fn supports_tool_image_input(&self) -> bool {
        self.supports_tool_image_input
    }

    /// Returns the most recent typed conversation message in model context.
    #[must_use]
    pub fn latest_message(&self) -> Option<MessageEvent> {
        self.input.iter().rev().find_map(message_metadata)
    }

    /// Hides registered tools for this boundary.
    pub fn hide(&mut self, names: &[&str]) {
        for name in names {
            self.available.remove(*name);
        }
    }
}

impl ModelContext<'_> {
    /// Prevents provider-hosted tools for this model step.
    pub fn disable_hosted_tools(&mut self) {
        *self.allow_hosted_tools = false;
    }

    /// Returns durable provider-neutral model context.
    #[must_use]
    pub fn input(&self) -> &[Value] {
        self.durable_input
    }

    /// Replaces active model context and advances its rewrite epoch once per boundary.
    /// # Errors
    ///
    /// Returns an error if validation or an operation required by this function fails.
    pub fn rewrite_input(
        &mut self,
        reason: ContextRewriteReason,
        mut input: Vec<Value>,
    ) -> Result<()> {
        if *self.durable_input == input {
            return Ok(());
        }
        if self.rewrite_reasons.is_empty() {
            *self.context_epoch = self
                .context_epoch
                .checked_add(1)
                .ok_or_else(|| Error::Checkpoint("context rewrite epoch overflow".into()))?;
        }
        if !self.rewrite_reasons.contains(&reason) {
            self.rewrite_reasons.push(reason);
        }
        crate::backend::model::reset_prompt_cache_breakpoint(&mut input);
        *self.durable_input = input;
        self.last_usage = None;
        *self.checkpoint_changed = true;
        Ok(())
    }

    /// Appends a durable replay item without adding it to provider context.
    pub(crate) fn record_transcript_item(&mut self, item: Value) {
        self.transcript_delta.push(item);
        *self.checkpoint_changed = true;
    }

    /// Appends durable provider context without adding synthetic replay history.
    pub fn append_model_input(&mut self, item: Value) {
        self.durable_input.push(item);
        *self.checkpoint_changed = true;
    }

    /// Appends durable input to model context and its transcript journal.
    /// # Errors
    ///
    /// Returns an error if validation or an operation required by this function fails.
    pub fn push_input(&mut self, item: Value) -> Result<MessageTarget> {
        self.durable_input.push(item.clone());
        self.transcript_delta.push(item);
        *self.checkpoint_changed = true;
        provisional_message_target(self.checkpoint_sequence, self.transcript_delta.len())
    }

    /// Estimates visible history, instructions, and tool schemas at four bytes per token.
    #[must_use]
    pub fn estimated_input_tokens(&self) -> i64 {
        let Ok(tools) = self
            .tools
            .prepare(self.input(), self.available_tools.clone())
        else {
            return i64::MAX;
        };
        let visible = tools
            .direct()
            .iter()
            .chain(
                tools
                    .deferred()
                    .iter()
                    .filter(|tool| tools.materialized().contains(&tool.name)),
            )
            .collect::<Vec<_>>();
        let Some(tool_bytes) = serialized_len(&visible) else {
            return i64::MAX;
        };
        let history = self
            .durable_input
            .iter()
            .map(approximate_item_tokens)
            .fold(0usize, usize::saturating_add);
        i64::try_from(history.saturating_add(approximate_tokens(
            tool_bytes.saturating_add(self.instructions.len()),
        )))
        .unwrap_or(i64::MAX)
    }

    pub(crate) async fn pre_compact(&mut self) -> Result<()> {
        let hooks = self.hooks;
        let stop_reason = hooks
            .pre_compact(CompactContext {
                session_id: self.session_id,
                turn_id: self.turn_id,
                model: &self.runtime.model,
                input: self.durable_input,
                events: self.events,
                stop_reason: None,
            })
            .await?;
        set_first(self.turn_stop, stop_reason);
        Ok(())
    }

    pub(crate) async fn post_compact(&mut self) -> Result<()> {
        let hooks = self.hooks;
        let stop_reason = hooks
            .post_compact(CompactContext {
                session_id: self.session_id,
                turn_id: self.turn_id,
                model: &self.runtime.model,
                input: self.durable_input,
                events: self.events,
                stop_reason: None,
            })
            .await?;
        set_first(self.turn_stop, stop_reason);
        if self.turn_stop.is_some() {
            return Ok(());
        }
        let start = hooks
            .session_start(
                self.runtime,
                &self.queued_messages,
                SessionStartSource::Compact,
                self.durable_input,
            )
            .await?;
        set_first(self.turn_stop, start.stop_reason);
        Ok(())
    }

    #[must_use]
    pub(crate) fn turn_stopped(&self) -> bool {
        self.turn_stop.is_some()
    }
}

/// Request-only model input exposed after every durable `PreModel` hook.
pub struct ModelRequestContext<'a> {
    /// Trusted provenance of the message that initiated the active turn.
    pub author: &'a MessageAuthor,
    /// The role.
    pub role: &'a AgentRole,
    /// The model.
    pub model: &'a ModelRouter,
    /// The provider.
    pub provider: &'a str,
    /// The session identifier.
    pub session_id: &'a str,
    /// The turn identifier.
    pub turn_id: &'a str,
    /// The model step.
    pub model_step: usize,
    pub(crate) input: Cow<'a, [Value]>,
}

impl ModelRequestContext<'_> {
    /// Returns the input currently prepared for this one model request.
    #[must_use]
    pub fn input(&self) -> &[Value] {
        self.input.as_ref()
    }

    /// Replaces only the input sent by this model request.
    pub fn replace_input(&mut self, input: Vec<Value>) {
        self.input = Cow::Owned(input);
    }
}

/// Mutable policy boundary for one normalized model-requested tool call.
pub struct PreToolUseContext<'a> {
    /// The turn.
    pub turn: TurnIdentity<'a>,
    /// The events.
    pub events: &'a mut Vec<EventMsg>,
    pub(crate) tools: &'a Catalog,
    pub(crate) call: &'a mut ToolCall,
    pub(crate) input: Vec<Value>,
    pub(crate) denial: Option<String>,
}

impl PreToolUseContext<'_> {
    /// Returns the call after any earlier middleware rewrites.
    #[must_use]
    pub fn call(&self) -> &ToolCall {
        self.call
    }

    /// Replaces the tool name and arguments while preserving the provider call ID.
    /// # Errors
    ///
    /// Returns an error if validation or an operation required by this function fails.
    pub fn replace(&mut self, name: impl Into<String>, arguments: Value) -> Result<()> {
        self.call.replace(name.into(), arguments)
    }

    /// Adds durable provider-neutral context before this call at a tool-complete boundary.
    pub fn push_input(&mut self, item: Value) {
        self.input.push(item);
    }

    /// Denies the call. Later middleware may observe but cannot undo the denial.
    /// # Errors
    ///
    /// Returns an error if validation or an operation required by this function fails.
    pub fn deny(&mut self, reason: impl Into<String>) -> Result<()> {
        let reason = hook_message("tool denial", reason)?;
        if self.denial.is_none() {
            self.denial = Some(reason);
        }
        Ok(())
    }

    /// Returns the first denial made by the ordered middleware chain.
    #[must_use]
    pub fn denial(&self) -> Option<&str> {
        self.denial.as_deref()
    }
}

/// Mutable policy boundary for a sandbox approval request.
pub struct PermissionRequestContext<'a> {
    /// The turn.
    pub turn: TurnIdentity<'a>,
    /// The calls.
    pub calls: &'a [ToolCall],
    /// The requested call identifiers.
    pub requested_call_ids: &'a [String],
    /// The reason.
    pub reason: &'a str,
    /// The events.
    pub events: &'a mut Vec<EventMsg>,
    pub(crate) tools: &'a Catalog,
    pub(crate) decision: Option<ReviewDecision>,
}

impl PermissionRequestContext<'_> {
    /// Returns the decision accumulated from earlier middleware.
    #[must_use]
    pub fn decision(&self) -> Option<&ReviewDecision> {
        self.decision.as_ref()
    }

    /// Allows this request unless an earlier middleware denied it.
    pub fn allow(&mut self) {
        if !matches!(self.decision, Some(ReviewDecision::Denied { .. })) {
            self.decision = Some(ReviewDecision::Approved);
        }
    }

    /// Denies this request. The decision cannot be weakened by later middleware.
    /// # Errors
    ///
    /// Returns an error if validation or an operation required by this function fails.
    pub fn deny(&mut self, reason: impl Into<String>) -> Result<()> {
        let reason = hook_message("permission denial", reason)?;
        if !matches!(self.decision, Some(ReviewDecision::Denied { .. })) {
            self.decision = Some(ReviewDecision::Denied { rejection: reason });
        }
        Ok(())
    }
}

/// Mutable model-visible result exposed after an executed tool call.
pub struct PostToolUseContext<'a> {
    /// The turn.
    pub turn: TurnIdentity<'a>,
    /// The call.
    pub call: &'a ToolCall,
    /// The events.
    pub events: &'a mut Vec<EventMsg>,
    pub(crate) tools: &'a Catalog,
    pub(crate) result: &'a mut ToolResult,
}

impl PostToolUseContext<'_> {
    /// Returns the result after any earlier middleware changes.
    #[must_use]
    pub fn result(&self) -> &ToolResult {
        self.result
    }

    /// Replaces the feedback returned to the model without changing past side effects.
    pub fn replace(&mut self, output: impl Into<String>) {
        self.result.replace(output.into());
    }

    /// Adds provider-neutral context immediately after this tool output.
    pub fn push_input(&mut self, item: Value) {
        self.result.additional_input.push(item);
    }
}

/// State exposed immediately before or after context compaction.
pub struct CompactContext<'a> {
    /// The session identifier.
    pub session_id: &'a str,
    /// The turn identifier.
    pub turn_id: &'a str,
    /// The model.
    pub model: &'a str,
    /// The input.
    pub input: &'a [Value],
    /// The events.
    pub events: &'a mut Vec<EventMsg>,
    pub(crate) stop_reason: Option<String>,
}

impl CompactContext<'_> {
    /// Stops the active turn at this compaction boundary.
    /// # Errors
    ///
    /// Returns an error if validation or an operation required by this function fails.
    pub fn stop(&mut self, reason: impl Into<String>) -> Result<()> {
        set_stop_reason(&mut self.stop_reason, "compaction stop", reason)
    }

    /// Returns the first stop requested by the ordered middleware chain.
    #[must_use]
    pub fn stop_reason(&self) -> Option<&str> {
        self.stop_reason.as_deref()
    }
}

/// Mutable policy boundary immediately before normal turn completion.
pub struct StopContext<'a> {
    /// The turn.
    pub turn: TurnIdentity<'a>,
    /// The events.
    pub events: &'a mut Vec<EventMsg>,
    pub(crate) role: &'a AgentRole,
    pub(crate) stop_hook_active: bool,
    pub(crate) last_assistant_message: Option<&'a str>,
    pub(crate) continuation: Option<String>,
}

impl StopContext<'_> {
    #[must_use]
    /// Returns the agent role.
    pub fn role(&self) -> &AgentRole {
        self.role
    }

    #[must_use]
    /// Stops hook active.
    pub fn stop_hook_active(&self) -> bool {
        self.stop_hook_active
    }

    #[must_use]
    /// Returns the last assistant message.
    pub fn last_assistant_message(&self) -> Option<&str> {
        self.last_assistant_message
    }

    /// Returns the first continuation requested by the middleware chain.
    #[must_use]
    pub fn continuation(&self) -> Option<&str> {
        self.continuation.as_deref()
    }

    /// Requests one more model step with hidden context.
    /// # Errors
    ///
    /// Returns an error if validation or an operation required by this function fails.
    pub fn continue_with(&mut self, prompt: impl Into<String>) -> Result<()> {
        if self.stop_hook_active {
            return Err(Error::Config(
                "a stop hook may continue a turn only once".into(),
            ));
        }
        let prompt = hook_message("stop continuation prompt", prompt)?;
        if self.continuation.is_none() {
            self.continuation = Some(prompt);
        }
        Ok(())
    }
}

fn hook_message(name: &str, value: impl Into<String>) -> Result<String> {
    let value = value.into();
    if value.trim().is_empty() || value.len() > MAX_CAPABILITY_INPUT_BYTES {
        return Err(Error::Config(format!("{name} is empty or too long")));
    }
    Ok(value)
}

fn set_stop_reason(
    target: &mut Option<String>,
    name: &str,
    reason: impl Into<String>,
) -> Result<()> {
    let reason = hook_message(name, reason)?;
    if target.is_none() {
        *target = Some(reason);
    }
    Ok(())
}

fn set_first(target: &mut Option<String>, value: Option<String>) {
    if target.is_none() {
        *target = value;
    }
}

pub(super) fn provisional_message_target(
    checkpoint_sequence: u64,
    batch_item_count: usize,
) -> Result<MessageTarget> {
    Ok(MessageTarget {
        checkpoint_sequence: checkpoint_sequence
            .checked_add(1)
            .ok_or_else(|| Error::Checkpoint("checkpoint sequence overflow".into()))?,
        batch_item_count,
    })
}

/// Mutable state exposed to the middleware preparing conversation messages.
pub struct MessageRouteContext<'a> {
    /// The submission identifier.
    pub submission_id: &'a str,
    /// The message.
    pub message: &'a MessageSubmission,
    /// The active turn identifier.
    pub active_turn_id: Option<&'a str>,
    /// The queued messages.
    pub queued_messages: MessageQueue<'a>,
    /// The events.
    pub events: &'a mut Vec<EventMsg>,
}

/// Mutable turn state exposed to a capability command that can run immediately.
pub struct ActiveCommandContext<'a> {
    /// The checkpoints.
    pub checkpoints: &'a dyn CheckpointStore,
    /// The submission identifier.
    pub submission_id: &'a str,
    /// The session identifier.
    pub session_id: &'a str,
    /// The metadata.
    pub metadata: &'a BTreeMap<String, Value>,
    /// The active turn identifier.
    pub active_turn_id: &'a str,
    /// The command.
    pub command: &'a str,
    /// The arguments.
    pub arguments: &'a str,
    /// The input.
    pub input: Option<&'a str>,
    /// The target.
    pub target: Option<MessageTarget>,
    /// The queued messages.
    pub queued_messages: MessageQueue<'a>,
    /// The events.
    pub events: &'a mut Vec<EventMsg>,
}

/// Result of one middleware-owned submission.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum SubmissionResult {
    /// Selects the accepted case.
    Accepted {
        /// The input changed.
        input_changed: bool,
    },
    /// The operation completed without changing durable turn state; publish its events now.
    Handled,
    /// Selects the rejected case.
    Rejected(String),
}

/// State exposed when the loop finishes or aborts a turn.
pub struct TurnEndContext<'a> {
    /// The session identifier.
    pub session_id: &'a str,
    /// The turn identifier.
    pub turn_id: &'a str,
    pub(crate) outcome: ExecutionOutcome,
    pub(crate) queued_messages: &'a [DurableQueuedMessage],
    pub(crate) owner: Option<&'static str>,
    /// The events.
    pub events: &'a mut Vec<EventMsg>,
}

impl TurnEndContext<'_> {
    #[must_use]
    /// Returns the turn outcome.
    pub fn outcome(&self) -> ExecutionOutcome {
        self.outcome
    }

    /// Returns queued messages still pending for this middleware, oldest first.
    pub fn queued_messages(&self) -> impl Iterator<Item = QueuedMessageView<'_>> {
        let owner = self.owner;
        self.queued_messages
            .iter()
            .filter(move |item| owner.is_some_and(|owner| item.owner() == owner))
            .map(|item| QueuedMessageView { item })
    }
}

/// State available to a middleware-owned frontend command.
pub struct MiddlewareCommandContext<'a> {
    /// The command.
    pub command: &'a str,
    /// The arguments.
    pub arguments: &'a str,
    /// The input.
    pub input: Option<&'a str>,
    /// The target.
    pub target: Option<MessageTarget>,
    /// The session identifier.
    pub session_id: &'a str,
    /// The session context.
    pub session_context: &'a SessionContext,
    /// The checkpoint.
    pub checkpoint: &'a Checkpoint,
    /// The checkpoints.
    pub checkpoints: Arc<dyn CheckpointStore>,
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::backend::model::{Model, ModelEventSink, ModelOutput, ModelRequest};

    struct NoModel;

    impl Model for NoModel {
        fn respond<'a>(
            &'a self,
            _request: ModelRequest<'a>,
            _events: ModelEventSink,
        ) -> crate::BoxFuture<'a, crate::Result<ModelOutput>> {
            Box::pin(async { Err(crate::Error::Provider("unused".into())) })
        }
    }

    #[test]
    fn request_input_is_borrowed_until_replaced() {
        let original = vec![Value::String("original".into())];
        let role = AgentRole::Main;
        let router = ModelRouter::new("test", Arc::new(NoModel));
        let mut context = ModelRequestContext {
            author: &MessageAuthor::User,
            role: &role,
            model: &router,
            provider: "test",
            session_id: "session",
            turn_id: "turn",
            model_step: 0,
            input: Cow::Borrowed(&original),
        };

        assert!(matches!(&context.input, Cow::Borrowed(_)));
        context.replace_input(vec![Value::String("replacement".into())]);
        assert!(matches!(&context.input, Cow::Owned(_)));
        assert_eq!(original, [Value::String("original".into())]);
    }
}