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