aion_integration_cli/
mock.rs

1//! [`MockAgentHarness`] — the second, INTERVENEABLE shape driven in NOI-8.
2//!
3//! Where [`crate::CliHarness`] proves the **observability-only** tier (empty capability set), this
4//! deterministic in-crate harness proves the **interveneable** tier: it advertises
5//! `{inject_message, cancel}`, accepts those commands (recording them so a test can assert the
6//! neutral command reached the agent), and REJECTS at least one advertised-unsupported primitive
7//! (`RespondToApproval`, and the other two it does not advertise) with a clean
8//! [`HarnessError::CapabilityNotSupported`] NACK.
9//!
10//! Together the two harnesses exercise BOTH branches of the [`AgentSession::intervene`] contract —
11//! the empty-set rejection and the advertised-set accept/gate — through the same neutral seam,
12//! proving `aion-integrations` is a real SDK with two independent implementations, not a Norn
13//! wrapper (§9.1, NOI-8).
14//!
15//! It is a genuine `AgentHarness` implementation (not a stub returning canned outcomes): the worker
16//! driver ([`aion_worker::spawn_agent`]) drives it live in the crate's e2e test, feeding real
17//! commands through the real control channel and observing the real accept/gate results.
18
19use std::sync::Arc;
20use std::sync::Mutex;
21
22use aion_core::{
23    ActivityEvent, ActivityEventKind, ContentType, InterventionCapabilities, InterventionCommand,
24    InterventionKind, InterventionPrimitive, MessageRole, Payload, StopKind,
25};
26use aion_integrations::contract::{AgentHarness, AgentSession};
27use aion_integrations::error::HarnessError;
28use aion_integrations::spec::AgentRunSpec;
29use async_trait::async_trait;
30use chrono::Utc;
31use futures::stream::BoxStream;
32use tokio::sync::mpsc;
33use uuid::Uuid;
34
35/// A record of the commands a [`MockAgentSession`] accepted, shared with the test that drives it.
36pub type AppliedLog = Arc<Mutex<Vec<InterventionKind>>>;
37
38/// A handle that keeps a [`MockAgentSession`]'s run alive and ends it on demand.
39///
40/// The session's event stream drains a channel; holding this sender keeps the stream open (the run
41/// stays live so a driver can deliver commands mid-run), and dropping it — via [`Self::end_run`] or
42/// simply letting it fall out of scope — ends the stream, at which point the worker driver takes the
43/// terminal result. This mirrors the worker's real lifecycle, where the run ends when the harness's
44/// event stream closes, not when a command arrives.
45pub struct RunHandle {
46    keep_alive: mpsc::UnboundedSender<ActivityEvent>,
47}
48
49impl RunHandle {
50    /// Ends the run by closing the event stream, so the driving worker takes the terminal result.
51    pub fn end_run(self) {
52        drop(self.keep_alive);
53    }
54}
55
56/// An interveneable mock harness advertising `{inject_message, cancel}`.
57///
58/// Deterministic: it replays a fixed two-event transcript (a message + a terminal stop), keeps the
59/// run live until the paired [`RunHandle`] is dropped, and yields a fixed result. The `applied` log
60/// is shared so a driver can assert an accepted command reached the session; commands outside the
61/// advertised set are gated and never logged.
62#[derive(Clone)]
63pub struct MockAgentHarness {
64    applied: AppliedLog,
65    run_handle: Arc<Mutex<Option<RunHandle>>>,
66}
67
68impl Default for MockAgentHarness {
69    fn default() -> Self {
70        Self::new()
71    }
72}
73
74impl MockAgentHarness {
75    /// A mock harness with a fresh applied-command log.
76    #[must_use]
77    pub fn new() -> Self {
78        Self {
79            applied: Arc::new(Mutex::new(Vec::new())),
80            run_handle: Arc::new(Mutex::new(None)),
81        }
82    }
83
84    /// The shared log of commands sessions from this harness have accepted.
85    #[must_use]
86    pub fn applied(&self) -> AppliedLog {
87        Arc::clone(&self.applied)
88    }
89
90    /// Takes the [`RunHandle`] for the most recently started run, if a run has started and its
91    /// handle has not already been taken.
92    ///
93    /// A driver ends the run by dropping this handle (or calling [`RunHandle::end_run`]); until then
94    /// the run stays live so commands can be delivered mid-run.
95    #[must_use]
96    pub fn take_run_handle(&self) -> Option<RunHandle> {
97        self.run_handle
98            .lock()
99            .ok()
100            .and_then(|mut guard| guard.take())
101    }
102
103    /// The capability set this harness advertises: `{inject_message, cancel}` — a strict subset of
104    /// the five neutral primitives, so `pause_resume`, `update_budget`, and `respond_to_approval`
105    /// are all advertised-unsupported and gated.
106    #[must_use]
107    pub fn advertised_capabilities() -> InterventionCapabilities {
108        InterventionCapabilities::from_primitives([
109            InterventionPrimitive::InjectMessage,
110            InterventionPrimitive::Cancel,
111        ])
112    }
113}
114
115/// The live session the mock produces: it streams a fixed transcript, gates commands on the
116/// advertised set, and drains an event channel whose sender (the [`RunHandle`]) a driver holds so
117/// the run stays live until the driver ends it.
118pub struct MockAgentSession {
119    capabilities: InterventionCapabilities,
120    applied: AppliedLog,
121    /// A channel draining into the events stream. The paired sender is held by the [`RunHandle`];
122    /// dropping it ends the stream (and the run). A fixed transcript is pre-loaded before the
123    /// handle is published.
124    events: Option<mpsc::UnboundedReceiver<ActivityEvent>>,
125    result: Payload,
126}
127
128#[async_trait]
129impl AgentHarness for MockAgentHarness {
130    type Session = MockAgentSession;
131
132    async fn start(&self, spec: AgentRunSpec) -> Result<Self::Session, HarnessError> {
133        let (tx, rx) = mpsc::unbounded_channel();
134        // Pre-load a fixed transcript stamped with the spec identity. The run stays live until the
135        // RunHandle (holding `tx`) is dropped by the driver, at which point the stream ends.
136        for event in fixed_transcript(&spec) {
137            let _ = tx.send(event);
138        }
139        if let Ok(mut guard) = self.run_handle.lock() {
140            *guard = Some(RunHandle { keep_alive: tx });
141        }
142        Ok(MockAgentSession {
143            capabilities: MockAgentHarness::advertised_capabilities(),
144            applied: Arc::clone(&self.applied),
145            events: Some(rx),
146            result: Payload::new(ContentType::Json, b"{\"result\":\"completed\"}".to_vec()),
147        })
148    }
149}
150
151/// A fixed two-event transcript stamped with the run identity: one assistant message + a stop.
152fn fixed_transcript(spec: &AgentRunSpec) -> Vec<ActivityEvent> {
153    let base = ActivityEvent {
154        workflow_id: spec.workflow_id.clone(),
155        activity_id: spec.activity_id.clone(),
156        attempt: spec.attempt,
157        agent_id: Uuid::nil(),
158        agent_role: "mock".to_owned(),
159        emitted_at: Utc::now(),
160        worker_seq: 0,
161        store_seq: None,
162        ephemeral: false,
163        kind: ActivityEventKind::Message {
164            role: MessageRole::Assistant,
165            text: "working".to_owned(),
166        },
167    };
168    let stop = ActivityEvent {
169        worker_seq: 1,
170        kind: ActivityEventKind::Stop {
171            reason: StopKind::EndTurn,
172        },
173        ..base.clone()
174    };
175    vec![base, stop]
176}
177
178#[async_trait]
179impl AgentSession for MockAgentSession {
180    fn capabilities(&self) -> &InterventionCapabilities {
181        &self.capabilities
182    }
183
184    fn events(&mut self) -> BoxStream<'static, ActivityEvent> {
185        match self.events.take() {
186            Some(receiver) => Box::pin(tokio_stream::wrappers::UnboundedReceiverStream::new(
187                receiver,
188            )),
189            None => Box::pin(futures::stream::empty()),
190        }
191    }
192
193    async fn intervene(&self, cmd: InterventionCommand) -> Result<(), HarnessError> {
194        // Capability-gate on the advertised set: inject_message/cancel are accepted; the three
195        // unadvertised primitives (pause_resume/update_budget/respond_to_approval) are cleanly
196        // rejected as capability-not-supported and NEVER recorded as applied.
197        if !self.capabilities.supports(&cmd.kind) {
198            return Err(HarnessError::capability_not_supported(format!(
199                "{:?}",
200                cmd.kind.primitive()
201            )));
202        }
203        self.applied
204            .lock()
205            .map_err(|_poison| HarnessError::harness("applied-command lock poisoned"))?
206            .push(cmd.kind);
207        Ok(())
208    }
209
210    async fn wait_result(self) -> Result<Payload, HarnessError> {
211        // The run has already ended (the driver dropped the RunHandle, closing the event stream);
212        // the terminal result is the fixed completed payload.
213        Ok(self.result)
214    }
215}
216
217#[cfg(test)]
218#[allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
219mod tests {
220    use super::*;
221    use aion_core::{ActivityId, ApprovalDecision, InjectPriority, WorkflowId};
222    use futures::StreamExt;
223
224    fn spec() -> AgentRunSpec {
225        AgentRunSpec::new(
226            WorkflowId::new(Uuid::nil()),
227            ActivityId::from_sequence_position(4),
228            2,
229            "mock-activity",
230            Payload::new(ContentType::Json, b"\"in\"".to_vec()),
231        )
232    }
233
234    fn command(kind: InterventionKind) -> InterventionCommand {
235        InterventionCommand {
236            workflow_id: WorkflowId::new(Uuid::nil()),
237            activity_id: ActivityId::from_sequence_position(4),
238            attempt: 2,
239            issued_by: Some("operator".to_owned()),
240            issued_at: Utc::now(),
241            kind,
242        }
243    }
244
245    #[tokio::test]
246    async fn advertises_inject_message_and_cancel_only() {
247        let harness = MockAgentHarness::new();
248        let session = harness.start(spec()).await.unwrap();
249        let caps = session.capabilities();
250        assert!(caps.supports_primitive(InterventionPrimitive::InjectMessage));
251        assert!(caps.supports_primitive(InterventionPrimitive::Cancel));
252        assert!(!caps.supports_primitive(InterventionPrimitive::PauseResume));
253        assert!(!caps.supports_primitive(InterventionPrimitive::UpdateBudget));
254        assert!(!caps.supports_primitive(InterventionPrimitive::RespondToApproval));
255    }
256
257    #[tokio::test]
258    async fn accepts_an_advertised_command_and_records_it() {
259        let harness = MockAgentHarness::new();
260        let applied = harness.applied();
261        let session = harness.start(spec()).await.unwrap();
262
263        session
264            .intervene(command(InterventionKind::InjectMessage {
265                text: "steer".to_owned(),
266                priority: InjectPriority::Interrupt,
267            }))
268            .await
269            .unwrap();
270        session
271            .intervene(command(InterventionKind::Cancel {
272                reason: "stop".to_owned(),
273            }))
274            .await
275            .unwrap();
276
277        assert_eq!(applied.lock().unwrap().len(), 2);
278    }
279
280    #[tokio::test]
281    async fn rejects_respond_to_approval_as_capability_not_supported() {
282        let harness = MockAgentHarness::new();
283        let applied = harness.applied();
284        let session = harness.start(spec()).await.unwrap();
285
286        let error = session
287            .intervene(command(InterventionKind::RespondToApproval {
288                call_id: "c1".to_owned(),
289                decision: ApprovalDecision::Approve,
290                note: None,
291            }))
292            .await
293            .unwrap_err();
294        assert!(
295            matches!(error, HarnessError::CapabilityNotSupported { .. }),
296            "an advertised-unsupported primitive is cleanly NACKed, got {error:?}"
297        );
298        assert!(
299            applied.lock().unwrap().is_empty(),
300            "a gated command is never recorded as applied"
301        );
302    }
303
304    #[tokio::test]
305    async fn streams_a_transcript_and_returns_a_result() {
306        let harness = MockAgentHarness::new();
307        let mut session = harness.start(spec()).await.unwrap();
308        // End the run so the pre-loaded transcript stream terminates and can be collected.
309        harness.take_run_handle().expect("a run handle").end_run();
310        let events: Vec<ActivityEvent> = session.events().collect().await;
311        assert_eq!(events.len(), 2, "the fixed transcript streams");
312        assert!(matches!(events[0].kind, ActivityEventKind::Message { .. }));
313        assert!(matches!(events[1].kind, ActivityEventKind::Stop { .. }));
314        let payload = session.wait_result().await.unwrap();
315        assert_eq!(payload.content_type(), &ContentType::Json);
316    }
317}
aion_integration_cli/mock.rs

aion_integration_cli/
mock.rs
