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        run_id: spec.run_id.clone(),
156        activity_id: spec.activity_id.clone(),
157        attempt: spec.attempt,
158        agent_id: Uuid::nil(),
159        agent_role: "mock".to_owned(),
160        emitted_at: Utc::now(),
161        worker_seq: 0,
162        store_seq: None,
163        ephemeral: false,
164        kind: ActivityEventKind::Message {
165            role: MessageRole::Assistant,
166            text: "working".to_owned(),
167        },
168    };
169    let stop = ActivityEvent {
170        worker_seq: 1,
171        kind: ActivityEventKind::Stop {
172            reason: StopKind::EndTurn,
173        },
174        ..base.clone()
175    };
176    vec![base, stop]
177}
178
179#[async_trait]
180impl AgentSession for MockAgentSession {
181    fn capabilities(&self) -> &InterventionCapabilities {
182        &self.capabilities
183    }
184
185    fn events(&mut self) -> BoxStream<'static, ActivityEvent> {
186        match self.events.take() {
187            Some(receiver) => Box::pin(tokio_stream::wrappers::UnboundedReceiverStream::new(
188                receiver,
189            )),
190            None => Box::pin(futures::stream::empty()),
191        }
192    }
193
194    async fn intervene(&self, cmd: InterventionCommand) -> Result<(), HarnessError> {
195        // Capability-gate on the advertised set: inject_message/cancel are accepted; the three
196        // unadvertised primitives (pause_resume/update_budget/respond_to_approval) are cleanly
197        // rejected as capability-not-supported and NEVER recorded as applied.
198        if !self.capabilities.supports(&cmd.kind) {
199            return Err(HarnessError::capability_not_supported(format!(
200                "{:?}",
201                cmd.kind.primitive()
202            )));
203        }
204        self.applied
205            .lock()
206            .map_err(|_poison| HarnessError::harness("applied-command lock poisoned"))?
207            .push(cmd.kind);
208        Ok(())
209    }
210
211    async fn wait_result(self) -> Result<Payload, HarnessError> {
212        // The run has already ended (the driver dropped the RunHandle, closing the event stream);
213        // the terminal result is the fixed completed payload.
214        Ok(self.result)
215    }
216}
217
218#[cfg(test)]
219#[allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
220mod tests {
221    use super::*;
222    use aion_core::{ActivityId, ApprovalDecision, InjectPriority, RunId, WorkflowId};
223    use futures::StreamExt;
224
225    fn spec() -> AgentRunSpec {
226        AgentRunSpec::new(
227            WorkflowId::new(Uuid::nil()),
228            RunId::new_v4(),
229            ActivityId::from_sequence_position(4),
230            2,
231            "mock-activity",
232            Payload::new(ContentType::Json, b"\"in\"".to_vec()),
233        )
234    }
235
236    fn command(kind: InterventionKind) -> InterventionCommand {
237        InterventionCommand {
238            workflow_id: WorkflowId::new(Uuid::nil()),
239            run_id: RunId::new_v4(),
240            activity_id: ActivityId::from_sequence_position(4),
241            attempt: 2,
242            issued_by: Some("operator".to_owned()),
243            issued_at: Utc::now(),
244            kind,
245        }
246    }
247
248    #[tokio::test]
249    async fn advertises_inject_message_and_cancel_only() {
250        let harness = MockAgentHarness::new();
251        let session = harness.start(spec()).await.unwrap();
252        let caps = session.capabilities();
253        assert!(caps.supports_primitive(InterventionPrimitive::InjectMessage));
254        assert!(caps.supports_primitive(InterventionPrimitive::Cancel));
255        assert!(!caps.supports_primitive(InterventionPrimitive::PauseResume));
256        assert!(!caps.supports_primitive(InterventionPrimitive::UpdateBudget));
257        assert!(!caps.supports_primitive(InterventionPrimitive::RespondToApproval));
258    }
259
260    #[tokio::test]
261    async fn accepts_an_advertised_command_and_records_it() {
262        let harness = MockAgentHarness::new();
263        let applied = harness.applied();
264        let session = harness.start(spec()).await.unwrap();
265
266        session
267            .intervene(command(InterventionKind::InjectMessage {
268                text: "steer".to_owned(),
269                priority: InjectPriority::Interrupt,
270            }))
271            .await
272            .unwrap();
273        session
274            .intervene(command(InterventionKind::Cancel {
275                reason: "stop".to_owned(),
276            }))
277            .await
278            .unwrap();
279
280        assert_eq!(applied.lock().unwrap().len(), 2);
281    }
282
283    #[tokio::test]
284    async fn rejects_respond_to_approval_as_capability_not_supported() {
285        let harness = MockAgentHarness::new();
286        let applied = harness.applied();
287        let session = harness.start(spec()).await.unwrap();
288
289        let error = session
290            .intervene(command(InterventionKind::RespondToApproval {
291                call_id: "c1".to_owned(),
292                decision: ApprovalDecision::Approve,
293                note: None,
294            }))
295            .await
296            .unwrap_err();
297        assert!(
298            matches!(error, HarnessError::CapabilityNotSupported { .. }),
299            "an advertised-unsupported primitive is cleanly NACKed, got {error:?}"
300        );
301        assert!(
302            applied.lock().unwrap().is_empty(),
303            "a gated command is never recorded as applied"
304        );
305    }
306
307    #[tokio::test]
308    async fn streams_a_transcript_and_returns_a_result() {
309        let harness = MockAgentHarness::new();
310        let mut session = harness.start(spec()).await.unwrap();
311        // End the run so the pre-loaded transcript stream terminates and can be collected.
312        harness.take_run_handle().expect("a run handle").end_run();
313        let events: Vec<ActivityEvent> = session.events().collect().await;
314        assert_eq!(events.len(), 2, "the fixed transcript streams");
315        assert!(matches!(events[0].kind, ActivityEventKind::Message { .. }));
316        assert!(matches!(events[1].kind, ActivityEventKind::Stop { .. }));
317        let payload = session.wait_result().await.unwrap();
318        assert_eq!(payload.content_type(), &ContentType::Json);
319    }
320}
aion_integration_cli/mock.rs

aion_integration_cli/
mock.rs
