adk_managed/
default_runtime.rs

1//! Default implementation of the [`ManagedAgentRuntime`] trait.
2//!
3//! [`DefaultManagedAgentRuntime`] composes existing ADK crates (`Runner`,
4//! `SessionService`, optional sandbox and memory) behind the unified lifecycle
5//! trait. It manages active sessions as supervised background tasks with
6//! durable checkpointing, event streaming, and custom tool parking.
7//!
8//! # Architecture
9//!
10//! The runtime is a library, not a service. The platform hosts it:
11//!
12//! - **Testable in isolation**: Zero HTTP/auth/billing dependencies
13//! - **Embeddable**: Self-hosted deployments use the runtime trait directly
14//! - **Swappable platform**: Different platforms can host the same runtime
15//! - **Provider-neutral**: Identical event sequences regardless of model provider
16//!
17//! # Example
18//!
19//! ```rust,ignore
20//! use std::sync::Arc;
21//! use adk_managed::default_runtime::DefaultManagedAgentRuntime;
22//! use adk_managed::resolver::DefaultModelResolver;
23//! use adk_session::InMemorySessionService;
24//!
25//! let resolver = Arc::new(DefaultModelResolver::new());
26//! let sessions = Arc::new(InMemorySessionService::new());
27//!
28//! let runtime = DefaultManagedAgentRuntime::new(resolver, sessions);
29//! ```
30
31use std::collections::HashMap;
32use std::sync::Arc;
33use std::time::Duration;
34
35use async_trait::async_trait;
36use futures::stream::BoxStream;
37use tokio::sync::{Mutex, Notify, RwLock, broadcast, mpsc};
38use tokio_util::sync::CancellationToken;
39use tracing::{debug, info};
40
41use adk_core::Agent;
42#[cfg(feature = "memory")]
43use adk_core::Memory;
44#[cfg(feature = "sandbox")]
45use adk_sandbox::SandboxBackend;
46use adk_session::service::{CreateRequest, SessionService};
47
48use crate::agent_builder::{BuildError, build_agent};
49use crate::checkpoint::CheckpointManager;
50use crate::parking::ToolParkingLot;
51use crate::replay::create_event_stream;
52use crate::resolver::ModelResolver;
53use crate::runtime::{AgentHandle, EnvironmentConfig, ManagedAgentRuntime, SessionHandle};
54use crate::session_loop::SessionLoop;
55use crate::types::{ManagedAgentDef, RuntimeError, SessionEvent, SessionStatus, UserEvent};
56
57// ─── ActiveSession ───────────────────────────────────────────────────────────
58
59/// Internal state for an active (or recently active) session.
60///
61/// Each session spawns a background task running the [`SessionLoop`](crate::session_loop::SessionLoop).
62/// This struct holds the communication handles and control primitives needed
63/// to interact with that background task from the runtime methods.
64#[allow(dead_code)] // Fields are retained for the session lifecycle
65pub(crate) struct ActiveSession {
66    /// The built agent driving this session.
67    pub(crate) agent: Arc<dyn Agent>,
68    /// Sender for user events into the session loop.
69    pub(crate) event_tx: mpsc::Sender<crate::types::UserEvent>,
70    /// Broadcast sender for session events (fan-out to stream subscribers).
71    pub(crate) broadcast_tx: broadcast::Sender<crate::types::SessionEvent>,
72    /// Cancellation token for interrupt handling.
73    pub(crate) cancel_token: CancellationToken,
74    /// Pause flag — when true, the session loop parks until resumed.
75    pub(crate) pause_flag: Arc<Mutex<bool>>,
76    /// Notify used to wake the session loop after resume.
77    pub(crate) pause_notify: Arc<Notify>,
78    /// Current session status (shared with the session loop).
79    pub(crate) status: Arc<RwLock<SessionStatus>>,
80    /// Checkpoint manager for durable state.
81    pub(crate) checkpoint: Arc<RwLock<CheckpointManager>>,
82}
83
84// ─── DefaultManagedAgentRuntime ──────────────────────────────────────────────
85
86/// Default implementation of the managed agent runtime.
87///
88/// Composed from a [`ModelResolver`] + a pluggable [`SessionService`] +
89/// optional sandbox factory and memory service. Has no platform dependencies —
90/// the platform injects its own implementations of these traits.
91///
92/// # Fields
93///
94/// - `model_resolver` — resolves [`ModelRef`](crate::types::ModelRef) into `Arc<dyn Llm>`
95/// - `session_service` — persistent session storage backend
96/// - `sandbox_factory` — optional sandbox for built-in tool execution
97/// - `memory` — optional cross-session memory service
98/// - `sessions` — active session registry
99///
100/// # Example
101///
102/// ```rust,ignore
103/// use std::sync::Arc;
104/// use adk_managed::default_runtime::DefaultManagedAgentRuntime;
105/// use adk_managed::resolver::DefaultModelResolver;
106/// use adk_session::InMemorySessionService;
107///
108/// // Minimal runtime with defaults
109/// let runtime = DefaultManagedAgentRuntime::new(
110///     Arc::new(DefaultModelResolver::new()),
111///     Arc::new(InMemorySessionService::new()),
112/// );
113///
114/// // With sandbox and memory (feature-gated)
115/// let runtime = DefaultManagedAgentRuntime::new(
116///     Arc::new(DefaultModelResolver::new()),
117///     Arc::new(InMemorySessionService::new()),
118/// )
119/// .with_sandbox(my_sandbox)
120/// .with_memory(my_memory_service);
121/// ```
122pub struct DefaultManagedAgentRuntime {
123    /// Resolves ModelRef → `Arc<dyn Llm>`.
124    model_resolver: Arc<dyn ModelResolver>,
125    /// Persistent session storage.
126    session_service: Arc<dyn SessionService>,
127    /// Optional sandbox backend for isolated built-in tool execution.
128    ///
129    /// When set, built-in tools (bash, code_execution, etc.) execute inside
130    /// this sandbox. When `None`, built-in tools execute in-process.
131    #[cfg(feature = "sandbox")]
132    sandbox: Option<Arc<dyn SandboxBackend>>,
133    /// Optional memory service for cross-session persistent memory.
134    ///
135    /// Passed to the Runner's `memory_service` field so agents can search
136    /// and store semantic memories across sessions.
137    #[cfg(feature = "memory")]
138    memory: Option<Arc<dyn Memory>>,
139    /// Registered agents keyed by agent handle ID.
140    agents: Arc<RwLock<HashMap<String, RegisteredAgent>>>,
141    /// Active session registry keyed by session ID.
142    sessions: Arc<RwLock<HashMap<String, ActiveSession>>>,
143}
144
145/// Internal state for a registered agent.
146#[allow(dead_code)] // `def` is retained for future session creation
147struct RegisteredAgent {
148    /// The built agent instance.
149    agent: Arc<dyn Agent>,
150    /// The original definition (retained for session creation).
151    def: ManagedAgentDef,
152}
153
154impl DefaultManagedAgentRuntime {
155    /// Create a new `DefaultManagedAgentRuntime` with injected services.
156    ///
157    /// # Arguments
158    ///
159    /// * `model_resolver` - Resolves `ModelRef` declarations into callable LLM instances.
160    /// * `session_service` - Persistent storage backend for sessions and checkpoints.
161    ///
162    /// Use `.with_sandbox()` and `.with_memory()` builder methods to inject
163    /// optional sandbox and memory services (feature-gated).
164    ///
165    /// # Example
166    ///
167    /// ```rust,ignore
168    /// use std::sync::Arc;
169    /// use adk_managed::default_runtime::DefaultManagedAgentRuntime;
170    /// use adk_managed::resolver::DefaultModelResolver;
171    /// use adk_session::InMemorySessionService;
172    ///
173    /// let runtime = DefaultManagedAgentRuntime::new(
174    ///     Arc::new(DefaultModelResolver::new()),
175    ///     Arc::new(InMemorySessionService::new()),
176    /// );
177    /// ```
178    pub fn new(
179        model_resolver: Arc<dyn ModelResolver>,
180        session_service: Arc<dyn SessionService>,
181    ) -> Self {
182        Self {
183            model_resolver,
184            session_service,
185            #[cfg(feature = "sandbox")]
186            sandbox: None,
187            #[cfg(feature = "memory")]
188            memory: None,
189            agents: Arc::new(RwLock::new(HashMap::new())),
190            sessions: Arc::new(RwLock::new(HashMap::new())),
191        }
192    }
193
194    /// Set the sandbox backend for isolated built-in tool execution.
195    #[cfg(feature = "sandbox")]
196    pub fn with_sandbox(mut self, sandbox: Arc<dyn SandboxBackend>) -> Self {
197        self.sandbox = Some(sandbox);
198        self
199    }
200
201    /// Set the memory service for cross-session persistent memory.
202    #[cfg(feature = "memory")]
203    pub fn with_memory(mut self, memory: Arc<dyn Memory>) -> Self {
204        self.memory = Some(memory);
205        self
206    }
207
208    /// Get a reference to the model resolver.
209    pub fn model_resolver(&self) -> &Arc<dyn ModelResolver> {
210        &self.model_resolver
211    }
212
213    /// Get a reference to the session service.
214    pub fn session_service(&self) -> &Arc<dyn SessionService> {
215        &self.session_service
216    }
217
218    /// Get a reference to the optional sandbox backend.
219    #[cfg(feature = "sandbox")]
220    pub fn sandbox(&self) -> Option<&Arc<dyn SandboxBackend>> {
221        self.sandbox.as_ref()
222    }
223
224    /// Get a reference to the optional memory service.
225    #[cfg(feature = "memory")]
226    pub fn memory(&self) -> Option<&Arc<dyn Memory>> {
227        self.memory.as_ref()
228    }
229
230    /// Get a reference to the active sessions map.
231    #[cfg(test)]
232    pub(crate) fn sessions(&self) -> &Arc<RwLock<HashMap<String, ActiveSession>>> {
233        &self.sessions
234    }
235}
236
237// ─── Channel and timeout defaults ────────────────────────────────────────────
238
239/// Default capacity for the user event mpsc channel.
240const DEFAULT_EVENT_CHANNEL_CAPACITY: usize = 64;
241
242/// Default capacity for the session event broadcast channel.
243const DEFAULT_BROADCAST_CHANNEL_CAPACITY: usize = 256;
244
245/// Default timeout for custom tool parking (5 minutes).
246const DEFAULT_PARKING_TIMEOUT: Duration = Duration::from_secs(300);
247
248// ─── ManagedAgentRuntime implementation ──────────────────────────────────────
249
250#[async_trait]
251impl ManagedAgentRuntime for DefaultManagedAgentRuntime {
252    /// Create a managed agent from a declarative definition.
253    ///
254    /// Resolves the `ModelRef` into an `Arc<dyn Llm>`, builds a runnable agent,
255    /// stores it in the internal registry, and returns an opaque handle.
256    async fn create(&self, def: ManagedAgentDef) -> Result<AgentHandle, RuntimeError> {
257        // 1. Resolve model
258        let model = self.model_resolver.resolve(&def.model).await.map_err(|e| {
259            RuntimeError::ProviderError {
260                provider: format!("{:?}", def.model),
261                message: e.to_string(),
262            }
263        })?;
264
265        // 2. Build agent from definition
266        #[cfg(feature = "sandbox")]
267        let agent = build_agent(&def, model, self.sandbox.clone()).map_err(|e| match e {
268            BuildError::InvalidDef(msg) => RuntimeError::invalid_request(msg),
269            BuildError::BuildFailed(msg) => RuntimeError::internal(msg),
270        })?;
271        #[cfg(not(feature = "sandbox"))]
272        let agent = build_agent(&def, model).map_err(|e| match e {
273            BuildError::InvalidDef(msg) => RuntimeError::invalid_request(msg),
274            BuildError::BuildFailed(msg) => RuntimeError::internal(msg),
275        })?;
276
277        // 3. Generate handle ID
278        let handle_id = uuid::Uuid::new_v4().to_string();
279
280        info!(agent_handle = %handle_id, agent_name = %def.name, "agent created");
281
282        // 4. Store in registry
283        let registered = RegisteredAgent { agent, def };
284        self.agents.write().await.insert(handle_id.clone(), registered);
285
286        Ok(AgentHandle(handle_id))
287    }
288
289    /// Start a new session for the given agent.
290    ///
291    /// Creates internal communication channels, spawns the session loop as a
292    /// background task, and stores the active session handle. Initial status
293    /// is `Queued`.
294    async fn start_session(
295        &self,
296        agent: &AgentHandle,
297        _env: Option<EnvironmentConfig>,
298    ) -> Result<SessionHandle, RuntimeError> {
299        // 1. Look up agent from registry
300        let agents = self.agents.read().await;
301        let registered = agents
302            .get(&agent.0)
303            .ok_or_else(|| RuntimeError::NotFound { session_id: agent.0.clone() })?;
304        let agent_arc = Arc::clone(&registered.agent);
305        drop(agents);
306
307        // 2. Generate session ID
308        let session_id = uuid::Uuid::new_v4().to_string();
309
310        // 3. Create mpsc channel for user events
311        let (event_tx, event_rx) = mpsc::channel(DEFAULT_EVENT_CHANNEL_CAPACITY);
312
313        // 4. Create broadcast channel for session events
314        let (broadcast_tx, _) = broadcast::channel(DEFAULT_BROADCAST_CHANNEL_CAPACITY);
315
316        // 5. Create control primitives
317        let cancel_token = CancellationToken::new();
318        let pause_flag = Arc::new(Mutex::new(false));
319        let pause_notify = Arc::new(Notify::new());
320
321        // 6. Create ToolParkingLot and CheckpointManager
322        let parking = Arc::new(ToolParkingLot::new(DEFAULT_PARKING_TIMEOUT));
323        let checkpoint = Arc::new(RwLock::new(CheckpointManager::new(session_id.clone())));
324
325        // 7. Seed the session in the SessionService.
326        //    The Runner's run() calls session_service.get() which requires the
327        //    session to exist. We create it here with the same triple
328        //    (app_name="managed", user_id="managed_user", session_id) that
329        //    build_runner/run_str use in the session loop.
330        self.session_service
331            .create(CreateRequest {
332                app_name: "managed".to_string(),
333                user_id: "managed_user".to_string(),
334                session_id: Some(session_id.clone()),
335                state: std::collections::HashMap::new(),
336            })
337            .await
338            .map_err(|e| RuntimeError::internal(format!("failed to seed session: {e}")))?;
339
340        // 8. Spawn SessionLoop as background task
341        #[cfg(feature = "memory")]
342        let session_loop = SessionLoop::with_pause_controls(
343            session_id.clone(),
344            event_rx,
345            broadcast_tx.clone(),
346            Arc::clone(&parking),
347            cancel_token.clone(),
348            Arc::clone(&pause_flag),
349            Arc::clone(&pause_notify),
350            Arc::clone(&checkpoint),
351            Arc::clone(&agent_arc),
352            Arc::clone(&self.session_service),
353            self.memory.clone(),
354        );
355        #[cfg(not(feature = "memory"))]
356        let session_loop = SessionLoop::with_pause_controls(
357            session_id.clone(),
358            event_rx,
359            broadcast_tx.clone(),
360            Arc::clone(&parking),
361            cancel_token.clone(),
362            Arc::clone(&pause_flag),
363            Arc::clone(&pause_notify),
364            Arc::clone(&checkpoint),
365            Arc::clone(&agent_arc),
366            Arc::clone(&self.session_service),
367        );
368        tokio::spawn(session_loop.run());
369
370        // 9. Set initial status to Queued
371        let status = Arc::new(RwLock::new(SessionStatus::Queued));
372
373        // 10. Create and store ActiveSession
374        let active_session = ActiveSession {
375            agent: agent_arc,
376            event_tx,
377            broadcast_tx,
378            cancel_token,
379            pause_flag,
380            pause_notify,
381            status,
382            checkpoint,
383        };
384
385        self.sessions.write().await.insert(session_id.clone(), active_session);
386
387        info!(session_id = %session_id, "session started");
388
389        Ok(SessionHandle(session_id))
390    }
391
392    /// Send a user event to the session.
393    ///
394    /// Dispatches the event to the session loop's input channel.
395    async fn send_event(
396        &self,
397        session: &SessionHandle,
398        event: UserEvent,
399    ) -> Result<(), RuntimeError> {
400        let sessions = self.sessions.read().await;
401        let active = sessions
402            .get(&session.0)
403            .ok_or_else(|| RuntimeError::NotFound { session_id: session.0.clone() })?;
404
405        active
406            .event_tx
407            .send(event)
408            .await
409            .map_err(|_| RuntimeError::conflict("session loop channel closed"))?;
410
411        Ok(())
412    }
413
414    /// Subscribe to the session's event stream.
415    ///
416    /// If `from_seq` is provided, replays historical events first, then attaches
417    /// to the live broadcast.
418    async fn stream_events(
419        &self,
420        session: &SessionHandle,
421        from_seq: Option<u64>,
422    ) -> Result<BoxStream<'static, SessionEvent>, RuntimeError> {
423        let sessions = self.sessions.read().await;
424        let active = sessions
425            .get(&session.0)
426            .ok_or_else(|| RuntimeError::NotFound { session_id: session.0.clone() })?;
427
428        // Subscribe to broadcast channel
429        let broadcast_rx = active.broadcast_tx.subscribe();
430
431        // Read checkpoint for replay
432        let checkpoint = active.checkpoint.read().await;
433        let stream = create_event_stream(&checkpoint, broadcast_rx, from_seq);
434
435        Ok(stream)
436    }
437
438    /// Interrupt the session at the next safe boundary.
439    async fn interrupt(&self, session: &SessionHandle) -> Result<(), RuntimeError> {
440        let sessions = self.sessions.read().await;
441        let active = sessions
442            .get(&session.0)
443            .ok_or_else(|| RuntimeError::NotFound { session_id: session.0.clone() })?;
444
445        debug!(session_id = %session.0, "interrupting session");
446        active.cancel_token.cancel();
447
448        Ok(())
449    }
450
451    /// Pause the session, checkpointing current state.
452    async fn pause(&self, session: &SessionHandle) -> Result<(), RuntimeError> {
453        let sessions = self.sessions.read().await;
454        let active = sessions
455            .get(&session.0)
456            .ok_or_else(|| RuntimeError::NotFound { session_id: session.0.clone() })?;
457
458        debug!(session_id = %session.0, "pausing session");
459        *active.pause_flag.lock().await = true;
460        *active.status.write().await = SessionStatus::Paused;
461
462        Ok(())
463    }
464
465    /// Resume a paused session.
466    async fn resume(&self, session: &SessionHandle) -> Result<(), RuntimeError> {
467        let sessions = self.sessions.read().await;
468        let active = sessions
469            .get(&session.0)
470            .ok_or_else(|| RuntimeError::NotFound { session_id: session.0.clone() })?;
471
472        debug!(session_id = %session.0, "resuming session");
473        *active.pause_flag.lock().await = false;
474        *active.status.write().await = SessionStatus::Running;
475        active.pause_notify.notify_one();
476
477        Ok(())
478    }
479
480    /// Query the current status of a session.
481    async fn status(&self, session: &SessionHandle) -> Result<SessionStatus, RuntimeError> {
482        let sessions = self.sessions.read().await;
483        let active = sessions
484            .get(&session.0)
485            .ok_or_else(|| RuntimeError::NotFound { session_id: session.0.clone() })?;
486
487        Ok(*active.status.read().await)
488    }
489
490    /// Archive a session (terminal state).
491    async fn archive(&self, session: &SessionHandle) -> Result<(), RuntimeError> {
492        let sessions = self.sessions.read().await;
493        let active = sessions
494            .get(&session.0)
495            .ok_or_else(|| RuntimeError::NotFound { session_id: session.0.clone() })?;
496
497        debug!(session_id = %session.0, "archiving session");
498        *active.status.write().await = SessionStatus::Archived;
499        active.cancel_token.cancel();
500
501        Ok(())
502    }
503
504    /// Delete a session and its associated data.
505    async fn delete_session(&self, session: &SessionHandle) -> Result<(), RuntimeError> {
506        // First archive (set terminal state and cancel loop)
507        {
508            let sessions = self.sessions.read().await;
509            if let Some(active) = sessions.get(&session.0) {
510                *active.status.write().await = SessionStatus::Archived;
511                active.cancel_token.cancel();
512            }
513        }
514
515        // Remove from sessions map
516        let removed = self.sessions.write().await.remove(&session.0);
517        if removed.is_none() {
518            return Err(RuntimeError::NotFound { session_id: session.0.clone() });
519        }
520
521        debug!(session_id = %session.0, "session deleted");
522        Ok(())
523    }
524}
525
526#[cfg(test)]
527mod tests {
528    use super::*;
529    use crate::resolver::DefaultModelResolver;
530    use crate::types::{ContentBlock, ModelRef};
531    use adk_core::{Content, FinishReason, Llm, LlmRequest, LlmResponse, LlmResponseStream};
532    use async_stream::stream;
533    use futures::StreamExt;
534    use std::time::Duration;
535
536    /// A minimal in-memory session service for testing.
537    /// Uses adk-session's InMemorySessionService.
538    fn mock_session_service() -> Arc<dyn SessionService> {
539        Arc::new(adk_session::InMemorySessionService::new())
540    }
541
542    /// Mock LLM for testing the full runtime lifecycle.
543    struct MockLlm {
544        name: String,
545    }
546
547    impl MockLlm {
548        fn new(name: &str) -> Self {
549            Self { name: name.to_string() }
550        }
551    }
552
553    #[async_trait]
554    impl Llm for MockLlm {
555        fn name(&self) -> &str {
556            &self.name
557        }
558
559        async fn generate_content(
560            &self,
561            _request: LlmRequest,
562            _stream: bool,
563        ) -> adk_core::Result<LlmResponseStream> {
564            let s = stream! {
565                yield Ok(LlmResponse {
566                    content: Some(Content::new("model").with_text("Hello from mock")),
567                    partial: false,
568                    turn_complete: true,
569                    finish_reason: Some(FinishReason::Stop),
570                    ..Default::default()
571                });
572            };
573            Ok(Box::pin(s))
574        }
575    }
576
577    /// A mock resolver that returns a MockLlm for any model ref.
578    struct MockResolver;
579
580    #[async_trait]
581    impl ModelResolver for MockResolver {
582        async fn resolve(
583            &self,
584            _model_ref: &ModelRef,
585        ) -> crate::resolver::ResolverResult<Arc<dyn Llm>> {
586            Ok(Arc::new(MockLlm::new("mock-model")))
587        }
588    }
589
590    fn create_test_runtime() -> DefaultManagedAgentRuntime {
591        let resolver: Arc<dyn ModelResolver> = Arc::new(MockResolver);
592        let sessions = mock_session_service();
593        DefaultManagedAgentRuntime::new(resolver, sessions)
594    }
595
596    #[test]
597    fn test_new_with_minimal_config() {
598        let resolver = Arc::new(DefaultModelResolver::new());
599        let sessions = mock_session_service();
600
601        let _runtime = DefaultManagedAgentRuntime::new(resolver, sessions);
602
603        #[cfg(feature = "sandbox")]
604        assert!(_runtime.sandbox().is_none());
605        #[cfg(feature = "memory")]
606        assert!(_runtime.memory().is_none());
607    }
608
609    #[cfg(all(feature = "sandbox", feature = "memory"))]
610    #[test]
611    fn test_new_with_sandbox_and_memory() {
612        use adk_sandbox::{
613            BackendCapabilities, EnforcedLimits, ExecRequest, ExecResult, Language, SandboxBackend,
614            SandboxError,
615        };
616
617        struct FakeSandbox;
618
619        #[async_trait]
620        impl SandboxBackend for FakeSandbox {
621            fn name(&self) -> &str {
622                "fake"
623            }
624            fn capabilities(&self) -> BackendCapabilities {
625                BackendCapabilities {
626                    supported_languages: vec![Language::Python],
627                    isolation_class: "fake".to_string(),
628                    enforced_limits: EnforcedLimits {
629                        timeout: true,
630                        memory: false,
631                        network_isolation: false,
632                        filesystem_isolation: false,
633                        environment_isolation: false,
634                    },
635                }
636            }
637            async fn execute(&self, _request: ExecRequest) -> Result<ExecResult, SandboxError> {
638                Ok(ExecResult {
639                    stdout: "ok".to_string(),
640                    stderr: String::new(),
641                    exit_code: 0,
642                    duration: std::time::Duration::from_millis(1),
643                })
644            }
645        }
646
647        struct FakeMemory;
648
649        #[async_trait]
650        impl adk_core::Memory for FakeMemory {
651            async fn search(&self, _query: &str) -> adk_core::Result<Vec<adk_core::MemoryEntry>> {
652                Ok(vec![])
653            }
654        }
655
656        let resolver = Arc::new(DefaultModelResolver::new());
657        let sessions = mock_session_service();
658
659        let runtime = DefaultManagedAgentRuntime::new(resolver, sessions)
660            .with_sandbox(Arc::new(FakeSandbox))
661            .with_memory(Arc::new(FakeMemory));
662
663        assert!(runtime.sandbox().is_some());
664        assert!(runtime.memory().is_some());
665    }
666
667    #[test]
668    fn test_sessions_map_starts_empty() {
669        let resolver = Arc::new(DefaultModelResolver::new());
670        let sessions = mock_session_service();
671
672        let runtime = DefaultManagedAgentRuntime::new(resolver, sessions);
673
674        let sessions = runtime.sessions().try_read().unwrap();
675        assert!(sessions.is_empty());
676    }
677
678    #[test]
679    fn test_accessors_return_injected_services() {
680        let resolver: Arc<dyn ModelResolver> = Arc::new(DefaultModelResolver::new());
681        let session_service = mock_session_service();
682
683        let runtime =
684            DefaultManagedAgentRuntime::new(Arc::clone(&resolver), Arc::clone(&session_service));
685
686        // Verify we get references back (type-level verification)
687        let _r: &Arc<dyn ModelResolver> = runtime.model_resolver();
688        let _s: &Arc<dyn SessionService> = runtime.session_service();
689    }
690
691    // ─── Task 7.2: create() method tests ─────────────────────────────────────
692
693    #[tokio::test]
694    async fn test_create_agent_returns_handle() {
695        let runtime = create_test_runtime();
696
697        let def = ManagedAgentDef {
698            name: "test-agent".to_string(),
699            model: ModelRef::Shorthand("gemini-2.5-flash".to_string()),
700            system: Some("You are helpful.".to_string()),
701            description: None,
702            tools: vec![],
703            mcp_servers: vec![],
704            skills: vec![],
705            permission_policy: None,
706            metadata: None,
707        };
708
709        let handle = runtime.create(def).await.unwrap();
710        assert!(!handle.0.is_empty());
711    }
712
713    #[tokio::test]
714    async fn test_create_agent_stores_in_registry() {
715        let runtime = create_test_runtime();
716
717        let def = ManagedAgentDef {
718            name: "stored-agent".to_string(),
719            model: ModelRef::Shorthand("gemini-2.5-flash".to_string()),
720            system: None,
721            description: None,
722            tools: vec![],
723            mcp_servers: vec![],
724            skills: vec![],
725            permission_policy: None,
726            metadata: None,
727        };
728
729        let handle = runtime.create(def).await.unwrap();
730        let agents = runtime.agents.read().await;
731        assert!(agents.contains_key(&handle.0));
732    }
733
734    #[tokio::test]
735    async fn test_create_multiple_agents() {
736        let runtime = create_test_runtime();
737
738        let make_def = |name: &str| ManagedAgentDef {
739            name: name.to_string(),
740            model: ModelRef::Shorthand("gemini-2.5-flash".to_string()),
741            system: None,
742            description: None,
743            tools: vec![],
744            mcp_servers: vec![],
745            skills: vec![],
746            permission_policy: None,
747            metadata: None,
748        };
749
750        let h1 = runtime.create(make_def("agent-1")).await.unwrap();
751        let h2 = runtime.create(make_def("agent-2")).await.unwrap();
752
753        assert_ne!(h1.0, h2.0);
754        assert_eq!(runtime.agents.read().await.len(), 2);
755    }
756
757    // ─── Task 7.3: start_session() method tests ──────────────────────────────
758
759    #[tokio::test]
760    async fn test_start_session_returns_handle() {
761        let runtime = create_test_runtime();
762
763        let def = ManagedAgentDef {
764            name: "session-agent".to_string(),
765            model: ModelRef::Shorthand("gemini-2.5-flash".to_string()),
766            system: None,
767            description: None,
768            tools: vec![],
769            mcp_servers: vec![],
770            skills: vec![],
771            permission_policy: None,
772            metadata: None,
773        };
774
775        let agent = runtime.create(def).await.unwrap();
776        let session = runtime.start_session(&agent, None).await.unwrap();
777        assert!(!session.0.is_empty());
778    }
779
780    #[tokio::test]
781    async fn test_start_session_initial_status_queued() {
782        let runtime = create_test_runtime();
783
784        let def = ManagedAgentDef {
785            name: "status-agent".to_string(),
786            model: ModelRef::Shorthand("gemini-2.5-flash".to_string()),
787            system: None,
788            description: None,
789            tools: vec![],
790            mcp_servers: vec![],
791            skills: vec![],
792            permission_policy: None,
793            metadata: None,
794        };
795
796        let agent = runtime.create(def).await.unwrap();
797        let session = runtime.start_session(&agent, None).await.unwrap();
798
799        let status = runtime.status(&session).await.unwrap();
800        assert_eq!(status, SessionStatus::Queued);
801    }
802
803    #[tokio::test]
804    async fn test_start_session_unknown_agent_returns_error() {
805        let runtime = create_test_runtime();
806
807        let fake_agent = AgentHandle("nonexistent".to_string());
808        let result = runtime.start_session(&fake_agent, None).await;
809        assert!(result.is_err());
810    }
811
812    // ─── Task 7.4: send_event() method tests ─────────────────────────────────
813
814    #[tokio::test]
815    async fn test_send_event_message() {
816        let runtime = create_test_runtime();
817
818        let def = ManagedAgentDef {
819            name: "event-agent".to_string(),
820            model: ModelRef::Shorthand("gemini-2.5-flash".to_string()),
821            system: None,
822            description: None,
823            tools: vec![],
824            mcp_servers: vec![],
825            skills: vec![],
826            permission_policy: None,
827            metadata: None,
828        };
829
830        let agent = runtime.create(def).await.unwrap();
831        let session = runtime.start_session(&agent, None).await.unwrap();
832
833        let event =
834            UserEvent::Message { content: vec![ContentBlock::Text { text: "Hello".to_string() }] };
835
836        let result = runtime.send_event(&session, event).await;
837        assert!(result.is_ok());
838    }
839
840    #[tokio::test]
841    async fn test_send_event_unknown_session_returns_error() {
842        let runtime = create_test_runtime();
843
844        let fake_session = SessionHandle("nonexistent".to_string());
845        let event =
846            UserEvent::Message { content: vec![ContentBlock::Text { text: "Hello".to_string() }] };
847
848        let result = runtime.send_event(&fake_session, event).await;
849        assert!(result.is_err());
850    }
851
852    // ─── Task 7.5: stream_events() method tests ──────────────────────────────
853
854    #[tokio::test]
855    async fn test_stream_events_receives_broadcast() {
856        let runtime = create_test_runtime();
857
858        let def = ManagedAgentDef {
859            name: "stream-agent".to_string(),
860            model: ModelRef::Shorthand("gemini-2.5-flash".to_string()),
861            system: None,
862            description: None,
863            tools: vec![],
864            mcp_servers: vec![],
865            skills: vec![],
866            permission_policy: None,
867            metadata: None,
868        };
869
870        let agent = runtime.create(def).await.unwrap();
871        let session = runtime.start_session(&agent, None).await.unwrap();
872
873        // Subscribe to stream
874        let mut stream = runtime.stream_events(&session, None).await.unwrap();
875
876        // Send a message (the session loop will process it and emit events)
877        let event =
878            UserEvent::Message { content: vec![ContentBlock::Text { text: "Test".to_string() }] };
879        runtime.send_event(&session, event).await.unwrap();
880
881        // We should receive at least a StatusRunning event
882        let first_event = tokio::time::timeout(Duration::from_secs(2), stream.next())
883            .await
884            .expect("timed out waiting for event")
885            .expect("stream ended unexpectedly");
886
887        match first_event {
888            SessionEvent::StatusRunning { .. } => {}
889            other => panic!("expected StatusRunning, got: {other:?}"),
890        }
891    }
892
893    #[tokio::test]
894    async fn test_stream_events_unknown_session_returns_error() {
895        let runtime = create_test_runtime();
896
897        let fake_session = SessionHandle("nonexistent".to_string());
898        let result = runtime.stream_events(&fake_session, None).await;
899        assert!(result.is_err());
900    }
901
902    // ─── Task 7.6: interrupt/pause/resume/status/archive/delete tests ────────
903
904    #[tokio::test]
905    async fn test_interrupt_cancels_session() {
906        let runtime = create_test_runtime();
907
908        let def = ManagedAgentDef {
909            name: "interrupt-agent".to_string(),
910            model: ModelRef::Shorthand("gemini-2.5-flash".to_string()),
911            system: None,
912            description: None,
913            tools: vec![],
914            mcp_servers: vec![],
915            skills: vec![],
916            permission_policy: None,
917            metadata: None,
918        };
919
920        let agent = runtime.create(def).await.unwrap();
921        let session = runtime.start_session(&agent, None).await.unwrap();
922
923        let result = runtime.interrupt(&session).await;
924        assert!(result.is_ok());
925    }
926
927    #[tokio::test]
928    async fn test_pause_sets_paused_status() {
929        let runtime = create_test_runtime();
930
931        let def = ManagedAgentDef {
932            name: "pause-agent".to_string(),
933            model: ModelRef::Shorthand("gemini-2.5-flash".to_string()),
934            system: None,
935            description: None,
936            tools: vec![],
937            mcp_servers: vec![],
938            skills: vec![],
939            permission_policy: None,
940            metadata: None,
941        };
942
943        let agent = runtime.create(def).await.unwrap();
944        let session = runtime.start_session(&agent, None).await.unwrap();
945
946        runtime.pause(&session).await.unwrap();
947        let status = runtime.status(&session).await.unwrap();
948        assert_eq!(status, SessionStatus::Paused);
949    }
950
951    #[tokio::test]
952    async fn test_resume_clears_pause() {
953        let runtime = create_test_runtime();
954
955        let def = ManagedAgentDef {
956            name: "resume-agent".to_string(),
957            model: ModelRef::Shorthand("gemini-2.5-flash".to_string()),
958            system: None,
959            description: None,
960            tools: vec![],
961            mcp_servers: vec![],
962            skills: vec![],
963            permission_policy: None,
964            metadata: None,
965        };
966
967        let agent = runtime.create(def).await.unwrap();
968        let session = runtime.start_session(&agent, None).await.unwrap();
969
970        runtime.pause(&session).await.unwrap();
971        assert_eq!(runtime.status(&session).await.unwrap(), SessionStatus::Paused);
972
973        runtime.resume(&session).await.unwrap();
974        assert_eq!(runtime.status(&session).await.unwrap(), SessionStatus::Running);
975    }
976
977    #[tokio::test]
978    async fn test_archive_sets_archived_status() {
979        let runtime = create_test_runtime();
980
981        let def = ManagedAgentDef {
982            name: "archive-agent".to_string(),
983            model: ModelRef::Shorthand("gemini-2.5-flash".to_string()),
984            system: None,
985            description: None,
986            tools: vec![],
987            mcp_servers: vec![],
988            skills: vec![],
989            permission_policy: None,
990            metadata: None,
991        };
992
993        let agent = runtime.create(def).await.unwrap();
994        let session = runtime.start_session(&agent, None).await.unwrap();
995
996        runtime.archive(&session).await.unwrap();
997        let status = runtime.status(&session).await.unwrap();
998        assert_eq!(status, SessionStatus::Archived);
999    }
1000
1001    #[tokio::test]
1002    async fn test_delete_session_removes_from_registry() {
1003        let runtime = create_test_runtime();
1004
1005        let def = ManagedAgentDef {
1006            name: "delete-agent".to_string(),
1007            model: ModelRef::Shorthand("gemini-2.5-flash".to_string()),
1008            system: None,
1009            description: None,
1010            tools: vec![],
1011            mcp_servers: vec![],
1012            skills: vec![],
1013            permission_policy: None,
1014            metadata: None,
1015        };
1016
1017        let agent = runtime.create(def).await.unwrap();
1018        let session = runtime.start_session(&agent, None).await.unwrap();
1019
1020        runtime.delete_session(&session).await.unwrap();
1021
1022        // Session should no longer be accessible
1023        let result = runtime.status(&session).await;
1024        assert!(result.is_err());
1025    }
1026
1027    #[tokio::test]
1028    async fn test_delete_nonexistent_session_returns_error() {
1029        let runtime = create_test_runtime();
1030
1031        let fake_session = SessionHandle("nonexistent".to_string());
1032        let result = runtime.delete_session(&fake_session).await;
1033        assert!(result.is_err());
1034    }
1035
1036    #[tokio::test]
1037    async fn test_interrupt_nonexistent_session_returns_error() {
1038        let runtime = create_test_runtime();
1039
1040        let fake_session = SessionHandle("nonexistent".to_string());
1041        let result = runtime.interrupt(&fake_session).await;
1042        assert!(result.is_err());
1043    }
1044}
adk_managed/default_runtime.rs

adk_managed/
default_runtime.rs
