a2a_agents/runtime/
memory.rs

1//! [`InMemoryAgentRuntime`] — a process-free fake for tests and composition.
2
3use std::collections::HashMap;
4use std::sync::Arc;
5
6use async_trait::async_trait;
7use tokio::sync::Mutex;
8
9use super::{AgentRuntime, AgentSpec, Recovered, RuntimeError, RuntimeHealth, RuntimeStatus};
10use crate::registry::AgentId;
11
12/// An [`AgentRuntime`] that tracks lifecycle state in a map **without spawning
13/// any processes**.
14///
15/// A first-class adapter (hex rule 6 — not test-only): services like
16/// [`ControlPlane`](crate::control_plane::ControlPlane) are unit-tested against
17/// it, and it serves as a dev substrate when real process isolation isn't wanted.
18/// `start` reports [`RuntimeHealth::Healthy`] immediately — there is no real
19/// process or card to probe, so it models the happy path the composition needs.
20#[derive(Clone, Default)]
21pub struct InMemoryAgentRuntime {
22    agents: Arc<Mutex<HashMap<AgentId, (AgentSpec, RuntimeHealth)>>>,
23    /// Output attributed to each agent, oldest first. Populated by tests via
24    /// [`push_log`](Self::push_log) — nothing here produces output on its own.
25    logs: Arc<Mutex<HashMap<AgentId, Vec<String>>>>,
26}
27
28impl InMemoryAgentRuntime {
29    /// Create an empty runtime.
30    pub fn new() -> Self {
31        Self::default()
32    }
33
34    /// Record a line as though the agent had printed it.
35    ///
36    /// The stand-in for a backend that captures output, so everything above the
37    /// port — [`ControlPlane::logs`](crate::ControlPlane::logs), the HTTP route,
38    /// the client — is exercisable without a container engine or a child
39    /// process. Lines are kept for unknown ids too, so a test can seed a log
40    /// before provisioning.
41    pub async fn push_log(&self, id: &AgentId, line: impl Into<String>) {
42        self.logs
43            .lock()
44            .await
45            .entry(id.clone())
46            .or_default()
47            .push(line.into());
48    }
49}
50
51#[async_trait]
52impl AgentRuntime for InMemoryAgentRuntime {
53    async fn provision(&self, spec: AgentSpec) -> Result<AgentId, RuntimeError> {
54        let id = spec.id.clone();
55        self.agents
56            .lock()
57            .await
58            .insert(id.clone(), (spec, RuntimeHealth::Provisioned));
59        Ok(id)
60    }
61
62    /// Reports everything it holds as [`Recovered::Adopted`].
63    ///
64    /// This fake stands in for a *durable* backend, so it answers like one: the
65    /// map plays the part the container engine plays for real. That is what lets
66    /// [`ControlPlane::recover`](crate::ControlPlane::recover) — which has to
67    /// re-register cards, tolerate unreachable agents, and stay idempotent — be
68    /// tested without Docker or a network.
69    async fn recover(&self) -> Result<Recovered<AgentId>, RuntimeError> {
70        let mut adopted: Vec<AgentId> = self.agents.lock().await.keys().cloned().collect();
71        adopted.sort_by(|a, b| a.as_str().cmp(b.as_str()));
72        Ok(Recovered::Adopted(adopted))
73    }
74
75    async fn start(&self, id: &AgentId) -> Result<(), RuntimeError> {
76        let mut guard = self.agents.lock().await;
77        let (_, health) = guard
78            .get_mut(id)
79            .ok_or_else(|| RuntimeError::NotFound(id.clone()))?;
80        if *health == RuntimeHealth::Healthy {
81            return Err(RuntimeError::AlreadyRunning(id.clone()));
82        }
83        *health = RuntimeHealth::Healthy;
84        Ok(())
85    }
86
87    async fn stop(&self, id: &AgentId) -> Result<(), RuntimeError> {
88        let mut guard = self.agents.lock().await;
89        let (_, health) = guard
90            .get_mut(id)
91            .ok_or_else(|| RuntimeError::NotFound(id.clone()))?;
92        *health = RuntimeHealth::Stopped;
93        Ok(())
94    }
95
96    async fn health(&self, id: &AgentId) -> Result<RuntimeHealth, RuntimeError> {
97        self.agents
98            .lock()
99            .await
100            .get(id)
101            .map(|(_, health)| *health)
102            .ok_or_else(|| RuntimeError::NotFound(id.clone()))
103    }
104
105    async fn list(&self) -> Result<Vec<RuntimeStatus>, RuntimeError> {
106        Ok(self
107            .agents
108            .lock()
109            .await
110            .values()
111            .map(|(spec, health)| RuntimeStatus {
112                id: spec.id.clone(),
113                health: *health,
114                endpoint: spec.endpoint.clone(),
115            })
116            .collect())
117    }
118
119    /// Serve back whatever [`push_log`](Self::push_log) recorded.
120    ///
121    /// Answers like a backend that *does* capture output — an empty list means
122    /// the agent printed nothing, never [`RuntimeError::Unsupported`] — for the
123    /// same reason [`recover`](Self::recover) answers `Adopted`: this fake
124    /// stands in for the durable case so the layers above it can be tested.
125    async fn logs(&self, id: &AgentId, tail: Option<usize>) -> Result<Vec<String>, RuntimeError> {
126        if !self.agents.lock().await.contains_key(id) {
127            return Err(RuntimeError::NotFound(id.clone()));
128        }
129        let lines = self.logs.lock().await.get(id).cloned().unwrap_or_default();
130        Ok(match tail {
131            Some(n) => lines[lines.len().saturating_sub(n)..].to_vec(),
132            None => lines,
133        })
134    }
135}
a2a_agents/runtime/memory.rs

a2a_agents/runtime/
memory.rs
