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

aion_integration_cli/
mock.rs
