1use std::path::PathBuf;
9use std::sync::{Arc, Mutex};
10
11use ag_protocol::{
12 AgentResponse, build_protocol_repair_prompt, format_protocol_parse_debug_details,
13 parse_agent_response_strict,
14};
15use async_trait::async_trait;
16
17use super::backend::{AgentBackend, BuildCommandRequest};
18use super::cli::error;
19use super::cli::execution::{self, CliExecutionError, CliExecutionObserver, CliExitStatus};
20use super::{
21 ParsedResponse, create_app_server_client, create_backend, parse_response, transport_mode,
22};
23use crate::app_server::{AppServerClient, AppServerTurnRequest};
24use crate::channel::AgentRequestKind;
25use crate::model::agent::{AgentKind, AgentModel, ReasoningLevel};
26use crate::model::permission::PermissionMode;
27use crate::model::session::{SessionDiffState, SessionStats};
28
29#[derive(Clone, Debug)]
32pub struct OneShotRequest {
33 pub agent_kind: AgentKind,
36 pub child_pid: Option<Arc<Mutex<Option<u32>>>>,
39 pub folder: PathBuf,
41 pub model: AgentModel,
43 pub prompt: String,
45 pub request_kind: AgentRequestKind,
47 pub reasoning_level: ReasoningLevel,
49}
50
51#[derive(Clone, Debug, PartialEq)]
53pub struct OneShotSubmission {
54 pub response: AgentResponse,
56 pub stats: SessionStats,
58}
59
60#[derive(Clone, Debug, Eq, PartialEq, thiserror::Error)]
66#[error("{message}")]
67pub struct OneShotError {
68 message: String,
69}
70
71impl OneShotError {
72 pub fn new(message: impl Into<String>) -> Self {
74 Self {
75 message: message.into(),
76 }
77 }
78}
79
80#[cfg_attr(any(test, feature = "test-utils"), mockall::automock)]
86#[async_trait]
87pub trait OneShotClient: Send + Sync {
88 async fn submit(&self, request: OneShotRequest) -> Result<OneShotSubmission, OneShotError>;
90}
91
92pub struct RealOneShotClient {
94 app_server_client_override: Option<Arc<dyn AppServerClient>>,
95}
96
97impl RealOneShotClient {
98 pub fn new(app_server_client_override: Option<Arc<dyn AppServerClient>>) -> Self {
103 Self {
104 app_server_client_override,
105 }
106 }
107}
108
109#[async_trait]
110impl OneShotClient for RealOneShotClient {
111 async fn submit(&self, request: OneShotRequest) -> Result<OneShotSubmission, OneShotError> {
112 let app_server_client_override = self.app_server_client_override.as_ref().map(Arc::clone);
113
114 submit_one_shot_with_stats_and_app_server_client(request, app_server_client_override)
115 .await
116 .map_err(OneShotError::new)
117 }
118}
119
120async fn submit_one_shot_with_stats_and_app_server_client(
128 request: OneShotRequest,
129 app_server_client_override: Option<Arc<dyn AppServerClient>>,
130) -> Result<OneShotSubmission, String> {
131 let backend = create_backend(request.agent_kind);
132
133 if transport_mode(request.agent_kind).uses_app_server() {
134 let app_server_client =
135 create_app_server_client(request.agent_kind, app_server_client_override).ok_or_else(
136 || {
137 format!(
138 "{} provider did not provide an app-server client",
139 request.agent_kind
140 )
141 },
142 )?;
143
144 return submit_one_shot_with_app_server_client(app_server_client.as_ref(), request).await;
145 }
146
147 submit_one_shot_with_backend(backend.as_ref(), request).await
148}
149
150async fn submit_one_shot_with_app_server_client(
160 app_server_client: &dyn AppServerClient,
161 request: OneShotRequest,
162) -> Result<OneShotSubmission, String> {
163 clear_child_pid_slot(request.child_pid.as_deref());
164
165 let session_id = format!("one-shot-{}", uuid::Uuid::new_v4());
166 let (stream_tx, _stream_rx) = tokio::sync::mpsc::unbounded_channel();
167 let turn_request = AppServerTurnRequest {
168 folder: request.folder.clone(),
169 live_transcript: None,
170 main_checkout_root: None,
171 model: request.model.provider_model_str().to_string(),
172 permission_mode: PermissionMode::AutoEdit,
173 personality: crate::channel::PersonalityPrompt::default(),
174 prompt: ag_protocol::TurnPrompt::from_agent_data(request.prompt.clone()),
175 request_kind: request.request_kind.clone(),
176 replay_transcript: None,
177 provider_conversation_id: None,
178 persisted_instruction_conversation_id: None,
179 reasoning_level: request.reasoning_level,
180 session_id: session_id.clone(),
181 speed_mode: crate::model::session::SpeedMode::default(),
182 };
183
184 let turn_result = app_server_client.run_turn(turn_request, stream_tx).await;
185
186 let child_pid = request.child_pid.as_ref().map(Arc::clone);
187
188 let turn_result = match turn_result {
189 Ok(result) => result,
190 Err(error) => {
191 app_server_client.shutdown_session(session_id).await;
192 clear_child_pid_slot(child_pid.as_deref());
193
194 return Err(format!(
195 "Failed to execute one-shot app-server turn: {error}"
196 ));
197 }
198 };
199
200 let parse_result = match parse_one_shot_response(&turn_result.assistant_message) {
201 Ok(response) => Ok((response, 0, 0)),
202 Err(parse_error) => {
203 attempt_one_shot_app_server_repair(
204 app_server_client,
205 &parse_error,
206 &turn_result.assistant_message,
207 request,
208 &session_id,
209 turn_result.provider_conversation_id.as_deref(),
210 )
211 .await
212 }
213 };
214
215 app_server_client.shutdown_session(session_id).await;
216 clear_child_pid_slot(child_pid.as_deref());
217
218 let (response, repair_input_tokens, repair_output_tokens) = parse_result?;
219
220 Ok(OneShotSubmission {
221 response,
222 stats: SessionStats {
223 added_lines: 0,
224 deleted_lines: 0,
225 diff_state: SessionDiffState::Unknown,
226 input_tokens: turn_result.input_tokens + repair_input_tokens,
227 output_tokens: turn_result.output_tokens + repair_output_tokens,
228 },
229 })
230}
231
232async fn submit_one_shot_with_backend(
242 backend: &dyn AgentBackend,
243 request: OneShotRequest,
244) -> Result<OneShotSubmission, String> {
245 let parsed_response =
246 execute_one_shot_command(backend, &request.prompt, request.clone()).await?;
247 let (agent_response, repair_stats) = match parse_one_shot_response(&parsed_response.content) {
248 Ok(response) => (response, None),
249 Err(parse_error) => {
250 let repair_prompt =
251 build_protocol_repair_prompt(&parse_error, &parsed_response.content);
252 let repair_response = execute_one_shot_command(backend, &repair_prompt, request)
253 .await
254 .map_err(|error| format!("{parse_error}\nrepair transport failed: {error}"))?;
255
256 let response = parse_one_shot_response(&repair_response.content).map_err(|error| {
257 format!(
258 "{parse_error}\nrepair retry also failed: {error}\nrepair_response:\n{}",
259 repair_response.content
260 )
261 })?;
262
263 (response, Some(repair_response.stats))
264 }
265 };
266
267 let mut stats = parsed_response.stats;
268 if let Some(repair) = repair_stats {
269 stats.input_tokens += repair.input_tokens;
270 stats.output_tokens += repair.output_tokens;
271 }
272
273 Ok(OneShotSubmission {
274 response: agent_response,
275 stats,
276 })
277}
278
279fn parse_one_shot_response(content: &str) -> Result<AgentResponse, String> {
286 parse_agent_response_strict(content).map_err(|error| {
287 format!(
288 "One-shot agent output did not match the required JSON schema: \
289 {error}\ndebug_details:\n{}",
290 format_protocol_parse_debug_details(content)
291 )
292 })
293}
294
295async fn attempt_one_shot_app_server_repair(
309 app_server_client: &dyn AppServerClient,
310 parse_error: &str,
311 malformed_response: &str,
312 request: OneShotRequest,
313 session_id: &str,
314 provider_conversation_id: Option<&str>,
315) -> Result<(AgentResponse, u64, u64), String> {
316 let repair_prompt = build_protocol_repair_prompt(parse_error, malformed_response);
317
318 let (repair_stream_tx, _repair_stream_rx) = tokio::sync::mpsc::unbounded_channel();
319 let repair_turn_request = AppServerTurnRequest {
320 folder: request.folder,
321 live_transcript: None,
322 main_checkout_root: None,
323 model: request.model.provider_model_str().to_string(),
324 permission_mode: PermissionMode::AutoEdit,
325 personality: crate::channel::PersonalityPrompt::default(),
326 prompt: ag_protocol::TurnPrompt::from_agent_data(repair_prompt),
327 request_kind: request.request_kind,
328 replay_transcript: None,
329 provider_conversation_id: provider_conversation_id.map(String::from),
330 persisted_instruction_conversation_id: None,
331 reasoning_level: request.reasoning_level,
332 session_id: session_id.to_string(),
333 speed_mode: crate::model::session::SpeedMode::default(),
334 };
335 let repair_result = app_server_client
336 .run_turn(repair_turn_request, repair_stream_tx)
337 .await
338 .map_err(|error| format!("{parse_error}\nrepair transport failed: {error}"))?;
339
340 let response = parse_one_shot_response(&repair_result.assistant_message).map_err(|error| {
341 format!(
342 "{parse_error}\nrepair retry also failed: {error}\nrepair_response:\n{}",
343 repair_result.assistant_message
344 )
345 })?;
346
347 Ok((
348 response,
349 repair_result.input_tokens,
350 repair_result.output_tokens,
351 ))
352}
353
354async fn execute_one_shot_command(
364 backend: &dyn AgentBackend,
365 prompt: &str,
366 request: OneShotRequest,
367) -> Result<ParsedResponse, String> {
368 let prompt_payload = ag_protocol::TurnPrompt::from_agent_data(prompt.to_string());
369 let build_request = BuildCommandRequest {
370 attachments: &prompt_payload.attachments,
371 folder: &request.folder,
372 main_checkout_root: None,
373 replay_transcript: None,
374 model: request.model.provider_model_str(),
375 permission_mode: PermissionMode::AutoEdit,
376 personality_prompt: None,
377 prompt,
378 reasoning_level: request.reasoning_level,
379 request_kind: &request.request_kind,
380 speed_mode: crate::model::session::SpeedMode::default(),
381 };
382 let observer = OneShotCliObserver {
383 child_pid: request.child_pid,
384 };
385 let output =
386 execution::execute_cli_command(backend, request.agent_kind, build_request, &observer, None)
387 .await
388 .map_err(format_one_shot_execution_error)?;
389
390 match output.exit_status {
391 CliExitStatus::Signaled(_) => {
392 return Err("One-shot agent command was interrupted".to_string());
393 }
394 CliExitStatus::NonZero(exit_code) => {
395 return Err(format_one_shot_exit_error(
396 request.agent_kind,
397 exit_code,
398 &output.stdout,
399 &output.stderr,
400 ));
401 }
402 CliExitStatus::Success => {}
403 }
404
405 let parsed_response = parse_response(request.agent_kind, &output.stdout, &output.stderr);
406
407 Ok(parsed_response)
408}
409
410fn format_one_shot_execution_error(error: CliExecutionError) -> String {
412 match error {
413 CliExecutionError::CommandBuild(error) => {
414 format!("Failed to build one-shot agent command: {error}")
415 }
416 CliExecutionError::StdinBuild(error) => {
417 format!("Failed to build one-shot agent stdin payload: {error}")
418 }
419 error => format!("Failed to execute one-shot agent command: {error}"),
420 }
421}
422
423fn format_one_shot_exit_error(
425 agent_kind: AgentKind,
426 exit_code: Option<i32>,
427 stdout: &str,
428 stderr: &str,
429) -> String {
430 error::format_agent_cli_exit_error(
431 agent_kind,
432 "One-shot agent command",
433 exit_code,
434 stdout,
435 stderr,
436 )
437}
438
439fn clear_child_pid_slot(child_pid: Option<&Mutex<Option<u32>>>) {
441 let Some(child_pid) = child_pid else {
442 return;
443 };
444
445 if let Ok(mut guard) = child_pid.lock() {
446 *guard = None;
447 }
448}
449
450struct OneShotCliObserver {
452 child_pid: Option<Arc<Mutex<Option<u32>>>>,
453}
454
455impl CliExecutionObserver for OneShotCliObserver {
456 fn pid_updated(&self, active_child_pid: Option<u32>) {
457 let Some(child_pid_slot) = self.child_pid.as_deref() else {
458 return;
459 };
460
461 if let Ok(mut guard) = child_pid_slot.lock() {
462 *guard = active_child_pid;
463 }
464 }
465
466 fn stdout_line(&self, _line: &str) {}
467}
468
469#[cfg(test)]
470mod tests {
471 use std::path::Path;
472 use std::process::Command;
473 use std::time::Duration;
474
475 use tempfile::tempdir;
476
477 use super::*;
478 use crate::MockAgentBackend;
479 use crate::app_server::{AppServerError, AppServerTurnResponse, MockAppServerClient};
480
481 fn mock_shell_command(stdout: &str, stderr: &str, exit_code: i32) -> Command {
483 let mut command = Command::new("sh");
484 command.arg("-c").arg(
485 "printf '%s' \"$ONE_SHOT_STDOUT\"; printf '%s' \"$ONE_SHOT_STDERR\" >&2; exit \
486 \"$ONE_SHOT_EXIT\"",
487 );
488 command.env("ONE_SHOT_STDOUT", stdout);
489 command.env("ONE_SHOT_STDERR", stderr);
490 command.env("ONE_SHOT_EXIT", exit_code.to_string());
491 command.stdout(std::process::Stdio::piped());
492 command.stderr(std::process::Stdio::piped());
493
494 command
495 }
496
497 fn stdin_capture_shell_command(capture_path: &Path) -> Command {
499 let mut command = Command::new("sh");
500 command.arg("-c").arg(
501 "cat > \"$ONE_SHOT_CAPTURE_PATH\"; printf '%s' \
502 '{\"answer\":\"captured\",\"questions\":[],\"summary\":null}'",
503 );
504 command.env("ONE_SHOT_CAPTURE_PATH", capture_path);
505 command.stdout(std::process::Stdio::piped());
506 command.stderr(std::process::Stdio::piped());
507
508 command
509 }
510
511 #[test]
512 fn test_format_one_shot_execution_error_preserves_build_context() {
513 let command_error = CliExecutionError::CommandBuild(
515 crate::agent::AgentBackendError::CommandBuild("command".to_string()),
516 );
517 let stdin_error = CliExecutionError::StdinBuild(
518 crate::agent::AgentBackendError::CommandBuild("stdin".to_string()),
519 );
520 let execution_error = CliExecutionError::StdinWrite("write".to_string());
521
522 let command_message = format_one_shot_execution_error(command_error);
524 let stdin_message = format_one_shot_execution_error(stdin_error);
525 let execution_message = format_one_shot_execution_error(execution_error);
526
527 assert_eq!(
529 command_message,
530 "Failed to build one-shot agent command: command"
531 );
532 assert_eq!(
533 stdin_message,
534 "Failed to build one-shot agent stdin payload: stdin"
535 );
536 assert_eq!(
537 execution_message,
538 "Failed to execute one-shot agent command: stdin delivery failed: write"
539 );
540 }
541
542 #[test]
543 fn test_one_shot_cli_observer_updates_child_pid_slot() {
544 let child_pid = Arc::new(Mutex::new(None));
546 let observer = OneShotCliObserver {
547 child_pid: Some(Arc::clone(&child_pid)),
548 };
549
550 observer.pid_updated(Some(42));
552 let active_pid = *child_pid.lock().expect("PID lock should be available");
553 observer.stdout_line("collected output");
554 observer.pid_updated(None);
555 let cleared_pid = *child_pid.lock().expect("PID lock should be available");
556
557 assert_eq!(active_pid, Some(42));
559 assert_eq!(cleared_pid, None);
560 }
561
562 #[tokio::test]
563 async fn test_submit_one_shot_with_backend_reports_signal_interruption() {
564 let temp_directory = tempdir().expect("failed to create temp dir");
566 let mut backend = MockAgentBackend::new();
567 backend.expect_build_command().returning(|_| {
568 let mut command = Command::new("sh");
569 command.arg("-c").arg("kill -9 $$");
570
571 Ok(command)
572 });
573
574 let error = submit_one_shot_with_backend(
576 &backend,
577 OneShotRequest {
578 agent_kind: AgentKind::Codex,
579 child_pid: None,
580 folder: temp_directory.path().to_path_buf(),
581 model: AgentModel::Gpt56Sol,
582 prompt: "Generate title".to_string(),
583 request_kind: AgentRequestKind::UtilityPrompt,
584 reasoning_level: ReasoningLevel::default(),
585 },
586 )
587 .await
588 .expect_err("signal termination should interrupt the one-shot command");
589
590 assert_eq!(error, "One-shot agent command was interrupted");
592 }
593
594 #[tokio::test]
595 async fn test_submit_one_shot_with_backend_returns_protocol_response() {
597 let temp_directory = tempdir().expect("failed to create temp dir");
599 let mut backend = MockAgentBackend::new();
600 backend.expect_build_command().returning(|request| {
601 assert!(matches!(
602 request.request_kind,
603 AgentRequestKind::UtilityPrompt
604 ));
605 assert_eq!(request.prompt, "Generate title");
606
607 Ok(mock_shell_command(
608 r#"{"answer":"Generated title","questions":[],"summary":null}"#,
609 "",
610 0,
611 ))
612 });
613
614 let response = submit_one_shot_with_backend(
616 &backend,
617 OneShotRequest {
618 agent_kind: AgentKind::Claude,
619 child_pid: None,
620 folder: temp_directory.path().to_path_buf(),
621 model: AgentModel::ClaudeSonnet5,
622 prompt: "Generate title".to_string(),
623 request_kind: AgentRequestKind::UtilityPrompt,
624 reasoning_level: ReasoningLevel::default(),
625 },
626 )
627 .await
628 .expect("one-shot prompt should succeed");
629
630 assert_eq!(
632 response.response.answers(),
633 vec!["Generated title".to_string()]
634 );
635 }
636
637 #[tokio::test]
638 async fn test_submit_one_shot_with_backend_rejects_plain_text_utility_output() {
641 let temp_directory = tempdir().expect("failed to create temp dir");
643 let mut backend = MockAgentBackend::new();
644 backend
645 .expect_build_command()
646 .times(2)
647 .returning(|request| {
648 assert!(matches!(
649 request.request_kind,
650 AgentRequestKind::UtilityPrompt
651 ));
652
653 Ok(mock_shell_command("plain text", "", 0))
654 });
655
656 let error = submit_one_shot_with_backend(
658 &backend,
659 OneShotRequest {
660 agent_kind: AgentKind::Codex,
661 child_pid: None,
662 folder: temp_directory.path().to_path_buf(),
663 model: AgentModel::Gpt56Sol,
664 prompt: "Generate title".to_string(),
665 request_kind: AgentRequestKind::UtilityPrompt,
666 reasoning_level: ReasoningLevel::default(),
667 },
668 )
669 .await
670 .expect_err("plain-text utility output should fail");
671
672 assert!(error.contains("did not match the required JSON schema"));
674 assert!(error.contains("debug_details:"));
675 assert!(error.contains("direct_json_error_location: line 1, column 1"));
676 assert!(error.contains("response:\nplain text"));
677 }
678
679 #[tokio::test]
680 async fn test_submit_one_shot_with_backend_rejects_wrapped_plain_text_utility_output() {
683 let temp_directory = tempdir().expect("failed to create temp dir");
685 let mut backend = MockAgentBackend::new();
686 backend
687 .expect_build_command()
688 .times(2)
689 .returning(|request| {
690 assert!(matches!(
691 request.request_kind,
692 AgentRequestKind::UtilityPrompt
693 ));
694
695 Ok(mock_shell_command(
696 r#"{"result":"plain text","usage":{"input_tokens":2,"output_tokens":1}}"#,
697 "",
698 0,
699 ))
700 });
701
702 let error = submit_one_shot_with_backend(
704 &backend,
705 OneShotRequest {
706 agent_kind: AgentKind::Claude,
707 child_pid: None,
708 folder: temp_directory.path().to_path_buf(),
709 model: AgentModel::ClaudeSonnet5,
710 prompt: "Generate title".to_string(),
711 request_kind: AgentRequestKind::UtilityPrompt,
712 reasoning_level: ReasoningLevel::default(),
713 },
714 )
715 .await
716 .expect_err("wrapped plain-text utility output should fail");
717
718 assert!(error.contains("did not match the required JSON schema"));
721 assert!(error.contains("direct_json_error:"));
722 assert!(error.contains("response:\nplain text"));
723 }
724
725 #[tokio::test]
726 async fn test_submit_one_shot_with_backend_recovers_wrapped_protocol_output() {
729 let temp_directory = tempdir().expect("failed to create temp dir");
731 let mut backend = MockAgentBackend::new();
732 backend
733 .expect_build_command()
734 .times(1)
735 .returning(|request| {
736 assert!(matches!(
737 request.request_kind,
738 AgentRequestKind::UtilityPrompt
739 ));
740 assert_eq!(request.prompt, "Generate title");
741
742 Ok(mock_shell_command(
743 concat!(
744 "Now I have full context.\n",
745 r#"{"answer":"Generated title","questions":[],"summary":null}"#
746 ),
747 "",
748 0,
749 ))
750 });
751
752 let response = submit_one_shot_with_backend(
754 &backend,
755 OneShotRequest {
756 agent_kind: AgentKind::Claude,
757 child_pid: None,
758 folder: temp_directory.path().to_path_buf(),
759 model: AgentModel::ClaudeSonnet5,
760 prompt: "Generate title".to_string(),
761 request_kind: AgentRequestKind::UtilityPrompt,
762 reasoning_level: ReasoningLevel::default(),
763 },
764 )
765 .await
766 .expect("wrapped protocol output should succeed");
767
768 assert_eq!(
770 response.response.answers(),
771 vec!["Generated title".to_string()]
772 );
773 }
774
775 #[tokio::test]
776 async fn test_submit_one_shot_with_backend_recovers_via_protocol_repair() {
779 let temp_directory = tempdir().expect("failed to create temp dir");
781 let call_counter = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
782 let mut backend = MockAgentBackend::new();
783 backend.expect_build_command().times(2).returning({
784 let counter = std::sync::Arc::clone(&call_counter);
785
786 move |_| {
787 let call_number = counter.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
788
789 if call_number == 0 {
790 Ok(mock_shell_command("plain text", "", 0))
791 } else {
792 Ok(mock_shell_command(
793 r#"{"answer":"Repaired title","questions":[],"summary":null}"#,
794 "",
795 0,
796 ))
797 }
798 }
799 });
800
801 let response = submit_one_shot_with_backend(
803 &backend,
804 OneShotRequest {
805 agent_kind: AgentKind::Codex,
806 child_pid: None,
807 folder: temp_directory.path().to_path_buf(),
808 model: AgentModel::Gpt56Sol,
809 prompt: "Generate title".to_string(),
810 request_kind: AgentRequestKind::UtilityPrompt,
811 reasoning_level: ReasoningLevel::default(),
812 },
813 )
814 .await
815 .expect("repair retry should succeed");
816
817 assert_eq!(
819 response.response.answers(),
820 vec!["Repaired title".to_string()]
821 );
822 }
823
824 #[tokio::test]
825 async fn test_submit_one_shot_with_backend_rejects_blank_utility_output() {
828 let temp_directory = tempdir().expect("failed to create temp dir");
830 let mut backend = MockAgentBackend::new();
831 backend.expect_build_command().returning(|request| {
832 assert!(matches!(
833 request.request_kind,
834 AgentRequestKind::UtilityPrompt
835 ));
836
837 Ok(mock_shell_command(" ", "", 0))
838 });
839
840 let error = submit_one_shot_with_backend(
842 &backend,
843 OneShotRequest {
844 agent_kind: AgentKind::Codex,
845 child_pid: None,
846 folder: temp_directory.path().to_path_buf(),
847 model: AgentModel::Gpt56Sol,
848 prompt: "Generate title".to_string(),
849 request_kind: AgentRequestKind::UtilityPrompt,
850 reasoning_level: ReasoningLevel::default(),
851 },
852 )
853 .await
854 .expect_err("blank utility output should fail");
855
856 assert!(error.contains("did not match the required JSON schema"));
858 assert!(error.contains("trimmed_len: 0 chars"));
859 assert!(error.contains("response:\n"));
860 }
861
862 #[tokio::test]
863 async fn test_submit_one_shot_with_backend_writes_large_stdin_concurrently() {
866 let temp_directory = tempdir().expect("failed to create temp dir");
868 let large_prompt = "x".repeat(512 * 1024);
869 let mut backend = MockAgentBackend::new();
870 backend.expect_build_command().returning(|_| {
871 let mut command = Command::new("sh");
872 command.arg("-c").arg(
873 "printf 'warming up\\n' >&2; sleep 0.1; cat >/dev/null; printf '%s' \
874 '{\"answer\":\"done\",\"questions\":[],\"summary\":null}'",
875 );
876 command.stdout(std::process::Stdio::piped());
877 command.stderr(std::process::Stdio::piped());
878
879 Ok(command)
880 });
881
882 let response = tokio::time::timeout(
884 Duration::from_secs(5),
885 submit_one_shot_with_backend(
886 &backend,
887 OneShotRequest {
888 agent_kind: AgentKind::Claude,
889 child_pid: None,
890 folder: temp_directory.path().to_path_buf(),
891 model: AgentModel::ClaudeSonnet5,
892 prompt: large_prompt.clone(),
893 request_kind: AgentRequestKind::UtilityPrompt,
894 reasoning_level: ReasoningLevel::default(),
895 },
896 ),
897 )
898 .await
899 .expect("one-shot prompt should not deadlock")
900 .expect("one-shot prompt should succeed");
901
902 assert_eq!(response.response.answers(), vec!["done".to_string()]);
904 }
905
906 #[tokio::test]
907 async fn test_submit_one_shot_with_backend_writes_prompt_to_stdin() {
910 let temp_directory = tempdir().expect("failed to create temp dir");
912 let capture_path = temp_directory.path().join("stdin.txt");
913 let mut backend = MockAgentBackend::new();
914 backend.expect_build_command().returning({
915 let capture_path = capture_path.clone();
916
917 move |_| Ok(stdin_capture_shell_command(&capture_path))
918 });
919
920 let response = submit_one_shot_with_backend(
922 &backend,
923 OneShotRequest {
924 agent_kind: AgentKind::Claude,
925 child_pid: None,
926 folder: temp_directory.path().to_path_buf(),
927 model: AgentModel::ClaudeSonnet5,
928 prompt: "Generate title".to_string(),
929 request_kind: AgentRequestKind::UtilityPrompt,
930 reasoning_level: ReasoningLevel::default(),
931 },
932 )
933 .await
934 .expect("one-shot prompt should succeed");
935 let captured_prompt =
936 std::fs::read_to_string(&capture_path).expect("captured stdin payload should exist");
937
938 assert_eq!(response.response.answers(), vec!["captured".to_string()]);
940 assert!(captured_prompt.contains("Structured response protocol:"));
941 assert!(captured_prompt.contains("Generate title"));
942 }
943
944 #[tokio::test]
945 async fn test_submit_one_shot_with_backend_preserves_exit_error_after_broken_pipe() {
948 let temp_directory = tempdir().expect("failed to create temp dir");
950 let large_prompt = "x".repeat(512 * 1024);
951 let mut backend = MockAgentBackend::new();
952 backend.expect_build_command().returning(|_| {
953 let mut command = Command::new("sh");
954 command.arg("-c").arg("printf 'auth failed' >&2; exit 7");
955 command.stdout(std::process::Stdio::piped());
956 command.stderr(std::process::Stdio::piped());
957
958 Ok(command)
959 });
960
961 let error = submit_one_shot_with_backend(
963 &backend,
964 OneShotRequest {
965 agent_kind: AgentKind::Claude,
966 child_pid: None,
967 folder: temp_directory.path().to_path_buf(),
968 model: AgentModel::ClaudeSonnet5,
969 prompt: large_prompt,
970 request_kind: AgentRequestKind::UtilityPrompt,
971 reasoning_level: ReasoningLevel::default(),
972 },
973 )
974 .await
975 .expect_err("one-shot prompt should surface the child exit");
976
977 assert!(error.contains("exit code 7"), "error was: {error}");
979 assert!(error.contains("auth failed"), "error was: {error}");
980 assert!(
981 !error.contains("stdin payload"),
982 "stdin write error should not mask child failure: {error}"
983 );
984 }
985
986 #[tokio::test]
987 async fn test_submit_one_shot_with_backend_surfaces_claude_auth_guidance() {
990 let temp_directory = tempdir().expect("failed to create temp dir");
992 let mut backend = MockAgentBackend::new();
993 backend.expect_build_command().returning(|_| {
994 Ok(mock_shell_command(
995 r#"{"type":"error","error":{"type":"authentication_error","message":"OAuth token has expired. Please obtain a new token or refresh your existing token."}}"#,
996 "",
997 1,
998 ))
999 });
1000
1001 let error = submit_one_shot_with_backend(
1003 &backend,
1004 OneShotRequest {
1005 agent_kind: AgentKind::Claude,
1006 child_pid: None,
1007 folder: temp_directory.path().to_path_buf(),
1008 model: AgentModel::ClaudeSonnet5,
1009 prompt: "Generate title".to_string(),
1010 request_kind: AgentRequestKind::UtilityPrompt,
1011 reasoning_level: ReasoningLevel::default(),
1012 },
1013 )
1014 .await
1015 .expect_err("expired Claude auth should fail");
1016
1017 assert!(
1019 error.contains("One-shot agent command failed because Claude authentication expired")
1020 );
1021 assert!(error.contains("`claude auth login`"));
1022 assert!(error.contains("`claude auth status`"));
1023 }
1024
1025 #[tokio::test]
1026 async fn test_submit_one_shot_with_app_server_client_returns_protocol_response() {
1029 let temp_directory = tempdir().expect("failed to create temp dir");
1031 let mut app_server_client = MockAppServerClient::new();
1032 app_server_client
1033 .expect_run_turn()
1034 .times(1)
1035 .returning(|request, _| {
1036 assert_eq!(request.model, AgentModel::Gpt56Sol.as_str());
1037 assert!(matches!(
1038 request.request_kind,
1039 AgentRequestKind::UtilityPrompt
1040 ));
1041 assert_eq!(request.prompt.text, "Generate title");
1042
1043 Box::pin(async {
1044 Ok(AppServerTurnResponse {
1045 assistant_message:
1046 r#"{"answer":"Generated title","questions":[],"summary":null}"#
1047 .to_string(),
1048 context_reset: false,
1049 input_tokens: 11,
1050 output_tokens: 7,
1051 pid: Some(42),
1052 provider_conversation_id: Some("thread-1".to_string()),
1053 })
1054 })
1055 });
1056 app_server_client
1057 .expect_shutdown_session()
1058 .times(1)
1059 .returning(|_| Box::pin(async {}));
1060
1061 let response = submit_one_shot_with_app_server_client(
1063 &app_server_client,
1064 OneShotRequest {
1065 agent_kind: AgentKind::Codex,
1066 child_pid: None,
1067 folder: temp_directory.path().to_path_buf(),
1068 model: AgentModel::Gpt56Sol,
1069 prompt: "Generate title".to_string(),
1070 request_kind: AgentRequestKind::UtilityPrompt,
1071 reasoning_level: ReasoningLevel::default(),
1072 },
1073 )
1074 .await
1075 .expect("one-shot prompt should succeed");
1076
1077 assert_eq!(
1079 response.response.answers(),
1080 vec!["Generated title".to_string()]
1081 );
1082 assert_eq!(response.stats.input_tokens, 11);
1083 assert_eq!(response.stats.output_tokens, 7);
1084 }
1085
1086 #[tokio::test]
1087 async fn test_submit_one_shot_with_app_server_client_clears_pid_after_turn_failure() {
1090 let temp_directory = tempdir().expect("failed to create temp dir");
1092 let child_pid = Arc::new(Mutex::new(Some(42)));
1093 let mut app_server_client = MockAppServerClient::new();
1094 app_server_client
1095 .expect_run_turn()
1096 .times(1)
1097 .returning(|_, _| {
1098 Box::pin(async {
1099 Err(AppServerError::Provider(
1100 "app-server turn failed".to_string(),
1101 ))
1102 })
1103 });
1104 app_server_client
1105 .expect_shutdown_session()
1106 .times(1)
1107 .returning(|_| Box::pin(async {}));
1108
1109 let error = submit_one_shot_with_app_server_client(
1111 &app_server_client,
1112 OneShotRequest {
1113 agent_kind: AgentKind::Codex,
1114 child_pid: Some(Arc::clone(&child_pid)),
1115 folder: temp_directory.path().to_path_buf(),
1116 model: AgentModel::Gpt56Sol,
1117 prompt: "Generate title".to_string(),
1118 request_kind: AgentRequestKind::UtilityPrompt,
1119 reasoning_level: ReasoningLevel::default(),
1120 },
1121 )
1122 .await
1123 .expect_err("app-server turn failure should surface");
1124
1125 assert!(error.contains("app-server turn failed"));
1127 assert_eq!(
1128 *child_pid.lock().expect("child pid lock should succeed"),
1129 None
1130 );
1131 }
1132
1133 #[tokio::test]
1134 async fn test_submit_one_shot_with_app_server_client_rejects_plain_text_utility_output() {
1138 let temp_directory = tempdir().expect("failed to create temp dir");
1140 let mut app_server_client = MockAppServerClient::new();
1141 app_server_client
1142 .expect_run_turn()
1143 .times(2)
1144 .returning(|request, _| {
1145 assert_eq!(request.model, AgentModel::Gpt56Sol.as_str());
1146
1147 Box::pin(async {
1148 Ok(AppServerTurnResponse {
1149 assistant_message: "plain text".to_string(),
1150 context_reset: false,
1151 input_tokens: 2,
1152 output_tokens: 1,
1153 pid: None,
1154 provider_conversation_id: None,
1155 })
1156 })
1157 });
1158 app_server_client
1159 .expect_shutdown_session()
1160 .times(1)
1161 .returning(|_| Box::pin(async {}));
1162
1163 let error = submit_one_shot_with_app_server_client(
1165 &app_server_client,
1166 OneShotRequest {
1167 agent_kind: AgentKind::Codex,
1168 child_pid: None,
1169 folder: temp_directory.path().to_path_buf(),
1170 model: AgentModel::Gpt56Sol,
1171 prompt: "Generate title".to_string(),
1172 request_kind: AgentRequestKind::UtilityPrompt,
1173 reasoning_level: ReasoningLevel::default(),
1174 },
1175 )
1176 .await
1177 .expect_err("plain-text utility output should fail");
1178
1179 assert!(error.contains("did not match the required JSON schema"));
1181 assert!(error.contains("debug_details:"));
1182 assert!(error.contains("response:\nplain text"));
1183 }
1184
1185 #[tokio::test]
1186 async fn test_submit_one_shot_with_app_server_client_rejects_plain_text_non_utility_output() {
1190 let temp_directory = tempdir().expect("failed to create temp dir");
1192 let mut app_server_client = MockAppServerClient::new();
1193 app_server_client
1194 .expect_run_turn()
1195 .times(2)
1196 .returning(|request, _| {
1197 assert!(matches!(
1198 request.request_kind,
1199 AgentRequestKind::SessionStart
1200 ));
1201
1202 Box::pin(async {
1203 Ok(AppServerTurnResponse {
1204 assistant_message: "plain text".to_string(),
1205 context_reset: false,
1206 input_tokens: 2,
1207 output_tokens: 1,
1208 pid: None,
1209 provider_conversation_id: None,
1210 })
1211 })
1212 });
1213 app_server_client
1214 .expect_shutdown_session()
1215 .times(1)
1216 .returning(|_| Box::pin(async {}));
1217
1218 let error = submit_one_shot_with_app_server_client(
1220 &app_server_client,
1221 OneShotRequest {
1222 agent_kind: AgentKind::Codex,
1223 child_pid: None,
1224 folder: temp_directory.path().to_path_buf(),
1225 model: AgentModel::Gpt56Sol,
1226 prompt: "Generate title".to_string(),
1227 request_kind: AgentRequestKind::SessionStart,
1228 reasoning_level: ReasoningLevel::default(),
1229 },
1230 )
1231 .await
1232 .expect_err("invalid non-utility output should fail");
1233
1234 assert!(error.contains("did not match the required JSON schema"));
1236 assert!(error.contains("debug_details:"));
1237 assert!(error.contains("response:\nplain text"));
1238 }
1239}