a3s_code_core/
run_control.rs

1//! Typed, cooperative control for an in-flight Code run.
2//!
3//! A run is deliberately controlled through a small per-run inbox instead of
4//! mutating loop state from the host thread.  This keeps the execution loop as
5//! the sole owner of its transcript while still allowing an embedding host to
6//! steer or interrupt work at well-defined safe points.  The wire types in
7//! this module are also used by the language SDKs and can be persisted by a
8//! host without depending on Rust internals.
9
10use serde::{Deserialize, Serialize};
11use sha2::{Digest, Sha256};
12use std::collections::{HashMap, VecDeque};
13use std::sync::Arc;
14use thiserror::Error;
15use tokio::sync::{Mutex, Notify};
16use tokio_util::sync::CancellationToken;
17
18use crate::hooks::{HookExecutor, HookOutcome};
19
20/// Schema carried by a run-control request.
21pub const RUN_CONTROL_REQUEST_SCHEMA_V1: &str = "a3s.code.run-control-request.v1";
22/// Schema carried by a run-control receipt.
23pub const RUN_CONTROL_RECEIPT_SCHEMA_V1: &str = "a3s.code.run-control-receipt.v1";
24/// Maximum UTF-8 bytes accepted for a single steer message.
25pub const RUN_CONTROL_MAX_INPUT_BYTES: usize = 128 * 1024;
26/// Maximum UTF-8 bytes accepted for an interrupt reason.
27pub const RUN_CONTROL_MAX_REASON_BYTES: usize = 4 * 1024;
28/// Maximum number of controls waiting at one run safe point.
29pub const RUN_CONTROL_MAX_QUEUE: usize = 64;
30/// Number of request receipts retained for idempotent retries.
31pub const RUN_CONTROL_MAX_SEEN_REQUESTS: usize = 256;
32/// Maximum size of an externally supplied identifier.
33pub const RUN_CONTROL_MAX_ID_BYTES: usize = 512;
34
35fn default_request_schema() -> String {
36    RUN_CONTROL_REQUEST_SCHEMA_V1.to_string()
37}
38
39fn default_receipt_schema() -> String {
40    RUN_CONTROL_RECEIPT_SCHEMA_V1.to_string()
41}
42
43/// The operation requested by a host.
44#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
45#[serde(rename_all = "snake_case")]
46pub enum RunControlOperation {
47    /// Append a user-directed steering message at the next loop safe point.
48    Steer,
49    /// Cooperatively stop the active run.
50    Interrupt,
51}
52
53/// A typed control command.
54#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
55#[serde(tag = "kind", rename_all = "snake_case")]
56pub enum RunControlCommand {
57    /// Add a new user direction without starting a second turn.
58    Steer { input: String },
59    /// Stop the current run. `force` is advisory: the runtime remains
60    /// cooperative and will never skip cleanup or governance boundaries.
61    Interrupt {
62        reason: Option<String>,
63        #[serde(default)]
64        force: bool,
65    },
66}
67
68impl RunControlCommand {
69    pub fn operation(&self) -> RunControlOperation {
70        match self {
71            Self::Steer { .. } => RunControlOperation::Steer,
72            Self::Interrupt { .. } => RunControlOperation::Interrupt,
73        }
74    }
75}
76
77/// Versioned request accepted by a run-control inbox.
78#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
79#[serde(deny_unknown_fields)]
80pub struct RunControlRequest {
81    /// Protocol schema identifier.
82    #[serde(default = "default_request_schema")]
83    pub schema: String,
84    /// Idempotency key generated by the host.
85    pub request_id: String,
86    /// Optional session binding. When present it must match the target run.
87    #[serde(default, skip_serializing_if = "Option::is_none")]
88    pub session_id: Option<String>,
89    /// Immutable target run id.
90    pub run_id: String,
91    /// Expected logical turn. A mismatch is rejected as stale input.
92    #[serde(default, skip_serializing_if = "Option::is_none")]
93    pub expected_turn_id: Option<String>,
94    /// Monotonic control revision observed by the host.
95    #[serde(default, skip_serializing_if = "Option::is_none")]
96    pub expected_turn_revision: Option<u64>,
97    /// Control payload.
98    pub command: RunControlCommand,
99    /// Optional host deadline in Unix milliseconds.
100    #[serde(default, skip_serializing_if = "Option::is_none")]
101    pub deadline_ms: Option<u64>,
102}
103
104impl RunControlRequest {
105    /// Build a request with a fresh id. The target run is filled by the
106    /// session convenience APIs when omitted here.
107    pub fn new(run_id: impl Into<String>, command: RunControlCommand) -> Self {
108        Self {
109            schema: RUN_CONTROL_REQUEST_SCHEMA_V1.to_string(),
110            request_id: uuid::Uuid::new_v4().to_string(),
111            session_id: None,
112            run_id: run_id.into(),
113            expected_turn_id: None,
114            expected_turn_revision: None,
115            command,
116            deadline_ms: None,
117        }
118    }
119
120    pub fn steer(run_id: impl Into<String>, input: impl Into<String>) -> Self {
121        Self::new(
122            run_id,
123            RunControlCommand::Steer {
124                input: input.into(),
125            },
126        )
127    }
128
129    pub fn interrupt(run_id: impl Into<String>) -> Self {
130        Self::new(
131            run_id,
132            RunControlCommand::Interrupt {
133                reason: None,
134                force: false,
135            },
136        )
137    }
138
139    pub fn with_session_id(mut self, session_id: impl Into<String>) -> Self {
140        self.session_id = Some(session_id.into());
141        self
142    }
143
144    pub fn with_expected_turn(mut self, turn_id: impl Into<String>, revision: u64) -> Self {
145        self.expected_turn_id = Some(turn_id.into());
146        self.expected_turn_revision = Some(revision);
147        self
148    }
149
150    pub fn with_deadline_ms(mut self, deadline_ms: u64) -> Self {
151        self.deadline_ms = Some(deadline_ms);
152        self
153    }
154
155    /// Validate bounded input and protocol identity before enqueueing.
156    pub fn validate(&self) -> Result<(), RunControlError> {
157        if self.schema != RUN_CONTROL_REQUEST_SCHEMA_V1 {
158            return Err(RunControlError::InvalidRequest(format!(
159                "unsupported schema `{}`",
160                self.schema
161            )));
162        }
163        validate_id("request_id", &self.request_id)?;
164        validate_id("run_id", &self.run_id)?;
165        if let Some(session_id) = &self.session_id {
166            validate_id("session_id", session_id)?;
167        }
168        if let Some(turn_id) = &self.expected_turn_id {
169            validate_id("expected_turn_id", turn_id)?;
170        }
171        match &self.command {
172            RunControlCommand::Steer { input } => {
173                if input.trim().is_empty() {
174                    return Err(RunControlError::InvalidRequest(
175                        "steer input must not be empty".to_string(),
176                    ));
177                }
178                if input.len() > RUN_CONTROL_MAX_INPUT_BYTES {
179                    return Err(RunControlError::InvalidRequest(format!(
180                        "steer input exceeds {} bytes",
181                        RUN_CONTROL_MAX_INPUT_BYTES
182                    )));
183                }
184            }
185            RunControlCommand::Interrupt { reason, .. } => {
186                if let Some(reason) = reason {
187                    if reason.len() > RUN_CONTROL_MAX_REASON_BYTES {
188                        return Err(RunControlError::InvalidRequest(format!(
189                            "interrupt reason exceeds {} bytes",
190                            RUN_CONTROL_MAX_REASON_BYTES
191                        )));
192                    }
193                }
194            }
195        }
196        Ok(())
197    }
198}
199
200fn validate_id(name: &str, value: &str) -> Result<(), RunControlError> {
201    if value.trim().is_empty() {
202        return Err(RunControlError::InvalidRequest(format!(
203            "{name} must not be empty"
204        )));
205    }
206    if value.len() > RUN_CONTROL_MAX_ID_BYTES {
207        return Err(RunControlError::InvalidRequest(format!(
208            "{name} exceeds {RUN_CONTROL_MAX_ID_BYTES} bytes"
209        )));
210    }
211    if value
212        .chars()
213        .any(|character| character == '\0' || character == '\r' || character == '\n')
214    {
215        return Err(RunControlError::InvalidRequest(format!(
216            "{name} contains a control character"
217        )));
218    }
219    Ok(())
220}
221
222/// Convenience input for [`crate::AgentSession::steer`](crate::AgentSession::steer).
223#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
224#[serde(deny_unknown_fields)]
225pub struct SteerRequest {
226    pub input: String,
227    #[serde(default, skip_serializing_if = "Option::is_none")]
228    pub request_id: Option<String>,
229    #[serde(default, skip_serializing_if = "Option::is_none")]
230    pub run_id: Option<String>,
231    #[serde(default, skip_serializing_if = "Option::is_none")]
232    pub expected_turn_id: Option<String>,
233    #[serde(default, skip_serializing_if = "Option::is_none")]
234    pub expected_turn_revision: Option<u64>,
235    #[serde(default, skip_serializing_if = "Option::is_none")]
236    pub deadline_ms: Option<u64>,
237}
238
239impl SteerRequest {
240    pub fn new(input: impl Into<String>) -> Self {
241        Self {
242            input: input.into(),
243            request_id: None,
244            run_id: None,
245            expected_turn_id: None,
246            expected_turn_revision: None,
247            deadline_ms: None,
248        }
249    }
250
251    pub fn with_run_id(mut self, run_id: impl Into<String>) -> Self {
252        self.run_id = Some(run_id.into());
253        self
254    }
255
256    pub fn with_expected_turn(mut self, turn_id: impl Into<String>, revision: u64) -> Self {
257        self.expected_turn_id = Some(turn_id.into());
258        self.expected_turn_revision = Some(revision);
259        self
260    }
261
262    pub(crate) fn into_protocol(self, session_id: &str, active_run_id: &str) -> RunControlRequest {
263        RunControlRequest {
264            schema: RUN_CONTROL_REQUEST_SCHEMA_V1.to_string(),
265            request_id: self
266                .request_id
267                .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()),
268            session_id: Some(session_id.to_string()),
269            run_id: self.run_id.unwrap_or_else(|| active_run_id.to_string()),
270            expected_turn_id: self.expected_turn_id,
271            expected_turn_revision: self.expected_turn_revision,
272            command: RunControlCommand::Steer { input: self.input },
273            deadline_ms: self.deadline_ms,
274        }
275    }
276}
277
278/// Convenience input for [`crate::AgentSession::interrupt`](crate::AgentSession::interrupt).
279#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
280#[serde(deny_unknown_fields)]
281pub struct InterruptRequest {
282    #[serde(default, skip_serializing_if = "Option::is_none")]
283    pub reason: Option<String>,
284    #[serde(default)]
285    pub force: bool,
286    #[serde(default, skip_serializing_if = "Option::is_none")]
287    pub request_id: Option<String>,
288    #[serde(default, skip_serializing_if = "Option::is_none")]
289    pub run_id: Option<String>,
290    #[serde(default, skip_serializing_if = "Option::is_none")]
291    pub expected_turn_id: Option<String>,
292    #[serde(default, skip_serializing_if = "Option::is_none")]
293    pub expected_turn_revision: Option<u64>,
294    #[serde(default, skip_serializing_if = "Option::is_none")]
295    pub deadline_ms: Option<u64>,
296}
297
298impl InterruptRequest {
299    pub fn new() -> Self {
300        Self {
301            reason: None,
302            force: false,
303            request_id: None,
304            run_id: None,
305            expected_turn_id: None,
306            expected_turn_revision: None,
307            deadline_ms: None,
308        }
309    }
310
311    pub fn with_reason(mut self, reason: impl Into<String>) -> Self {
312        self.reason = Some(reason.into());
313        self
314    }
315
316    pub fn with_run_id(mut self, run_id: impl Into<String>) -> Self {
317        self.run_id = Some(run_id.into());
318        self
319    }
320
321    pub fn with_expected_turn(mut self, turn_id: impl Into<String>, revision: u64) -> Self {
322        self.expected_turn_id = Some(turn_id.into());
323        self.expected_turn_revision = Some(revision);
324        self
325    }
326
327    pub(crate) fn into_protocol(self, session_id: &str, active_run_id: &str) -> RunControlRequest {
328        RunControlRequest {
329            schema: RUN_CONTROL_REQUEST_SCHEMA_V1.to_string(),
330            request_id: self
331                .request_id
332                .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()),
333            session_id: Some(session_id.to_string()),
334            run_id: self.run_id.unwrap_or_else(|| active_run_id.to_string()),
335            expected_turn_id: self.expected_turn_id,
336            expected_turn_revision: self.expected_turn_revision,
337            command: RunControlCommand::Interrupt {
338                reason: self.reason,
339                force: self.force,
340            },
341            deadline_ms: self.deadline_ms,
342        }
343    }
344}
345
346impl Default for InterruptRequest {
347    fn default() -> Self {
348        Self::new()
349    }
350}
351
352/// Receipt state returned by the host boundary.
353#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
354#[serde(rename_all = "snake_case")]
355pub enum RunControlReceiptState {
356    /// The request passed validation and is waiting for a safe point.
357    Accepted,
358    /// The execution loop consumed the request.
359    Applied,
360    /// The request was rejected before it entered the inbox.
361    Rejected,
362    /// The run ended before an accepted request could be applied.
363    Settled,
364}
365
366/// Machine-readable rejection/settlement detail.
367#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
368pub struct RunControlErrorInfo {
369    pub code: String,
370    pub message: String,
371}
372
373/// Durable, idempotent acknowledgement for a control request.
374#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
375#[serde(deny_unknown_fields)]
376pub struct RunControlReceipt {
377    #[serde(default = "default_receipt_schema")]
378    pub schema: String,
379    pub request_id: String,
380    pub session_id: String,
381    pub run_id: String,
382    pub operation: RunControlOperation,
383    pub state: RunControlReceiptState,
384    pub sequence: u64,
385    pub turn_id: Option<String>,
386    pub turn_revision: u64,
387    pub accepted_at_ms: u64,
388    #[serde(default, skip_serializing_if = "Option::is_none")]
389    pub applied_at_ms: Option<u64>,
390    #[serde(default, skip_serializing_if = "Option::is_none")]
391    pub error: Option<RunControlErrorInfo>,
392}
393
394/// Read-only state exposed to SDK/UI clients for optimistic concurrency.
395#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
396pub struct RunControlSnapshot {
397    pub session_id: String,
398    pub run_id: String,
399    pub active: bool,
400    pub turn_id: Option<String>,
401    pub turn_revision: u64,
402    pub queued_controls: usize,
403    pub interrupt_requested: bool,
404    pub last_sequence: u64,
405}
406
407/// Errors returned before a control is accepted.
408#[derive(Debug, Clone, PartialEq, Eq, Error)]
409pub enum RunControlError {
410    #[error("invalid run-control request: {0}")]
411    InvalidRequest(String),
412    #[error("run-control target session does not match the active session")]
413    SessionMismatch,
414    #[error("run-control target run `{run_id}` is not the active run")]
415    RunMismatch { run_id: String },
416    #[error("there is no active run to control")]
417    NoActiveRun,
418    #[error(
419        "stale run-control request: expected turn={expected_turn_id:?}, revision={expected_revision:?}; current turn={actual_turn_id:?}, revision={actual_revision}"
420    )]
421    StaleTurn {
422        expected_turn_id: Option<String>,
423        expected_revision: Option<u64>,
424        actual_turn_id: Option<String>,
425        actual_revision: u64,
426    },
427    #[error("run-control request deadline has expired")]
428    DeadlineExceeded,
429    #[error("run-control inbox is full")]
430    QueueFull,
431    #[error("run-control inbox is closed")]
432    Closed,
433    #[error("request id `{request_id}` was already used for a different command")]
434    DuplicateRequest { request_id: String },
435    #[error("run-control request was denied by a hook: {reason}")]
436    HookDenied { reason: String },
437    #[error("run-control request must be retried after {retry_after_ms} ms: {reason}")]
438    HookRetry { reason: String, retry_after_ms: u64 },
439}
440
441impl RunControlError {
442    pub const fn code(&self) -> &'static str {
443        match self {
444            Self::InvalidRequest(_) => "INVALID_REQUEST",
445            Self::SessionMismatch => "SESSION_MISMATCH",
446            Self::RunMismatch { .. } => "RUN_MISMATCH",
447            Self::NoActiveRun => "NO_ACTIVE_RUN",
448            Self::StaleTurn { .. } => "STALE_TURN",
449            Self::DeadlineExceeded => "DEADLINE_EXCEEDED",
450            Self::QueueFull => "QUEUE_FULL",
451            Self::Closed => "CLOSED",
452            Self::DuplicateRequest { .. } => "DUPLICATE_REQUEST",
453            Self::HookDenied { .. } => "HOOK_DENIED",
454            Self::HookRetry { .. } => "HOOK_RETRY",
455        }
456    }
457}
458
459#[derive(Debug, Clone)]
460pub(crate) struct PendingRunControl {
461    pub(crate) request: RunControlRequest,
462    pub(crate) receipt: RunControlReceipt,
463}
464
465#[derive(Debug)]
466struct SeenRequest {
467    fingerprint: String,
468    /// Keep the immutable request alongside its receipt so shutdown can
469    /// settle controls that have already been drained by the loop but have
470    /// not reached `mark_applied` yet.  Without this, a close racing a safe
471    /// point could leave an `Accepted` receipt forever and never emit the
472    /// terminal post-hook event.
473    request: RunControlRequest,
474    receipt: RunControlReceipt,
475}
476
477#[derive(Debug)]
478struct InboxState {
479    session_id: String,
480    run_id: String,
481    active: bool,
482    closed: bool,
483    turn_id: Option<String>,
484    turn_revision: u64,
485    queue: VecDeque<PendingRunControl>,
486    /// Controls removed from `queue` by the loop and currently being applied.
487    /// They remain owned by the inbox until the loop acknowledges the
488    /// transition, which gives close a complete set to settle.
489    in_flight: HashMap<String, PendingRunControl>,
490    seen: HashMap<String, SeenRequest>,
491    seen_order: VecDeque<String>,
492    last_sequence: u64,
493    interrupt_requested: bool,
494}
495
496/// Internal per-run inbox. It is shared by the public session facade and the
497/// execution loop, but only the loop consumes pending steering messages.
498#[derive(Debug)]
499pub(crate) struct RunControlInbox {
500    run_id: String,
501    state: Mutex<InboxState>,
502    /// Serializes admission and receipt transitions with the loop's safe
503    /// point.  Without this gate two identical concurrent requests could both
504    /// execute a policy Hook, or the loop could emit `applied` before the
505    /// caller-visible `accepted` observation had been recorded.
506    admission: Mutex<()>,
507    notify: Notify,
508    cancellation: CancellationToken,
509    /// Run-frozen governance executor. Keeping this on the inbox prevents a
510    /// later Session hook mutation from changing the authority of an active
511    /// run-control request.
512    hook_executor: Option<Arc<dyn HookExecutor>>,
513}
514
515impl RunControlInbox {
516    #[cfg(test)]
517    pub(crate) fn new(
518        session_id: impl Into<String>,
519        run_id: impl Into<String>,
520        cancellation: CancellationToken,
521    ) -> Arc<Self> {
522        Self::new_with_hook_executor(session_id, run_id, cancellation, None)
523    }
524
525    pub(crate) fn new_with_hook_executor(
526        session_id: impl Into<String>,
527        run_id: impl Into<String>,
528        cancellation: CancellationToken,
529        hook_executor: Option<Arc<dyn HookExecutor>>,
530    ) -> Arc<Self> {
531        let run_id = run_id.into();
532        Arc::new(Self {
533            run_id: run_id.clone(),
534            state: Mutex::new(InboxState {
535                session_id: session_id.into(),
536                run_id,
537                active: true,
538                closed: false,
539                turn_id: None,
540                turn_revision: 0,
541                queue: VecDeque::new(),
542                in_flight: HashMap::new(),
543                seen: HashMap::new(),
544                seen_order: VecDeque::new(),
545                last_sequence: 0,
546                interrupt_requested: false,
547            }),
548            admission: Mutex::new(()),
549            notify: Notify::new(),
550            cancellation,
551            hook_executor,
552        })
553    }
554
555    pub(crate) fn is_cancelled(&self) -> bool {
556        self.cancellation.is_cancelled()
557    }
558
559    #[cfg(test)]
560    pub(crate) fn cancellation(&self) -> CancellationToken {
561        self.cancellation.clone()
562    }
563
564    /// Fast identity accessor used while a session slot is being swapped.
565    /// The run id is immutable after construction, so this avoids exposing
566    /// the mutable inbox state to callers.
567    pub(crate) fn snapshot_run_id(&self) -> String {
568        self.run_id.clone()
569    }
570
571    pub(crate) async fn update_turn(&self, turn: usize) -> RunControlSnapshot {
572        self.update_turn_id(format!("turn-{turn}")).await
573    }
574
575    pub(crate) async fn update_turn_id(&self, turn_id: String) -> RunControlSnapshot {
576        let _admission = self.admission.lock().await;
577        let mut state = self.state.lock().await;
578        if state.turn_id.as_deref() != Some(turn_id.as_str()) {
579            state.turn_id = Some(turn_id);
580            state.turn_revision = state.turn_revision.saturating_add(1);
581        }
582        snapshot(&state)
583    }
584
585    pub(crate) async fn snapshot(&self) -> RunControlSnapshot {
586        let state = self.state.lock().await;
587        snapshot(&state)
588    }
589
590    #[cfg(test)]
591    pub(crate) async fn submit(
592        &self,
593        request: RunControlRequest,
594        now_ms: u64,
595    ) -> Result<RunControlReceipt, RunControlError> {
596        request.validate()?;
597        let _admission = self.admission.lock().await;
598        let (receipt, cancel) = self.submit_locked(request, now_ms).await?;
599        drop(_admission);
600        if cancel {
601            // Cancellation is fired only after the request has been durably
602            // accepted in the inbox, so a caller can safely retry by id.
603            self.cancellation.cancel();
604        }
605        self.notify.notify_waiters();
606        Ok(receipt)
607    }
608
609    /// Admit one request while the caller holds the transition gate.  Keeping
610    /// the state mutation in one helper lets hook-backed admission publish its
611    /// `accepted` observation before the loop can drain the queue.
612    async fn submit_locked(
613        &self,
614        request: RunControlRequest,
615        now_ms: u64,
616    ) -> Result<(RunControlReceipt, bool), RunControlError> {
617        let fingerprint = request_fingerprint(&request)?;
618        let mut state = self.state.lock().await;
619
620        if let Some(previous) = state.seen.get(&request.request_id) {
621            if previous.fingerprint == fingerprint {
622                return Ok((previous.receipt.clone(), false));
623            }
624            return Err(RunControlError::DuplicateRequest {
625                request_id: request.request_id,
626            });
627        }
628        if request
629            .session_id
630            .as_deref()
631            .is_some_and(|id| id != state.session_id)
632        {
633            return Err(RunControlError::SessionMismatch);
634        }
635        if request.run_id != state.run_id {
636            return Err(RunControlError::RunMismatch {
637                run_id: request.run_id,
638            });
639        }
640        if state.closed || !state.active || self.is_cancelled() {
641            return Err(if state.closed {
642                RunControlError::Closed
643            } else {
644                RunControlError::NoActiveRun
645            });
646        }
647        if request
648            .deadline_ms
649            .is_some_and(|deadline| now_ms > deadline)
650        {
651            return Err(RunControlError::DeadlineExceeded);
652        }
653        if (request.expected_turn_id.is_some() && request.expected_turn_id != state.turn_id)
654            || request
655                .expected_turn_revision
656                .is_some_and(|revision| revision != state.turn_revision)
657        {
658            return Err(RunControlError::StaleTurn {
659                expected_turn_id: request.expected_turn_id,
660                expected_revision: request.expected_turn_revision,
661                actual_turn_id: state.turn_id.clone(),
662                actual_revision: state.turn_revision,
663            });
664        }
665        if state.queue.len() >= RUN_CONTROL_MAX_QUEUE {
666            return Err(RunControlError::QueueFull);
667        }
668
669        state.last_sequence = state.last_sequence.saturating_add(1);
670        let receipt = RunControlReceipt {
671            schema: RUN_CONTROL_RECEIPT_SCHEMA_V1.to_string(),
672            request_id: request.request_id.clone(),
673            session_id: state.session_id.clone(),
674            run_id: state.run_id.clone(),
675            operation: request.command.operation(),
676            state: RunControlReceiptState::Accepted,
677            sequence: state.last_sequence,
678            turn_id: state.turn_id.clone(),
679            turn_revision: state.turn_revision,
680            accepted_at_ms: now_ms,
681            applied_at_ms: None,
682            error: None,
683        };
684        let cancel = matches!(request.command, RunControlCommand::Interrupt { .. });
685        if cancel {
686            state.interrupt_requested = true;
687        }
688        state.queue.push_back(PendingRunControl {
689            request: request.clone(),
690            receipt: receipt.clone(),
691        });
692        state.seen.insert(
693            receipt.request_id.clone(),
694            SeenRequest {
695                fingerprint,
696                request,
697                receipt: receipt.clone(),
698            },
699        );
700        state.seen_order.push_back(receipt.request_id.clone());
701        while state.seen_order.len() > RUN_CONTROL_MAX_SEEN_REQUESTS {
702            if let Some(expired) = state.seen_order.pop_front() {
703                state.seen.remove(&expired);
704            }
705        }
706        Ok((receipt, cancel))
707    }
708
709    pub(crate) async fn drain(&self) -> Vec<PendingRunControl> {
710        let _admission = self.admission.lock().await;
711        let mut state = self.state.lock().await;
712        let pending: Vec<_> = state.queue.drain(..).collect();
713        for item in &pending {
714            state
715                .in_flight
716                .insert(item.receipt.request_id.clone(), item.clone());
717        }
718        pending
719    }
720
721    pub(crate) async fn mark_applied(
722        &self,
723        pending: &PendingRunControl,
724        turn_id: Option<String>,
725        turn_revision: u64,
726        now_ms: u64,
727    ) -> RunControlReceipt {
728        let _admission = self.admission.lock().await;
729        let receipt = {
730            let mut state = self.state.lock().await;
731            // Shutdown may win the race after `drain` and settle this item.
732            // Preserve that terminal state and avoid emitting a second post
733            // receipt or applying a control after the run has ended.
734            let Some(seen) = state.seen.get_mut(&pending.receipt.request_id) else {
735                return pending.receipt.clone();
736            };
737            if seen.receipt.state != RunControlReceiptState::Accepted {
738                return seen.receipt.clone();
739            }
740            let mut receipt = seen.receipt.clone();
741            receipt.state = RunControlReceiptState::Applied;
742            receipt.turn_id = turn_id;
743            receipt.turn_revision = turn_revision;
744            receipt.applied_at_ms = Some(now_ms);
745            seen.receipt = receipt.clone();
746            state.in_flight.remove(&receipt.request_id);
747            receipt
748        };
749        drop(_admission);
750        if receipt.state == RunControlReceiptState::Applied {
751            self.record_receipt(&pending.request, &receipt).await;
752        }
753        receipt
754    }
755
756    /// Mark accepted-but-unconsumed controls as settled when a run exits.
757    pub(crate) async fn close(&self, now_ms: u64) {
758        let _admission = self.admission.lock().await;
759        let settled = {
760            let mut state = self.state.lock().await;
761            state.active = false;
762            state.closed = true;
763            let mut settled = Vec::new();
764            state.queue.clear();
765            state.in_flight.clear();
766            // Iterate the retained receipts rather than only the queue: the
767            // loop may have drained a control immediately before close.
768            for seen in state.seen.values_mut() {
769                if seen.receipt.state == RunControlReceiptState::Accepted {
770                    seen.receipt.state = RunControlReceiptState::Settled;
771                    seen.receipt.applied_at_ms = Some(now_ms);
772                    seen.receipt.error = Some(RunControlErrorInfo {
773                        code: "RUN_ENDED".to_string(),
774                        message: "run ended before the control reached a safe point".to_string(),
775                    });
776                    settled.push((seen.request.clone(), seen.receipt.clone()));
777                }
778            }
779            settled
780        };
781        drop(_admission);
782        self.notify.notify_waiters();
783        for (request, receipt) in settled {
784            self.record_receipt(&request, &receipt).await;
785        }
786    }
787
788    pub(crate) async fn deactivate(&self, now_ms: u64) {
789        self.close(now_ms).await;
790    }
791
792    /// Submit a control through the run-frozen governance boundary. The hook
793    /// decision is made before queue admission; accepted, applied, and
794    /// settled receipts are then emitted as observational post events.
795    pub(crate) async fn submit_with_hooks(
796        &self,
797        request: RunControlRequest,
798        now_ms: u64,
799    ) -> Result<RunControlReceipt, RunControlError> {
800        request.validate()?;
801
802        // Serialize the policy decision, queue admission, and accepted
803        // observation.  This is required for both concurrent idempotency and
804        // the accepted-before-applied event ordering guarantee.
805        let _admission = self.admission.lock().await;
806
807        // Idempotent retries must not execute a policy callback twice. This is
808        // especially important for host hooks that charge a budget or create
809        // an approval record as a side effect.
810        if let Some(receipt) = self.known_receipt(&request).await? {
811            return Ok(receipt);
812        }
813
814        if let Some(executor) = &self.hook_executor {
815            match executor.before_run_control(&request).await {
816                HookOutcome::Continue(_) | HookOutcome::Skip => {}
817                outcome => {
818                    let error = hook_outcome_error(outcome);
819                    let rejected =
820                        rejected_receipt(&request, &self.snapshot().await, now_ms, &error);
821                    executor.record_run_control(&request, &rejected).await;
822                    return Err(error);
823                }
824            }
825        }
826
827        let (receipt, cancel) = self.submit_locked(request.clone(), now_ms).await?;
828        self.record_receipt(&request, &receipt).await;
829        drop(_admission);
830        if cancel {
831            self.cancellation.cancel();
832        }
833        self.notify.notify_waiters();
834        Ok(receipt)
835    }
836
837    async fn known_receipt(
838        &self,
839        request: &RunControlRequest,
840    ) -> Result<Option<RunControlReceipt>, RunControlError> {
841        let fingerprint = request_fingerprint(request)?;
842        let state = self.state.lock().await;
843        match state.seen.get(&request.request_id) {
844            None => Ok(None),
845            Some(previous) if previous.fingerprint == fingerprint => {
846                Ok(Some(previous.receipt.clone()))
847            }
848            Some(_) => Err(RunControlError::DuplicateRequest {
849                request_id: request.request_id.clone(),
850            }),
851        }
852    }
853
854    async fn record_receipt(&self, request: &RunControlRequest, receipt: &RunControlReceipt) {
855        if let Some(executor) = &self.hook_executor {
856            executor.record_run_control(request, receipt).await;
857        }
858    }
859}
860
861fn request_fingerprint(request: &RunControlRequest) -> Result<String, RunControlError> {
862    let encoded = serde_json::to_vec(request).map_err(|error| {
863        RunControlError::InvalidRequest(format!("could not encode request: {error}"))
864    })?;
865    Ok(format!("sha256:{:x}", Sha256::digest(encoded)))
866}
867
868fn hook_outcome_error(outcome: HookOutcome) -> RunControlError {
869    match outcome {
870        HookOutcome::Block { reason } => RunControlError::HookDenied { reason },
871        HookOutcome::Retry {
872            reason,
873            retry_after_ms,
874        } => RunControlError::HookRetry {
875            reason,
876            retry_after_ms,
877        },
878        HookOutcome::Escalate { reason, target } => RunControlError::HookDenied {
879            reason: target
880                .map(|target| format!("{reason} (escalate to {target})"))
881                .unwrap_or(reason),
882        },
883        HookOutcome::Continue(_) | HookOutcome::Skip => RunControlError::InvalidRequest(
884            "unexpected non-terminal run-control hook outcome".to_string(),
885        ),
886    }
887}
888
889fn rejected_receipt(
890    request: &RunControlRequest,
891    snapshot: &RunControlSnapshot,
892    now_ms: u64,
893    error: &RunControlError,
894) -> RunControlReceipt {
895    RunControlReceipt {
896        schema: RUN_CONTROL_RECEIPT_SCHEMA_V1.to_string(),
897        request_id: request.request_id.clone(),
898        session_id: snapshot.session_id.clone(),
899        run_id: request.run_id.clone(),
900        operation: request.command.operation(),
901        state: RunControlReceiptState::Rejected,
902        sequence: 0,
903        turn_id: snapshot.turn_id.clone(),
904        turn_revision: snapshot.turn_revision,
905        accepted_at_ms: now_ms,
906        applied_at_ms: Some(now_ms),
907        error: Some(RunControlErrorInfo {
908            code: error.code().to_string(),
909            message: error.to_string(),
910        }),
911    }
912}
913
914fn snapshot(state: &InboxState) -> RunControlSnapshot {
915    RunControlSnapshot {
916        session_id: state.session_id.clone(),
917        run_id: state.run_id.clone(),
918        active: state.active && !state.closed,
919        turn_id: state.turn_id.clone(),
920        turn_revision: state.turn_revision,
921        queued_controls: state.queue.len(),
922        interrupt_requested: state.interrupt_requested,
923        last_sequence: state.last_sequence,
924    }
925}
926
927#[cfg(test)]
928mod tests {
929    use super::*;
930    use crate::hooks::{HookEvent, HookResult};
931    use async_trait::async_trait;
932    use std::sync::Mutex as StdMutex;
933
934    fn inbox() -> Arc<RunControlInbox> {
935        RunControlInbox::new("session-1", "run-1", CancellationToken::new())
936    }
937
938    #[tokio::test]
939    async fn steer_is_idempotent_and_applies_at_safe_point() {
940        let inbox = inbox();
941        let turn = inbox.update_turn(1).await;
942        let mut request = RunControlRequest::steer("run-1", "focus on tests")
943            .with_session_id("session-1")
944            .with_expected_turn("turn-1", turn.turn_revision);
945        request.request_id = "req-1".to_string();
946
947        let accepted = inbox.submit(request.clone(), 10).await.unwrap();
948        assert_eq!(accepted.state, RunControlReceiptState::Accepted);
949        assert_eq!(inbox.submit(request, 11).await.unwrap(), accepted);
950
951        let pending = inbox.drain().await;
952        assert_eq!(pending.len(), 1);
953        let applied = inbox
954            .mark_applied(
955                &pending[0],
956                Some("turn-1".to_string()),
957                turn.turn_revision,
958                12,
959            )
960            .await;
961        assert_eq!(applied.state, RunControlReceiptState::Applied);
962        assert_eq!(inbox.snapshot().await.queued_controls, 0);
963    }
964
965    #[tokio::test]
966    async fn stale_turn_and_duplicate_conflict_are_rejected() {
967        let inbox = inbox();
968        let turn = inbox.update_turn(1).await;
969        let mut request = RunControlRequest::steer("run-1", "first")
970            .with_expected_turn("turn-1", turn.turn_revision);
971        request.request_id = "same-id".to_string();
972        inbox.submit(request.clone(), 1).await.unwrap();
973
974        let mut conflicting = request.clone();
975        conflicting.command = RunControlCommand::Steer {
976            input: "different".to_string(),
977        };
978        assert!(matches!(
979            inbox.submit(conflicting, 2).await,
980            Err(RunControlError::DuplicateRequest { .. })
981        ));
982
983        inbox.update_turn(2).await;
984        let stale = RunControlRequest::steer("run-1", "late")
985            .with_expected_turn("turn-1", turn.turn_revision);
986        assert!(matches!(
987            inbox.submit(stale, 3).await,
988            Err(RunControlError::StaleTurn { .. })
989        ));
990    }
991
992    #[tokio::test]
993    async fn interrupt_is_accepted_before_cancellation_fires() {
994        let inbox = inbox();
995        let request = RunControlRequest::interrupt("run-1");
996        let receipt = inbox.submit(request, 1).await.unwrap();
997        assert_eq!(receipt.state, RunControlReceiptState::Accepted);
998        assert!(inbox.cancellation().is_cancelled());
999        assert!(inbox.snapshot().await.interrupt_requested);
1000    }
1001
1002    #[tokio::test]
1003    async fn close_settles_pending_requests() {
1004        let inbox = inbox();
1005        let request = RunControlRequest::steer("run-1", "not applied");
1006        let accepted = inbox.submit(request, 1).await.unwrap();
1007        inbox.close(2).await;
1008        assert!(!inbox.snapshot().await.active);
1009        let retry = RunControlRequest {
1010            request_id: accepted.request_id.clone(),
1011            ..RunControlRequest::steer("run-1", "not applied")
1012        };
1013        let settled = inbox.submit(retry, 3).await.unwrap();
1014        assert_eq!(settled.state, RunControlReceiptState::Settled);
1015        assert_eq!(settled.error.unwrap().code, "RUN_ENDED");
1016    }
1017
1018    #[tokio::test]
1019    async fn close_settles_a_control_already_drained_by_the_loop() {
1020        let inbox = inbox();
1021        let request = RunControlRequest::steer("run-1", "close race");
1022        let accepted = inbox.submit(request, 1).await.unwrap();
1023        let pending = inbox.drain().await;
1024        assert_eq!(pending.len(), 1);
1025
1026        // The loop has taken ownership of the queue item, but has not yet
1027        // acknowledged application. Closing must still produce a terminal
1028        // receipt rather than leaving the request permanently Accepted.
1029        inbox.close(2).await;
1030        let retry = RunControlRequest {
1031            request_id: accepted.request_id.clone(),
1032            ..RunControlRequest::steer("run-1", "close race")
1033        };
1034        let settled = inbox.submit(retry, 3).await.unwrap();
1035        assert_eq!(settled.state, RunControlReceiptState::Settled);
1036
1037        // A late safe-point acknowledgement cannot resurrect the settled
1038        // request or emit another transition.
1039        let late = inbox
1040            .mark_applied(&pending[0], Some("turn-1".into()), 1, 4)
1041            .await;
1042        assert_eq!(late.state, RunControlReceiptState::Settled);
1043    }
1044
1045    #[derive(Debug, Default)]
1046    struct RecordingHook {
1047        events: StdMutex<Vec<HookEvent>>,
1048        deny: bool,
1049    }
1050
1051    #[async_trait]
1052    impl HookExecutor for RecordingHook {
1053        async fn fire(&self, event: &HookEvent) -> HookResult {
1054            self.events.lock().unwrap().push(event.clone());
1055            if self.deny && matches!(event, HookEvent::PreRunControl(_)) {
1056                HookResult::block("host policy denied control")
1057            } else {
1058                HookResult::continue_()
1059            }
1060        }
1061    }
1062
1063    #[tokio::test]
1064    async fn governance_hooks_observe_each_receipt_transition_once() {
1065        let hooks = Arc::new(RecordingHook::default());
1066        let inbox = RunControlInbox::new_with_hook_executor(
1067            "session-1",
1068            "run-1",
1069            CancellationToken::new(),
1070            Some(hooks.clone()),
1071        );
1072        let request = RunControlRequest::steer("run-1", "keep the answer concise")
1073            .with_session_id("session-1");
1074        let accepted = inbox.submit_with_hooks(request.clone(), 10).await.unwrap();
1075        let pending = inbox.drain().await;
1076        let _applied = inbox
1077            .mark_applied(&pending[0], Some("turn-1".to_string()), 1, 11)
1078            .await;
1079        inbox.close(12).await;
1080
1081        let events = hooks.events.lock().unwrap();
1082        assert_eq!(
1083            events
1084                .iter()
1085                .filter(|event| matches!(event, HookEvent::PreRunControl(_)))
1086                .count(),
1087            1,
1088        );
1089        assert_eq!(
1090            events
1091                .iter()
1092                .filter(|event| matches!(event, HookEvent::PostRunControl(_)))
1093                .count(),
1094            2,
1095            "accepted and applied receipts must both be observable",
1096        );
1097        assert_eq!(accepted.state, RunControlReceiptState::Accepted);
1098    }
1099
1100    #[tokio::test]
1101    async fn concurrent_duplicate_submission_runs_governance_once() {
1102        let hooks = Arc::new(RecordingHook::default());
1103        let inbox = RunControlInbox::new_with_hook_executor(
1104            "session-1",
1105            "run-1",
1106            CancellationToken::new(),
1107            Some(hooks.clone()),
1108        );
1109        let mut request = RunControlRequest::steer("run-1", "one admission");
1110        request.request_id = "concurrent-request".to_string();
1111
1112        let attempts = (0..16).map(|_| {
1113            let inbox = Arc::clone(&inbox);
1114            let request = request.clone();
1115            async move { inbox.submit_with_hooks(request, 10).await }
1116        });
1117        let receipts = futures::future::join_all(attempts).await;
1118        let first = receipts[0].as_ref().expect("submission should succeed");
1119        assert!(receipts.iter().all(|result| result.as_ref() == Ok(first)));
1120
1121        let events = hooks.events.lock().unwrap();
1122        assert_eq!(
1123            events
1124                .iter()
1125                .filter(|event| matches!(event, HookEvent::PreRunControl(_)))
1126                .count(),
1127            1,
1128            "a concurrent idempotent retry must not re-run the policy hook",
1129        );
1130        assert_eq!(
1131            events
1132                .iter()
1133                .filter(|event| matches!(event, HookEvent::PostRunControl(_)))
1134                .count(),
1135            1,
1136            "only the first admission emits an accepted observation",
1137        );
1138    }
1139
1140    #[tokio::test]
1141    async fn denied_control_never_enters_the_inbox() {
1142        let hooks = Arc::new(RecordingHook {
1143            deny: true,
1144            ..Default::default()
1145        });
1146        let inbox = RunControlInbox::new_with_hook_executor(
1147            "session-1",
1148            "run-1",
1149            CancellationToken::new(),
1150            Some(hooks.clone()),
1151        );
1152        let error = inbox
1153            .submit_with_hooks(RunControlRequest::interrupt("run-1"), 10)
1154            .await
1155            .unwrap_err();
1156        assert!(matches!(error, RunControlError::HookDenied { .. }));
1157        assert_eq!(inbox.snapshot().await.queued_controls, 0);
1158        let events = hooks.events.lock().unwrap();
1159        assert!(events.iter().any(|event| matches!(
1160            event,
1161            HookEvent::PostRunControl(crate::hooks::PostRunControlEvent {
1162                state: RunControlReceiptState::Rejected,
1163                ..
1164            })
1165        )));
1166    }
1167}
a3s_code_core/run_control.rs

a3s_code_core/
run_control.rs
