Skip to main content

mobius/middleware/
context.rs

1use std::collections::{BTreeMap, BTreeSet};
2use std::io::{self, Write};
3use std::sync::Arc;
4
5use serde_json::Value;
6
7use super::MiddlewareStack;
8use super::approximate_tokens;
9use super::tools::Catalog;
10use super::tools::ToolResult;
11use crate::agent::{AgentRole, WeakAgentSender};
12use crate::backend::checkpoint::{
13    Checkpoint, CheckpointStore, ContextRewriteReason, ExecutionOutcome, MAX_QUEUED_MESSAGES,
14    QueuedMessage as DurableQueuedMessage, QueuedMessageBoundary,
15};
16use crate::backend::model::{ModelRouter, ToolCall, message_input};
17use crate::backend::sandbox::ApprovalPolicy;
18use crate::protocol::{
19    EventMsg, FrontendEvent, MAX_CAPABILITY_INPUT_BYTES, MessageAuthor, MessageEvent,
20    MessageSubmission, MessageTarget, ReviewDecision, SessionContext, SessionFileReference,
21    TokenUsage, message_metadata,
22};
23use crate::{Error, Result};
24
25/// Sends middleware-owned UI updates without depending on a concrete frontend.
26pub type FrontendEventSink = Arc<dyn Fn(FrontendEvent) -> Result<()> + Send + Sync>;
27
28/// Read-only queued message owned by the middleware receiving it.
29#[derive(Debug, Clone, Copy, PartialEq)]
30pub struct QueuedMessageView<'a> {
31    item: &'a DurableQueuedMessage,
32}
33
34impl<'a> QueuedMessageView<'a> {
35    /// Returns the identity token required by a conditional queue mutation.
36    #[must_use]
37    pub fn id(&self) -> &'a str {
38        self.item.id()
39    }
40
41    /// Returns the prepared presentation event.
42    #[must_use]
43    pub fn event(&self) -> MessageEvent {
44        self.item.event()
45    }
46}
47
48/// Read-only startup snapshot containing only one middleware's queued messages.
49#[derive(Clone, Default)]
50pub struct QueuedMessageSnapshot {
51    items: Vec<DurableQueuedMessage>,
52}
53
54impl QueuedMessageSnapshot {
55    /// Returns every queued item owned by this middleware, oldest first.
56    pub fn views(&self) -> impl Iterator<Item = QueuedMessageView<'_>> {
57        self.items.iter().map(|item| QueuedMessageView { item })
58    }
59
60    pub(super) fn for_owner(owner: &str, items: &[DurableQueuedMessage]) -> Self {
61        Self {
62            items: items
63                .iter()
64                .filter(|item| item.owner() == owner)
65                .cloned()
66                .collect(),
67        }
68    }
69}
70
71/// Mutable scoped view of messages retained until their delivery boundary.
72pub struct MessageQueue<'a> {
73    items: &'a mut Vec<DurableQueuedMessage>,
74    owner: Option<&'static str>,
75}
76
77impl<'a> MessageQueue<'a> {
78    pub(crate) fn new(items: &'a mut Vec<DurableQueuedMessage>) -> Self {
79        Self { items, owner: None }
80    }
81
82    pub(super) fn scope(&mut self, owner: &'static str) {
83        self.owner = Some(owner);
84    }
85
86    fn owner(&self) -> Result<&'static str> {
87        self.owner
88            .ok_or_else(|| Error::Config("message queue is not scoped to a middleware".into()))
89    }
90
91    /// Returns the number of queued items owned by this middleware.
92    #[must_use]
93    pub fn count(&self) -> usize {
94        let Some(owner) = self.owner else {
95            return 0;
96        };
97        self.items
98            .iter()
99            .filter(|item| item.owner() == owner)
100            .count()
101    }
102
103    /// Returns the newest message available to this context.
104    #[must_use]
105    pub fn latest(&self) -> Option<QueuedMessageView<'_>> {
106        let owner = self.owner?;
107        self.items
108            .iter()
109            .rev()
110            .find(|item| item.owner() == owner)
111            .map(|item| QueuedMessageView { item })
112    }
113
114    /// Returns one owned item by its revision identity.
115    #[must_use]
116    pub fn find(&self, id: &str) -> Option<QueuedMessageView<'_>> {
117        let owner = self.owner?;
118        self.items
119            .iter()
120            .find(|item| item.owner() == owner && item.id() == id)
121            .map(|item| QueuedMessageView { item })
122    }
123
124    /// Appends one prepared message, or returns `false` when it is full or duplicated.
125    pub fn enqueue(
126        &mut self,
127        id: &str,
128        boundary: QueuedMessageBoundary,
129        event: MessageEvent,
130    ) -> Result<bool> {
131        let owner = self.owner()?;
132        let item = DurableQueuedMessage::new(owner, id, boundary, event)?;
133        if self.items.len() >= MAX_QUEUED_MESSAGES {
134            return Ok(false);
135        }
136        if self
137            .items
138            .iter()
139            .any(|item| item.owner() == owner && item.id() == id)
140        {
141            return Ok(false);
142        }
143        self.items.push(item);
144        Ok(true)
145    }
146
147    /// Atomically replaces one owned item while preserving its queue position.
148    pub fn replace(&mut self, id: &str, replacement_id: &str, event: MessageEvent) -> Result<bool> {
149        let owner = self.owner()?;
150        let Some(index) = self
151            .items
152            .iter()
153            .position(|item| item.owner() == owner && item.id() == id)
154        else {
155            return Ok(false);
156        };
157        if self.items.iter().enumerate().any(|(candidate, item)| {
158            candidate != index && item.owner() == owner && item.id() == replacement_id
159        }) {
160            return Ok(false);
161        }
162        self.items[index].replace(replacement_id, event)?;
163        Ok(true)
164    }
165
166    pub(crate) fn stage_model_messages(&mut self, turn_id: &str) -> Result<Vec<PreparedMessage>> {
167        let Some(owner) = self.owner else {
168            return Ok(Vec::new());
169        };
170        self.items
171            .extract_if(.., |item| {
172                item.owner() == owner
173                    && matches!(
174                        item.boundary(),
175                        QueuedMessageBoundary::Steer { turn_id: target }
176                            if target == turn_id
177                    )
178            })
179            .map(PreparedMessage::try_from)
180            .collect()
181    }
182
183    pub(crate) fn next_turn(&self) -> Result<Option<PreparedMessage>> {
184        let owner = self.owner()?;
185        self.items
186            .iter()
187            .find(|item| item.owner() == owner && item.boundary().starts_turn())
188            .cloned()
189            .map(PreparedMessage::try_from)
190            .transpose()
191    }
192
193    pub(crate) fn consume_next_turn(&mut self, id: &str) -> Result<()> {
194        let owner = self.owner()?;
195        let index = self
196            .items
197            .iter()
198            .position(|item| {
199                item.owner() == owner && item.id() == id && item.boundary().starts_turn()
200            })
201            .ok_or_else(|| Error::Checkpoint("prepared message is no longer queued".into()))?;
202        self.items.remove(index);
203        Ok(())
204    }
205
206    pub(crate) fn promote_failed_turn(&mut self, turn_id: &str) -> Result<()> {
207        let owner = self.owner()?;
208        for item in self.items.iter_mut().filter(|item| {
209            item.owner() == owner
210                && matches!(
211                    item.boundary(),
212                    QueuedMessageBoundary::Steer { turn_id: target }
213                        if target == turn_id
214                )
215        }) {
216            item.promote_to_next_turn()?;
217        }
218        Ok(())
219    }
220}
221
222/// One queued message prepared for its model boundary.
223pub(crate) struct PreparedMessage {
224    pub(crate) submission_id: String,
225    pub(crate) input: Value,
226    pub(crate) event: EventMsg,
227    pub(crate) title_seed: Option<String>,
228    pub(crate) boundary_events: Vec<EventMsg>,
229}
230
231impl TryFrom<DurableQueuedMessage> for PreparedMessage {
232    type Error = Error;
233
234    fn try_from(message: DurableQueuedMessage) -> Result<Self> {
235        let (submission_id, event) = message.into_parts();
236        let input = message_input(&event)?;
237        let title_seed = matches!(
238            event.author,
239            MessageAuthor::User | MessageAuthor::Peer { .. }
240        )
241        .then(|| event.text.trim().to_string())
242        .filter(|title| !title.is_empty());
243        Ok(Self {
244            submission_id,
245            input,
246            event: EventMsg::Message(event),
247            title_seed,
248            boundary_events: Vec::new(),
249        })
250    }
251}
252
253/// Durable runtime identity exposed while middleware starts a session.
254#[derive(Clone)]
255pub struct RuntimeContext {
256    pub sender: WeakAgentSender,
257    pub checkpoints: Arc<dyn CheckpointStore>,
258    pub session_id: String,
259    pub model_route: String,
260    pub model: String,
261    pub approval_policy: ApprovalPolicy,
262    pub session_context: SessionContext,
263    pub metadata: BTreeMap<String, Value>,
264    pub role: AgentRole,
265    pub frontend: FrontendEventSink,
266}
267
268impl RuntimeContext {
269    pub(crate) fn turn_identity<'a>(&'a self, turn_id: &'a str) -> TurnIdentity<'a> {
270        TurnIdentity {
271            session_id: &self.session_id,
272            turn_id,
273            model: &self.model,
274            approval_policy: self.approval_policy,
275        }
276    }
277}
278
279/// Stable facts shared by hooks that run within one active turn.
280#[derive(Debug, Clone, Copy, PartialEq, Eq)]
281pub struct TurnIdentity<'a> {
282    pub session_id: &'a str,
283    pub turn_id: &'a str,
284    pub model: &'a str,
285    pub approval_policy: ApprovalPolicy,
286}
287
288/// Why [`Middleware::session_start`](super::Middleware::session_start) is running.
289#[derive(Debug, Clone, Copy, PartialEq, Eq)]
290pub enum SessionStartSource {
291    Startup,
292    Resume,
293    Compact,
294}
295
296/// Mutable state shared by the declaration-ordered `SessionStart` hooks.
297pub struct SessionStartContext<'a> {
298    pub runtime: &'a RuntimeContext,
299    pub(crate) source: SessionStartSource,
300    pub(crate) queued_messages: QueuedMessageSnapshot,
301    pub(crate) input: &'a mut Vec<Value>,
302    pub(crate) input_changed: bool,
303    pub(crate) stop_reason: Option<String>,
304}
305
306impl SessionStartContext<'_> {
307    #[must_use]
308    pub fn source(&self) -> SessionStartSource {
309        self.source
310    }
311
312    #[must_use]
313    pub fn queued_messages(&self) -> &QueuedMessageSnapshot {
314        &self.queued_messages
315    }
316
317    /// Appends hidden provider context produced while the session starts.
318    pub fn push_input(&mut self, item: Value) {
319        self.input.push(item);
320        self.input_changed = true;
321    }
322
323    pub(crate) fn retain_input(&mut self, mut keep: impl FnMut(&Value) -> bool) {
324        let input_len = self.input.len();
325        self.input.retain(&mut keep);
326        self.input_changed |= self.input.len() != input_len;
327    }
328
329    /// Stops the active turn after session-start processing completes.
330    pub fn stop(&mut self, reason: impl Into<String>) -> Result<()> {
331        set_stop_reason(&mut self.stop_reason, "session-start stop", reason)
332    }
333
334    /// Returns the first stop requested by the ordered middleware chain.
335    #[must_use]
336    pub fn stop_reason(&self) -> Option<&str> {
337        self.stop_reason.as_deref()
338    }
339}
340
341/// Mutable state exposed before a prepared next-turn message enters durable context.
342pub struct MessageSubmitContext<'a> {
343    pub turn: TurnIdentity<'a>,
344    pub author: &'a MessageAuthor,
345    pub message: &'a str,
346    pub attachments: &'a [SessionFileReference],
347    pub events: &'a mut Vec<EventMsg>,
348    pub(crate) input: Vec<Value>,
349    pub(crate) rejection: Option<String>,
350}
351
352impl MessageSubmitContext<'_> {
353    /// Adds provider-neutral context immediately before the submitted message.
354    pub fn push_input(&mut self, item: Value) {
355        self.input.push(item);
356    }
357
358    /// Rejects the submission without treating the policy decision as a hook failure.
359    pub fn reject(&mut self, reason: impl Into<String>) -> Result<()> {
360        let reason = hook_message("prompt rejection", reason)?;
361        if self.rejection.is_none() {
362            self.rejection = Some(reason);
363        }
364        Ok(())
365    }
366}
367
368pub(crate) struct MessageSubmitResult {
369    pub(crate) input: Vec<Value>,
370    pub(crate) rejection: Option<String>,
371}
372
373/// Mutable state exposed immediately before a model request.
374pub struct ModelContext<'a> {
375    pub model: &'a ModelRouter,
376    pub provider: &'a str,
377    pub session_id: &'a str,
378    pub session_context: &'a SessionContext,
379    pub metadata: &'a BTreeMap<String, Value>,
380    pub turn_id: &'a str,
381    pub model_step: usize,
382    pub context_window: i64,
383    pub instructions: &'a str,
384    pub(crate) checkpoint_sequence: u64,
385    pub(crate) request_input: &'a mut Vec<Value>,
386    pub(crate) available_tools: &'a mut BTreeSet<String>,
387    pub(crate) durable_input: &'a mut Vec<Value>,
388    pub(crate) transcript_delta: &'a mut Vec<Value>,
389    pub(crate) context_epoch: &'a mut u64,
390    pub(crate) compaction_count: &'a mut u64,
391    pub(crate) rewrite_reasons: &'a mut Vec<ContextRewriteReason>,
392    pub(crate) turn_stop: &'a mut Option<String>,
393    pub(crate) queued_messages: Vec<DurableQueuedMessage>,
394    pub last_usage: Option<&'a TokenUsage>,
395    pub tools: &'a Catalog,
396    pub events: &'a mut Vec<EventMsg>,
397    pub usage: &'a mut Vec<TokenUsage>,
398    /// Set when this hook changes durable checkpoint state.
399    pub(crate) checkpoint_changed: &'a mut bool,
400    pub(crate) runtime: &'a RuntimeContext,
401    pub(crate) hooks: &'a MiddlewareStack,
402}
403
404/// Live capability state used to hide registered tools at a model boundary.
405pub struct ToolExposureContext<'a> {
406    pub session_id: &'a str,
407    pub(crate) input: &'a [Value],
408    pub(crate) available: &'a mut BTreeSet<String>,
409}
410
411impl ToolExposureContext<'_> {
412    /// Returns the most recent typed conversation message in model context.
413    #[must_use]
414    pub fn latest_message(&self) -> Option<MessageEvent> {
415        self.input.iter().rev().find_map(message_metadata)
416    }
417
418    /// Hides registered tools for this boundary.
419    pub fn hide(&mut self, names: &[&str]) {
420        for name in names {
421            self.available.remove(*name);
422        }
423    }
424}
425
426impl ModelContext<'_> {
427    /// Returns durable provider-neutral model context.
428    #[must_use]
429    pub fn input(&self) -> &[Value] {
430        self.durable_input
431    }
432
433    /// Returns the request input including earlier request-only middleware additions.
434    #[must_use]
435    pub fn request_input(&self) -> &[Value] {
436        self.request_input
437    }
438
439    /// Replaces active model context and advances its rewrite epoch once per boundary.
440    pub fn rewrite_input(&mut self, reason: ContextRewriteReason, input: Vec<Value>) -> Result<()> {
441        if *self.durable_input == input {
442            return Ok(());
443        }
444        if self.rewrite_reasons.is_empty() {
445            *self.context_epoch = self
446                .context_epoch
447                .checked_add(1)
448                .ok_or_else(|| Error::Checkpoint("context rewrite epoch overflow".into()))?;
449        }
450        if !self.rewrite_reasons.contains(&reason) {
451            self.rewrite_reasons.push(reason);
452        }
453        self.durable_input.clone_from(&input);
454        *self.request_input = input;
455        *self.checkpoint_changed = true;
456        Ok(())
457    }
458
459    /// Appends a durable replay item without adding it to provider context.
460    pub(crate) fn record_transcript_item(&mut self, item: Value) {
461        self.transcript_delta.push(item);
462        *self.checkpoint_changed = true;
463    }
464
465    /// Appends durable provider context without adding synthetic replay history.
466    pub fn append_model_input(&mut self, item: Value) {
467        self.request_input.push(item.clone());
468        self.durable_input.push(item);
469        *self.checkpoint_changed = true;
470    }
471
472    /// Appends durable input to model context and its transcript journal.
473    pub fn push_input(&mut self, item: Value) -> Result<MessageTarget> {
474        self.request_input.push(item.clone());
475        self.durable_input.push(item.clone());
476        self.transcript_delta.push(item);
477        *self.checkpoint_changed = true;
478        provisional_message_target(self.checkpoint_sequence, self.transcript_delta.len())
479    }
480
481    /// Estimates serialized model input at four bytes per token.
482    #[must_use]
483    pub fn estimated_input_tokens(&self) -> i64 {
484        let mut bytes = ByteCounter::default();
485        if serde_json::to_writer(&mut bytes, self.durable_input).is_err() {
486            return i64::MAX;
487        }
488        i64::try_from(approximate_tokens(bytes.0)).unwrap_or(i64::MAX)
489    }
490
491    pub(crate) async fn pre_compact(&mut self) -> Result<()> {
492        let hooks = self.hooks;
493        let stop_reason = hooks
494            .pre_compact(CompactContext {
495                session_id: self.session_id,
496                turn_id: self.turn_id,
497                model: &self.runtime.model,
498                input: self.durable_input,
499                events: self.events,
500                stop_reason: None,
501            })
502            .await?;
503        set_first(self.turn_stop, stop_reason);
504        Ok(())
505    }
506
507    pub(crate) async fn post_compact(&mut self) -> Result<()> {
508        let hooks = self.hooks;
509        let stop_reason = hooks
510            .post_compact(CompactContext {
511                session_id: self.session_id,
512                turn_id: self.turn_id,
513                model: &self.runtime.model,
514                input: self.durable_input,
515                events: self.events,
516                stop_reason: None,
517            })
518            .await?;
519        set_first(self.turn_stop, stop_reason);
520        if self.turn_stop.is_some() {
521            return Ok(());
522        }
523        let start = hooks
524            .session_start(
525                self.runtime,
526                &self.queued_messages,
527                SessionStartSource::Compact,
528                self.durable_input,
529            )
530            .await?;
531        set_first(self.turn_stop, start.stop_reason);
532        self.request_input.clone_from(self.durable_input);
533        Ok(())
534    }
535
536    #[must_use]
537    pub(crate) fn turn_stopped(&self) -> bool {
538        self.turn_stop.is_some()
539    }
540}
541
542/// Request-only model input exposed after every durable `PreModel` hook.
543pub struct ModelRequestContext<'a> {
544    pub model: &'a ModelRouter,
545    pub provider: &'a str,
546    pub session_id: &'a str,
547    pub turn_id: &'a str,
548    pub model_step: usize,
549    pub(crate) input: &'a mut Vec<Value>,
550}
551
552impl ModelRequestContext<'_> {
553    /// Returns the input currently prepared for this one model request.
554    #[must_use]
555    pub fn input(&self) -> &[Value] {
556        self.input
557    }
558
559    /// Replaces only the input sent by this model request.
560    pub fn replace_input(&mut self, input: Vec<Value>) {
561        *self.input = input;
562    }
563}
564
565/// Mutable policy boundary for one normalized model-requested tool call.
566pub struct PreToolUseContext<'a> {
567    pub turn: TurnIdentity<'a>,
568    pub events: &'a mut Vec<EventMsg>,
569    pub(crate) tools: &'a Catalog,
570    pub(crate) call: &'a mut ToolCall,
571    pub(crate) input: Vec<Value>,
572    pub(crate) denial: Option<String>,
573}
574
575impl PreToolUseContext<'_> {
576    /// Returns the call after any earlier middleware rewrites.
577    #[must_use]
578    pub fn call(&self) -> &ToolCall {
579        self.call
580    }
581
582    /// Replaces the tool name and arguments while preserving the provider call ID.
583    pub fn replace(&mut self, name: impl Into<String>, arguments: Value) -> Result<()> {
584        self.call.replace(name.into(), arguments)
585    }
586
587    /// Adds durable provider-neutral context before this call at a tool-complete boundary.
588    pub fn push_input(&mut self, item: Value) {
589        self.input.push(item);
590    }
591
592    /// Denies the call. Later middleware may observe but cannot undo the denial.
593    pub fn deny(&mut self, reason: impl Into<String>) -> Result<()> {
594        let reason = hook_message("tool denial", reason)?;
595        if self.denial.is_none() {
596            self.denial = Some(reason);
597        }
598        Ok(())
599    }
600
601    /// Returns the first denial made by the ordered middleware chain.
602    #[must_use]
603    pub fn denial(&self) -> Option<&str> {
604        self.denial.as_deref()
605    }
606}
607
608/// Mutable policy boundary for a sandbox approval request.
609pub struct PermissionRequestContext<'a> {
610    pub turn: TurnIdentity<'a>,
611    pub calls: &'a [ToolCall],
612    pub requested_call_ids: &'a [String],
613    pub reason: &'a str,
614    pub events: &'a mut Vec<EventMsg>,
615    pub(crate) tools: &'a Catalog,
616    pub(crate) decision: Option<ReviewDecision>,
617}
618
619impl PermissionRequestContext<'_> {
620    /// Returns the decision accumulated from earlier middleware.
621    #[must_use]
622    pub fn decision(&self) -> Option<&ReviewDecision> {
623        self.decision.as_ref()
624    }
625
626    /// Allows this request unless an earlier middleware denied it.
627    pub fn allow(&mut self) {
628        if !matches!(self.decision, Some(ReviewDecision::Denied { .. })) {
629            self.decision = Some(ReviewDecision::Approved);
630        }
631    }
632
633    /// Denies this request. The decision cannot be weakened by later middleware.
634    pub fn deny(&mut self, reason: impl Into<String>) -> Result<()> {
635        let reason = hook_message("permission denial", reason)?;
636        if !matches!(self.decision, Some(ReviewDecision::Denied { .. })) {
637            self.decision = Some(ReviewDecision::Denied { rejection: reason });
638        }
639        Ok(())
640    }
641}
642
643/// Mutable model-visible result exposed after an executed tool call.
644pub struct PostToolUseContext<'a> {
645    pub turn: TurnIdentity<'a>,
646    pub call: &'a ToolCall,
647    pub events: &'a mut Vec<EventMsg>,
648    pub(crate) tools: &'a Catalog,
649    pub(crate) result: &'a mut ToolResult,
650}
651
652impl PostToolUseContext<'_> {
653    /// Returns the result after any earlier middleware changes.
654    #[must_use]
655    pub fn result(&self) -> &ToolResult {
656        self.result
657    }
658
659    /// Replaces the feedback returned to the model without changing past side effects.
660    pub fn replace(&mut self, output: impl Into<String>) {
661        self.result.replace(output.into());
662    }
663
664    /// Adds provider-neutral context immediately after this tool output.
665    pub fn push_input(&mut self, item: Value) {
666        self.result.additional_input.push(item);
667    }
668}
669
670/// State exposed immediately before or after context compaction.
671pub struct CompactContext<'a> {
672    pub session_id: &'a str,
673    pub turn_id: &'a str,
674    pub model: &'a str,
675    pub input: &'a [Value],
676    pub events: &'a mut Vec<EventMsg>,
677    pub(crate) stop_reason: Option<String>,
678}
679
680impl CompactContext<'_> {
681    /// Stops the active turn at this compaction boundary.
682    pub fn stop(&mut self, reason: impl Into<String>) -> Result<()> {
683        set_stop_reason(&mut self.stop_reason, "compaction stop", reason)
684    }
685
686    /// Returns the first stop requested by the ordered middleware chain.
687    #[must_use]
688    pub fn stop_reason(&self) -> Option<&str> {
689        self.stop_reason.as_deref()
690    }
691}
692
693/// Mutable policy boundary immediately before normal turn completion.
694pub struct StopContext<'a> {
695    pub turn: TurnIdentity<'a>,
696    pub events: &'a mut Vec<EventMsg>,
697    pub(crate) role: &'a AgentRole,
698    pub(crate) stop_hook_active: bool,
699    pub(crate) last_assistant_message: Option<&'a str>,
700    pub(crate) continuation: Option<String>,
701}
702
703impl StopContext<'_> {
704    #[must_use]
705    pub fn role(&self) -> &AgentRole {
706        self.role
707    }
708
709    #[must_use]
710    pub fn stop_hook_active(&self) -> bool {
711        self.stop_hook_active
712    }
713
714    #[must_use]
715    pub fn last_assistant_message(&self) -> Option<&str> {
716        self.last_assistant_message
717    }
718
719    /// Returns the first continuation requested by the middleware chain.
720    #[must_use]
721    pub fn continuation(&self) -> Option<&str> {
722        self.continuation.as_deref()
723    }
724
725    /// Requests one more model step with hidden context.
726    pub fn continue_with(&mut self, prompt: impl Into<String>) -> Result<()> {
727        if self.stop_hook_active {
728            return Err(Error::Config(
729                "a stop hook may continue a turn only once".into(),
730            ));
731        }
732        let prompt = hook_message("stop continuation prompt", prompt)?;
733        if self.continuation.is_none() {
734            self.continuation = Some(prompt);
735        }
736        Ok(())
737    }
738}
739
740fn hook_message(name: &str, value: impl Into<String>) -> Result<String> {
741    let value = value.into();
742    if value.trim().is_empty() || value.len() > MAX_CAPABILITY_INPUT_BYTES {
743        return Err(Error::Config(format!("{name} is empty or too long")));
744    }
745    Ok(value)
746}
747
748fn set_stop_reason(
749    target: &mut Option<String>,
750    name: &str,
751    reason: impl Into<String>,
752) -> Result<()> {
753    let reason = hook_message(name, reason)?;
754    if target.is_none() {
755        *target = Some(reason);
756    }
757    Ok(())
758}
759
760fn set_first(target: &mut Option<String>, value: Option<String>) {
761    if target.is_none() {
762        *target = value;
763    }
764}
765
766pub(super) fn provisional_message_target(
767    checkpoint_sequence: u64,
768    batch_item_count: usize,
769) -> Result<MessageTarget> {
770    Ok(MessageTarget {
771        checkpoint_sequence: checkpoint_sequence
772            .checked_add(1)
773            .ok_or_else(|| Error::Checkpoint("checkpoint sequence overflow".into()))?,
774        batch_item_count,
775    })
776}
777
778#[derive(Default)]
779struct ByteCounter(usize);
780
781impl Write for ByteCounter {
782    fn write(&mut self, buffer: &[u8]) -> io::Result<usize> {
783        self.0 = self.0.saturating_add(buffer.len());
784        Ok(buffer.len())
785    }
786
787    fn flush(&mut self) -> io::Result<()> {
788        Ok(())
789    }
790}
791
792/// Mutable state exposed to the middleware preparing conversation messages.
793pub struct MessageRouteContext<'a> {
794    pub submission_id: &'a str,
795    pub message: &'a MessageSubmission,
796    pub active_turn_id: Option<&'a str>,
797    pub queued_messages: MessageQueue<'a>,
798    pub events: &'a mut Vec<EventMsg>,
799}
800
801/// Mutable turn state exposed to a capability command that can run immediately.
802pub struct ActiveCommandContext<'a> {
803    pub submission_id: &'a str,
804    pub session_id: &'a str,
805    pub metadata: &'a BTreeMap<String, Value>,
806    pub active_turn_id: &'a str,
807    pub command: &'a str,
808    pub arguments: &'a str,
809    pub input: Option<&'a str>,
810    pub target: Option<MessageTarget>,
811    pub queued_messages: MessageQueue<'a>,
812    pub events: &'a mut Vec<EventMsg>,
813}
814
815/// Result of one middleware-owned submission.
816#[derive(Debug, Clone, PartialEq, Eq)]
817pub enum SubmissionResult {
818    Accepted {
819        input_changed: bool,
820    },
821    /// The operation completed without changing durable turn state; publish its events now.
822    Handled,
823    Rejected(String),
824}
825
826/// State exposed when the loop finishes or aborts a turn.
827pub struct TurnEndContext<'a> {
828    pub session_id: &'a str,
829    pub turn_id: &'a str,
830    pub(crate) outcome: ExecutionOutcome,
831    pub(crate) queued_messages: &'a [DurableQueuedMessage],
832    pub(crate) owner: Option<&'static str>,
833    pub events: &'a mut Vec<EventMsg>,
834}
835
836impl TurnEndContext<'_> {
837    #[must_use]
838    pub fn outcome(&self) -> ExecutionOutcome {
839        self.outcome
840    }
841
842    /// Returns queued messages still pending for this middleware, oldest first.
843    pub fn queued_messages(&self) -> impl Iterator<Item = QueuedMessageView<'_>> {
844        let owner = self.owner;
845        self.queued_messages
846            .iter()
847            .filter(move |item| owner.is_some_and(|owner| item.owner() == owner))
848            .map(|item| QueuedMessageView { item })
849    }
850}
851
852/// State available to a middleware-owned frontend command.
853pub struct MiddlewareCommandContext<'a> {
854    pub command: &'a str,
855    pub arguments: &'a str,
856    pub input: Option<&'a str>,
857    pub target: Option<MessageTarget>,
858    pub session_id: &'a str,
859    pub session_context: &'a SessionContext,
860    pub checkpoint: &'a Checkpoint,
861    pub checkpoints: Arc<dyn CheckpointStore>,
862}