Skip to main content

steer_core/runners/
one_shot_runner.rs

1use serde::{Deserialize, Serialize};
2use tokio_util::sync::CancellationToken;
3use tracing::{error, info, warn};
4
5use crate::agents::default_agent_spec_id;
6use crate::app::conversation::{Message, UserContent};
7use crate::app::domain::event::SessionEvent;
8use crate::app::domain::runtime::{RuntimeError, RuntimeHandle};
9use crate::app::domain::types::SessionId;
10use crate::config::model::ModelId;
11use crate::error::{Error, Result};
12use crate::session::ToolApprovalPolicy;
13use crate::session::state::SessionConfig;
14use crate::tools::{DISPATCH_AGENT_TOOL_NAME, DispatchAgentParams, DispatchAgentTarget};
15use steer_tools::ToolCall;
16use steer_tools::tools::BASH_TOOL_NAME;
17use steer_tools::tools::bash::BashParams;
18
19#[derive(Debug, Clone, Serialize, Deserialize)]
20pub struct RunOnceResult {
21    pub final_message: Message,
22    pub session_id: SessionId,
23}
24
25pub struct OneShotRunner;
26
27impl Default for OneShotRunner {
28    fn default() -> Self {
29        Self::new()
30    }
31}
32
33impl OneShotRunner {
34    pub fn new() -> Self {
35        Self
36    }
37
38    pub async fn run_in_session(
39        runtime: &RuntimeHandle,
40        session_id: SessionId,
41        message: String,
42        model: ModelId,
43    ) -> Result<RunOnceResult> {
44        Self::run_in_session_with_cancel(
45            runtime,
46            session_id,
47            message,
48            model,
49            CancellationToken::new(),
50        )
51        .await
52    }
53
54    pub async fn run_in_session_with_cancel(
55        runtime: &RuntimeHandle,
56        session_id: SessionId,
57        message: String,
58        model: ModelId,
59        cancel_token: CancellationToken,
60    ) -> Result<RunOnceResult> {
61        runtime.resume_session(session_id).await.map_err(|e| {
62            Error::InvalidOperation(format!("Failed to resume session {session_id}: {e}"))
63        })?;
64
65        let subscription = runtime.subscribe_events(session_id).await.map_err(|e| {
66            Error::InvalidOperation(format!(
67                "Failed to subscribe to session {session_id} events: {e}"
68            ))
69        })?;
70
71        let approval_policy = match runtime.get_session_state(session_id).await {
72            Ok(state) => state
73                .session_config
74                .map(|config| config.tool_config.approval_policy)
75                .unwrap_or_default(),
76            Err(err) => {
77                warn!(
78                    session_id = %session_id,
79                    error = %err,
80                    "Failed to load session approval policy; defaulting to deny"
81                );
82                ToolApprovalPolicy::default()
83            }
84        };
85
86        info!(session_id = %session_id, message = %message, "Sending message to session");
87
88        let op_id = runtime
89            .submit_user_input(
90                session_id,
91                vec![UserContent::Text {
92                    text: message.clone(),
93                }],
94                model,
95            )
96            .await
97            .map_err(|e| {
98                Error::InvalidOperation(format!(
99                    "Failed to send message to session {session_id}: {e}"
100                ))
101            })?;
102
103        let cancel_task = {
104            let runtime = runtime.clone();
105            let cancel_token = cancel_token.clone();
106            tokio::spawn(async move {
107                cancel_token.cancelled().await;
108                if let Err(err) = runtime.cancel_operation(session_id, Some(op_id)).await {
109                    warn!(
110                        session_id = %session_id,
111                        error = %err,
112                        "Failed to cancel one-shot operation"
113                    );
114                }
115            })
116        };
117
118        let result =
119            Self::process_events(runtime, subscription, session_id, op_id, approval_policy).await;
120
121        cancel_task.abort();
122
123        if let Err(e) = runtime.suspend_session(session_id).await {
124            error!(session_id = %session_id, error = %e, "Failed to suspend session");
125        } else {
126            info!(session_id = %session_id, "Session suspended successfully");
127        }
128
129        result
130    }
131
132    pub async fn run_new_session(
133        runtime: &RuntimeHandle,
134        config: SessionConfig,
135        message: String,
136        model: ModelId,
137    ) -> Result<RunOnceResult> {
138        Self::run_new_session_with_cancel(runtime, config, message, model, CancellationToken::new())
139            .await
140    }
141
142    pub async fn run_new_session_with_cancel(
143        runtime: &RuntimeHandle,
144        config: SessionConfig,
145        message: String,
146        model: ModelId,
147        cancel_token: CancellationToken,
148    ) -> Result<RunOnceResult> {
149        let session_id = runtime
150            .create_session(config)
151            .await
152            .map_err(|e| Error::InvalidOperation(format!("Failed to create session: {e}")))?;
153
154        info!(session_id = %session_id, "Created new session for one-shot run");
155
156        Self::run_in_session_with_cancel(runtime, session_id, message, model, cancel_token).await
157    }
158
159    async fn process_events(
160        runtime: &RuntimeHandle,
161        mut subscription: crate::app::domain::runtime::SessionEventSubscription,
162        session_id: SessionId,
163        op_id: crate::app::domain::types::OpId,
164        approval_policy: ToolApprovalPolicy,
165    ) -> Result<RunOnceResult> {
166        let mut messages = Vec::new();
167        info!(session_id = %session_id, "Starting event processing loop");
168
169        while let Some(envelope) = subscription.recv().await {
170            match envelope.event {
171                SessionEvent::AssistantMessageAdded { message, model: _ } => {
172                    info!(
173                        session_id = %session_id,
174                        role = ?message.role(),
175                        id = %message.id(),
176                        "AssistantMessageAdded event"
177                    );
178                    messages.push(message);
179                }
180
181                SessionEvent::MessageUpdated { message } => {
182                    info!(
183                        session_id = %session_id,
184                        id = %message.id(),
185                        "MessageUpdated event"
186                    );
187                }
188
189                SessionEvent::OperationCompleted {
190                    op_id: completed_op,
191                } => {
192                    if completed_op != op_id {
193                        continue;
194                    }
195                    info!(
196                        session_id = %session_id,
197                        op_id = %completed_op,
198                        "OperationCompleted event received"
199                    );
200                    if !messages.is_empty() {
201                        info!(session_id = %session_id, "Final message received, exiting event loop");
202                        break;
203                    }
204                }
205
206                SessionEvent::OperationCancelled {
207                    op_id: cancelled_op,
208                    ..
209                } => {
210                    if cancelled_op != op_id {
211                        continue;
212                    }
213                    warn!(
214                        session_id = %session_id,
215                        op_id = %cancelled_op,
216                        "OperationCancelled event received"
217                    );
218                    return Err(Error::Cancelled);
219                }
220
221                SessionEvent::Error { message } => {
222                    error!(session_id = %session_id, error = %message, "Error event");
223                    return Err(Error::InvalidOperation(format!(
224                        "Error during processing: {message}"
225                    )));
226                }
227
228                SessionEvent::ApprovalRequested {
229                    request_id,
230                    tool_call,
231                } => {
232                    let approved = tool_is_preapproved(&tool_call, &approval_policy);
233                    if approved {
234                        info!(
235                            session_id = %session_id,
236                            request_id = %request_id,
237                            tool = %tool_call.name,
238                            "Auto-approving preapproved tool"
239                        );
240                    } else {
241                        warn!(
242                            session_id = %session_id,
243                            request_id = %request_id,
244                            tool = %tool_call.name,
245                            "Auto-denying unapproved tool"
246                        );
247                    }
248
249                    runtime
250                        .submit_tool_approval(session_id, request_id, approved, None)
251                        .await
252                        .map_err(|e| {
253                            Error::InvalidOperation(format!(
254                                "Failed to submit tool approval decision: {e}"
255                            ))
256                        })?;
257                }
258
259                _ => {}
260            }
261        }
262
263        match messages.last() {
264            Some(final_message) => {
265                info!(
266                    session_id = %session_id,
267                    message_count = messages.len(),
268                    "Returning final result"
269                );
270                Ok(RunOnceResult {
271                    final_message: final_message.clone(),
272                    session_id,
273                })
274            }
275            None => Err(Error::InvalidOperation("No message received".to_string())),
276        }
277    }
278}
279
280fn tool_is_preapproved(tool_call: &ToolCall, policy: &ToolApprovalPolicy) -> bool {
281    if policy.preapproved.tools.contains(&tool_call.name) {
282        return true;
283    }
284
285    if tool_call.name == DISPATCH_AGENT_TOOL_NAME {
286        let params = serde_json::from_value::<DispatchAgentParams>(tool_call.parameters.clone());
287        if let Ok(params) = params {
288            return match params.target {
289                DispatchAgentTarget::Resume { .. } => true,
290                DispatchAgentTarget::New { agent, .. } => {
291                    let agent_id = agent
292                        .as_deref()
293                        .filter(|value| !value.trim().is_empty())
294                        .map_or_else(|| default_agent_spec_id().to_string(), str::to_string);
295                    policy.is_dispatch_agent_pattern_preapproved(&agent_id)
296                }
297            };
298        }
299    }
300
301    if tool_call.name == BASH_TOOL_NAME {
302        let params = serde_json::from_value::<BashParams>(tool_call.parameters.clone());
303        if let Ok(params) = params {
304            return policy.is_bash_pattern_preapproved(&params.command);
305        }
306    }
307
308    false
309}
310
311impl From<RuntimeError> for Error {
312    fn from(e: RuntimeError) -> Self {
313        match e {
314            RuntimeError::SessionNotFound { session_id } => {
315                Error::InvalidOperation(format!("Session not found: {session_id}"))
316            }
317            RuntimeError::SessionAlreadyExists { session_id } => {
318                Error::InvalidOperation(format!("Session already exists: {session_id}"))
319            }
320            RuntimeError::InvalidInput { message } => Error::InvalidOperation(message),
321            RuntimeError::ChannelClosed => {
322                Error::InvalidOperation("Runtime channel closed".to_string())
323            }
324            RuntimeError::ShuttingDown => {
325                Error::InvalidOperation("Runtime is shutting down".to_string())
326            }
327            RuntimeError::Session(e) => Error::InvalidOperation(format!("Session error: {e}")),
328            RuntimeError::EventStore(e) => {
329                Error::InvalidOperation(format!("Event store error: {e}"))
330            }
331        }
332    }
333}
334
335#[cfg(test)]
336mod tests {
337    use super::*;
338    use crate::api::Client as ApiClient;
339    use crate::api::{ApiError, CompletionResponse, Provider};
340    use crate::app::conversation::{AssistantContent, Message, MessageData};
341    use crate::app::domain::action::ApprovalDecision;
342    use crate::app::domain::runtime::RuntimeService;
343    use crate::app::domain::session::event_store::InMemoryEventStore;
344    use crate::app::validation::ValidatorRegistry;
345    use crate::config::model::builtin;
346    use crate::session::SessionPolicyOverrides;
347    use crate::session::ToolApprovalPolicy;
348    use crate::session::state::{
349        ApprovalRules, SessionToolConfig, UnapprovedBehavior, WorkspaceConfig,
350    };
351    use crate::tools::static_tools::READ_ONLY_TOOL_NAMES;
352    use crate::tools::{BackendRegistry, ToolExecutor};
353    use dotenvy::dotenv;
354    use serde_json::json;
355    use std::collections::{HashMap, HashSet};
356    use std::sync::Arc;
357    use std::sync::Mutex as StdMutex;
358    use steer_tools::ToolCall;
359    use steer_tools::tools::BASH_TOOL_NAME;
360    use tokio_util::sync::CancellationToken;
361
362    #[derive(Clone)]
363    struct ToolCallThenTextProvider {
364        tool_call: ToolCall,
365        final_text: String,
366        call_count: Arc<StdMutex<usize>>,
367    }
368
369    impl ToolCallThenTextProvider {
370        fn new(tool_call: ToolCall, final_text: impl Into<String>) -> Self {
371            Self {
372                tool_call,
373                final_text: final_text.into(),
374                call_count: Arc::new(StdMutex::new(0)),
375            }
376        }
377    }
378
379    #[async_trait::async_trait]
380    impl Provider for ToolCallThenTextProvider {
381        fn name(&self) -> &'static str {
382            "stub-tool-call"
383        }
384
385        async fn complete(
386            &self,
387            _model_id: &crate::config::model::ModelId,
388            _messages: Vec<Message>,
389            _system: Option<crate::app::SystemContext>,
390            _tools: Option<Vec<steer_tools::ToolSchema>>,
391            _call_options: Option<crate::config::model::ModelParameters>,
392            _token: CancellationToken,
393        ) -> std::result::Result<CompletionResponse, ApiError> {
394            let mut count = self
395                .call_count
396                .lock()
397                .expect("tool call counter lock poisoned");
398            let response = if *count == 0 {
399                CompletionResponse {
400                    content: vec![AssistantContent::ToolCall {
401                        tool_call: self.tool_call.clone(),
402                        thought_signature: None,
403                    }],
404                    usage: None,
405                }
406            } else {
407                CompletionResponse {
408                    content: vec![AssistantContent::Text {
409                        text: self.final_text.clone(),
410                    }],
411                    usage: None,
412                }
413            };
414            *count += 1;
415            Ok(response)
416        }
417    }
418
419    async fn create_test_runtime() -> RuntimeService {
420        let event_store = Arc::new(InMemoryEventStore::new());
421        let model_registry = Arc::new(crate::model_registry::ModelRegistry::load(&[]).unwrap());
422        let provider_registry = Arc::new(crate::auth::ProviderRegistry::load(&[]).unwrap());
423        let api_client = Arc::new(ApiClient::new_with_deps(
424            crate::test_utils::test_llm_config_provider().unwrap(),
425            provider_registry,
426            model_registry,
427        ));
428
429        let tool_executor = Arc::new(ToolExecutor::with_components(
430            Arc::new(BackendRegistry::new()),
431            Arc::new(ValidatorRegistry::new()),
432        ));
433
434        RuntimeService::spawn(event_store, api_client, tool_executor)
435    }
436
437    fn create_test_session_config() -> SessionConfig {
438        SessionConfig {
439            default_model: builtin::claude_sonnet_4_5(),
440            workspace: WorkspaceConfig::default(),
441            workspace_ref: None,
442            workspace_id: None,
443            repo_ref: None,
444            parent_session_id: None,
445            workspace_name: None,
446            tool_config: SessionToolConfig::default(),
447            system_prompt: None,
448            primary_agent_id: None,
449            policy_overrides: SessionPolicyOverrides::empty(),
450            metadata: std::collections::HashMap::new(),
451            auto_compaction: crate::session::state::AutoCompactionConfig::default(),
452        }
453    }
454
455    fn create_test_tool_approval_policy() -> ToolApprovalPolicy {
456        let tool_names = READ_ONLY_TOOL_NAMES
457            .iter()
458            .map(|name| (*name).to_string())
459            .collect();
460        ToolApprovalPolicy {
461            default_behavior: UnapprovedBehavior::Prompt,
462            preapproved: ApprovalRules {
463                tools: tool_names,
464                per_tool: std::collections::HashMap::new(),
465            },
466        }
467    }
468
469    #[test]
470    fn tool_is_preapproved_allows_whitelisted_tool() {
471        let policy = create_test_tool_approval_policy();
472        let tool_call = ToolCall {
473            id: "tc_read".to_string(),
474            name: READ_ONLY_TOOL_NAMES[0].to_string(),
475            parameters: json!({}),
476        };
477
478        assert!(tool_is_preapproved(&tool_call, &policy));
479    }
480
481    #[test]
482    fn tool_is_preapproved_allows_bash_pattern() {
483        use crate::session::state::{ApprovalRules, ToolRule, UnapprovedBehavior};
484
485        let mut per_tool = HashMap::new();
486        per_tool.insert(
487            "bash".to_string(),
488            ToolRule::Bash {
489                patterns: vec!["echo *".to_string()],
490            },
491        );
492
493        let policy = ToolApprovalPolicy {
494            default_behavior: UnapprovedBehavior::Prompt,
495            preapproved: ApprovalRules {
496                tools: HashSet::new(),
497                per_tool,
498            },
499        };
500
501        let tool_call = ToolCall {
502            id: "tc_bash".to_string(),
503            name: BASH_TOOL_NAME.to_string(),
504            parameters: json!({ "command": "echo hello" }),
505        };
506
507        assert!(tool_is_preapproved(&tool_call, &policy));
508    }
509
510    #[test]
511    fn tool_is_preapproved_allows_dispatch_agent_pattern() {
512        use crate::session::state::{ApprovalRules, ToolRule, UnapprovedBehavior};
513
514        let mut per_tool = HashMap::new();
515        per_tool.insert(
516            "dispatch_agent".to_string(),
517            ToolRule::DispatchAgent {
518                agent_patterns: vec!["explore".to_string()],
519            },
520        );
521
522        let policy = ToolApprovalPolicy {
523            default_behavior: UnapprovedBehavior::Prompt,
524            preapproved: ApprovalRules {
525                tools: HashSet::new(),
526                per_tool,
527            },
528        };
529
530        let tool_call = ToolCall {
531            id: "tc_dispatch".to_string(),
532            name: DISPATCH_AGENT_TOOL_NAME.to_string(),
533            parameters: json!({
534                "prompt": "find files",
535                "target": {
536                    "session": "new",
537                    "workspace": {
538                        "location": "current"
539                    },
540                    "agent": "explore"
541                }
542            }),
543        };
544
545        assert!(tool_is_preapproved(&tool_call, &policy));
546    }
547
548    #[test]
549    fn tool_is_preapproved_denies_unlisted_tool() {
550        let policy = create_test_tool_approval_policy();
551        let tool_call = ToolCall {
552            id: "tc_other".to_string(),
553            name: "bash".to_string(),
554            parameters: json!({ "command": "rm -rf /" }),
555        };
556
557        assert!(!tool_is_preapproved(&tool_call, &policy));
558    }
559
560    #[tokio::test]
561    async fn run_new_session_denies_unapproved_tool_requests() {
562        let event_store = Arc::new(InMemoryEventStore::new());
563        let model_registry = Arc::new(crate::model_registry::ModelRegistry::load(&[]).unwrap());
564        let provider_registry = Arc::new(crate::auth::ProviderRegistry::load(&[]).unwrap());
565        let api_client = Arc::new(ApiClient::new_with_deps(
566            crate::test_utils::test_llm_config_provider().unwrap(),
567            provider_registry,
568            model_registry.clone(),
569        ));
570
571        let tool_call = ToolCall {
572            id: "tc_1".to_string(),
573            name: "bash".to_string(),
574            parameters: json!({ "command": "echo denied" }),
575        };
576        api_client.insert_test_provider(
577            builtin::claude_sonnet_4_5().provider.clone(),
578            Arc::new(ToolCallThenTextProvider::new(tool_call, "done")),
579        );
580
581        let tool_executor = Arc::new(ToolExecutor::with_components(
582            Arc::new(BackendRegistry::new()),
583            Arc::new(ValidatorRegistry::new()),
584        ));
585        let runtime = RuntimeService::spawn(event_store, api_client, tool_executor);
586
587        let mut config = create_test_session_config();
588        config.tool_config.approval_policy = ToolApprovalPolicy {
589            default_behavior: UnapprovedBehavior::Prompt,
590            preapproved: ApprovalRules {
591                tools: HashSet::new(),
592                per_tool: HashMap::new(),
593            },
594        };
595
596        let model = builtin::claude_sonnet_4_5();
597        let result = OneShotRunner::run_new_session(
598            &runtime.handle,
599            config,
600            "Trigger tool call".to_string(),
601            model,
602        )
603        .await
604        .expect("run_new_session should complete");
605
606        let events = runtime
607            .handle
608            .load_events_after(result.session_id, 0)
609            .await
610            .expect("load events");
611
612        let mut saw_request = false;
613        let mut saw_decision = false;
614        let mut saw_denied = false;
615
616        for (_, event) in events {
617            match event {
618                SessionEvent::ApprovalRequested { .. } => saw_request = true,
619                SessionEvent::ApprovalDecided { decision, .. } => {
620                    saw_decision = true;
621                    if decision == ApprovalDecision::Denied {
622                        saw_denied = true;
623                    }
624                }
625                _ => {}
626            }
627        }
628
629        assert!(saw_request, "expected ApprovalRequested event");
630        assert!(saw_decision, "expected ApprovalDecided event");
631        assert!(saw_denied, "expected denied decision");
632
633        runtime.shutdown().await;
634    }
635
636    #[tokio::test]
637    #[ignore = "Requires API keys and network access"]
638    async fn test_run_new_session_basic() {
639        dotenv().ok();
640        let runtime = create_test_runtime().await;
641
642        let mut config = create_test_session_config();
643        config.tool_config = SessionToolConfig::read_only();
644        config.tool_config.approval_policy = create_test_tool_approval_policy();
645        config
646            .metadata
647            .insert("mode".to_string(), "headless".to_string());
648
649        let model = builtin::claude_sonnet_4_5();
650        let result = OneShotRunner::run_new_session(
651            &runtime.handle,
652            config,
653            "What is 2 + 2?".to_string(),
654            model,
655        )
656        .await;
657
658        let result = tokio::time::timeout(std::time::Duration::from_secs(30), async { result })
659            .await
660            .expect("Timed out waiting for response")
661            .expect("run_new_session failed");
662
663        assert!(!result.final_message.id().is_empty());
664        println!("New session run succeeded: {:?}", result.final_message);
665
666        let content = match &result.final_message.data {
667            MessageData::Assistant { content, .. } => content,
668            _ => panic!("expected assistant message, got {:?}", result.final_message),
669        };
670        let text_content = content.iter().find_map(|c| match c {
671            AssistantContent::Text { text } => Some(text),
672            _ => None,
673        });
674        let content = text_content.expect("No text content found in assistant message");
675        assert!(!content.is_empty(), "Response should not be empty");
676        assert!(
677            content.contains('4'),
678            "Expected response to contain '4', got: {content}"
679        );
680
681        runtime.shutdown().await;
682    }
683
684    #[tokio::test]
685    async fn test_session_creation() {
686        let runtime = create_test_runtime().await;
687
688        let mut config = create_test_session_config();
689        config.tool_config.approval_policy = create_test_tool_approval_policy();
690        config
691            .metadata
692            .insert("test".to_string(), "value".to_string());
693
694        let session_id = runtime.handle.create_session(config).await.unwrap();
695
696        assert!(runtime.handle.is_session_active(session_id).await.unwrap());
697
698        let state = runtime.handle.get_session_state(session_id).await.unwrap();
699        assert_eq!(
700            state.session_config.as_ref().unwrap().metadata.get("test"),
701            Some(&"value".to_string())
702        );
703
704        runtime.shutdown().await;
705    }
706
707    #[tokio::test]
708    async fn test_run_in_session_nonexistent_session() {
709        let runtime = create_test_runtime().await;
710
711        let fake_session_id = SessionId::new();
712        let model = builtin::claude_sonnet_4_5();
713        let result = OneShotRunner::run_in_session(
714            &runtime.handle,
715            fake_session_id,
716            "Test message".to_string(),
717            model,
718        )
719        .await;
720
721        assert!(result.is_err());
722        let err = result.err().unwrap().to_string();
723        assert!(
724            err.contains("not found") || err.contains("Session"),
725            "Expected session not found error, got: {err}"
726        );
727
728        runtime.shutdown().await;
729    }
730
731    #[tokio::test]
732    #[ignore = "Requires API keys and network access"]
733    async fn test_run_in_session_with_real_api() {
734        dotenv().ok();
735        let runtime = create_test_runtime().await;
736
737        let mut config = create_test_session_config();
738        config.tool_config = SessionToolConfig::read_only();
739        config.tool_config.approval_policy = create_test_tool_approval_policy();
740        config
741            .metadata
742            .insert("test".to_string(), "api_test".to_string());
743
744        let session_id = runtime.handle.create_session(config).await.unwrap();
745        let model = builtin::claude_sonnet_4_5();
746
747        let result = OneShotRunner::run_in_session(
748            &runtime.handle,
749            session_id,
750            "What is the capital of France?".to_string(),
751            model,
752        )
753        .await;
754
755        match result {
756            Ok(run_result) => {
757                println!("Session run succeeded: {:?}", run_result.final_message);
758
759                let content = match &run_result.final_message.data {
760                    MessageData::Assistant { content, .. } => content.clone(),
761                    _ => panic!(
762                        "expected assistant message, got {:?}",
763                        run_result.final_message
764                    ),
765                };
766                let text_content = content.iter().find_map(|c| match c {
767                    AssistantContent::Text { text } => Some(text),
768                    _ => None,
769                });
770                let content = text_content.expect("expected text response in assistant message");
771                assert!(!content.is_empty(), "Response should not be empty");
772                assert!(
773                    content.to_lowercase().contains("paris"),
774                    "Expected response to contain 'Paris', got: {content}"
775                );
776            }
777            Err(e) => {
778                println!("Session run failed (expected if no API key): {e}");
779                assert!(
780                    e.to_string().contains("API key")
781                        || e.to_string().contains("authentication")
782                        || e.to_string().contains("timed out"),
783                    "Unexpected error: {e}"
784                );
785            }
786        }
787
788        runtime.shutdown().await;
789    }
790
791    #[tokio::test]
792    #[ignore = "Requires API keys and network access"]
793    async fn test_run_in_session_preserves_context() {
794        dotenv().ok();
795        let runtime = create_test_runtime().await;
796
797        let mut config = create_test_session_config();
798        config.tool_config = SessionToolConfig::read_only();
799        config.tool_config.approval_policy = create_test_tool_approval_policy();
800        config
801            .metadata
802            .insert("test".to_string(), "context_test".to_string());
803
804        let session_id = runtime.handle.create_session(config).await.unwrap();
805        let model = builtin::claude_sonnet_4_5();
806
807        let result1 = OneShotRunner::run_in_session(
808            &runtime.handle,
809            session_id,
810            "My name is Alice and I like pizza.".to_string(),
811            model.clone(),
812        )
813        .await
814        .expect("First session run should succeed");
815
816        println!("First interaction: {:?}", result1.final_message);
817
818        runtime.handle.resume_session(session_id).await.unwrap();
819
820        let result2 = OneShotRunner::run_in_session(
821            &runtime.handle,
822            session_id,
823            "What is my name and what do I like?".to_string(),
824            model,
825        )
826        .await
827        .expect("Second session run should succeed");
828
829        println!("Second interaction: {:?}", result2.final_message);
830
831        match &result2.final_message.data {
832            MessageData::Assistant { content, .. } => {
833                let text_content = content.iter().find_map(|c| match c {
834                    AssistantContent::Text { text } => Some(text),
835                    _ => None,
836                });
837
838                match text_content {
839                    Some(content) => {
840                        assert!(!content.is_empty(), "Response should not be empty");
841                        let content_lower = content.to_lowercase();
842
843                        assert!(
844                            content_lower.contains("alice") || content_lower.contains("name"),
845                            "Expected response to reference the name or context, got: {content}"
846                        );
847                    }
848                    None => {
849                        panic!("expected text response in assistant message");
850                    }
851                }
852            }
853            _ => {
854                panic!(
855                    "expected assistant message, got {:?}",
856                    result2.final_message
857                );
858            }
859        }
860
861        runtime.shutdown().await;
862    }
863
864    #[tokio::test]
865    #[ignore = "Requires API keys and network access"]
866    async fn test_run_new_session_with_tool_usage() {
867        dotenv().ok();
868        let runtime = create_test_runtime().await;
869
870        let mut config = create_test_session_config();
871        config.tool_config = SessionToolConfig::read_only();
872        config.tool_config.approval_policy = create_test_tool_approval_policy();
873        let model = builtin::claude_sonnet_4_5();
874
875        let result = OneShotRunner::run_new_session(
876            &runtime.handle,
877            config,
878            "List the files in the current directory".to_string(),
879            model,
880        )
881        .await
882        .expect("New session run with tools should succeed with valid API key");
883
884        assert!(!result.final_message.id().is_empty());
885        println!(
886            "New session run with tools succeeded: {:?}",
887            result.final_message
888        );
889
890        let has_content = match &result.final_message.data {
891            MessageData::Assistant { content, .. } => content.iter().any(|c| match c {
892                AssistantContent::Text { text } => !text.is_empty(),
893                _ => true,
894            }),
895            _ => false,
896        };
897        assert!(has_content, "Response should have some content");
898
899        runtime.shutdown().await;
900    }
901}