1use std::os::unix::process::ExitStatusExt as _;
9use std::path::PathBuf;
10use std::sync::{Arc, Mutex};
11
12use ag_protocol::{
13 AgentResponse, build_protocol_repair_prompt, format_protocol_parse_debug_details,
14 parse_agent_response_strict,
15};
16use async_trait::async_trait;
17
18use super::backend::{AgentBackend, BuildCommandRequest};
19use super::cli::{error, stdin};
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 command = backend
372 .build_command(build_request)
373 .map_err(|error| format!("Failed to build one-shot agent command: {error}"))?;
374 let stdin_payload = super::build_command_stdin_payload(request.agent_kind, build_request)
375 .map_err(|error| format!("Failed to build one-shot agent stdin payload: {error}"))?;
376 let mut tokio_command = tokio::process::Command::from(command);
377 tokio_command
378 .stdin(if stdin_payload.is_some() {
379 std::process::Stdio::piped()
380 } else {
381 std::process::Stdio::null()
382 })
383 .kill_on_drop(true);
384 let mut pid_guard = ChildPidGuard::new(request.child_pid);
385 let mut child = tokio_command
386 .spawn()
387 .map_err(|error| format!("Failed to execute one-shot agent command: {error}"))?;
388 pid_guard.update_from_child(&child);
389 let stdin_write_task = stdin::spawn_optional_stdin_write(
390 child.stdin.take(),
391 stdin_payload,
392 "one-shot stdin pipe unavailable after spawn",
393 std::convert::identity,
394 );
395 let output = child
396 .wait_with_output()
397 .await
398 .map_err(|error| format!("Failed to execute one-shot agent command: {error}"))?;
399 stdin::await_optional_stdin_write(
400 stdin_write_task,
401 "One-shot stdin write task failed",
402 std::convert::identity,
403 )
404 .await?;
405
406 if output.status.signal().is_some() {
407 return Err("One-shot agent command was interrupted".to_string());
408 }
409
410 let stdout_text = String::from_utf8_lossy(&output.stdout).into_owned();
411 let stderr_text = String::from_utf8_lossy(&output.stderr).into_owned();
412 if !output.status.success() {
413 return Err(format_one_shot_exit_error(
414 request.agent_kind,
415 output.status.code(),
416 &stdout_text,
417 &stderr_text,
418 ));
419 }
420
421 let parsed_response = parse_response(request.agent_kind, &stdout_text, &stderr_text);
422
423 Ok(parsed_response)
424}
425
426fn format_one_shot_exit_error(
428 agent_kind: AgentKind,
429 exit_code: Option<i32>,
430 stdout: &str,
431 stderr: &str,
432) -> String {
433 error::format_agent_cli_exit_error(
434 agent_kind,
435 "One-shot agent command",
436 exit_code,
437 stdout,
438 stderr,
439 )
440}
441
442fn clear_child_pid_slot(child_pid: Option<&Mutex<Option<u32>>>) {
444 let Some(child_pid) = child_pid else {
445 return;
446 };
447
448 if let Ok(mut guard) = child_pid.lock() {
449 *guard = None;
450 }
451}
452
453struct ChildPidGuard {
455 child_pid: Option<Arc<Mutex<Option<u32>>>>,
456}
457
458impl ChildPidGuard {
459 fn new(child_pid: Option<Arc<Mutex<Option<u32>>>>) -> Self {
461 Self { child_pid }
462 }
463
464 fn update_from_child(&mut self, child: &tokio::process::Child) {
466 let Some(pid) = child.id() else {
467 return;
468 };
469
470 let Some(child_pid) = self.child_pid.as_deref() else {
471 return;
472 };
473
474 if let Ok(mut guard) = child_pid.lock() {
475 *guard = Some(pid);
476 }
477 }
478}
479
480impl Drop for ChildPidGuard {
481 fn drop(&mut self) {
482 let Some(child_pid) = self.child_pid.as_deref() else {
483 return;
484 };
485
486 if let Ok(mut guard) = child_pid.lock() {
487 *guard = None;
488 }
489 }
490}
491
492#[cfg(test)]
493mod tests {
494 use std::path::Path;
495 use std::process::Command;
496 use std::time::Duration;
497
498 use tempfile::tempdir;
499
500 use super::*;
501 use crate::MockAgentBackend;
502 use crate::app_server::{AppServerError, AppServerTurnResponse, MockAppServerClient};
503
504 fn mock_shell_command(stdout: &str, stderr: &str, exit_code: i32) -> Command {
506 let mut command = Command::new("sh");
507 command.arg("-c").arg(
508 "printf '%s' \"$ONE_SHOT_STDOUT\"; printf '%s' \"$ONE_SHOT_STDERR\" >&2; exit \
509 \"$ONE_SHOT_EXIT\"",
510 );
511 command.env("ONE_SHOT_STDOUT", stdout);
512 command.env("ONE_SHOT_STDERR", stderr);
513 command.env("ONE_SHOT_EXIT", exit_code.to_string());
514 command.stdout(std::process::Stdio::piped());
515 command.stderr(std::process::Stdio::piped());
516
517 command
518 }
519
520 fn stdin_capture_shell_command(capture_path: &Path) -> Command {
522 let mut command = Command::new("sh");
523 command.arg("-c").arg(
524 "cat > \"$ONE_SHOT_CAPTURE_PATH\"; printf '%s' \
525 '{\"answer\":\"captured\",\"questions\":[],\"summary\":null}'",
526 );
527 command.env("ONE_SHOT_CAPTURE_PATH", capture_path);
528 command.stdout(std::process::Stdio::piped());
529 command.stderr(std::process::Stdio::piped());
530
531 command
532 }
533
534 #[tokio::test]
535 async fn test_submit_one_shot_with_backend_returns_protocol_response() {
537 let temp_directory = tempdir().expect("failed to create temp dir");
539 let mut backend = MockAgentBackend::new();
540 backend.expect_build_command().returning(|request| {
541 assert!(matches!(
542 request.request_kind,
543 AgentRequestKind::UtilityPrompt
544 ));
545 assert_eq!(request.prompt, "Generate title");
546
547 Ok(mock_shell_command(
548 r#"{"answer":"Generated title","questions":[],"summary":null}"#,
549 "",
550 0,
551 ))
552 });
553
554 let response = submit_one_shot_with_backend(
556 &backend,
557 OneShotRequest {
558 agent_kind: AgentKind::Claude,
559 child_pid: None,
560 folder: temp_directory.path().to_path_buf(),
561 model: AgentModel::ClaudeSonnet5,
562 prompt: "Generate title".to_string(),
563 request_kind: AgentRequestKind::UtilityPrompt,
564 reasoning_level: ReasoningLevel::default(),
565 },
566 )
567 .await
568 .expect("one-shot prompt should succeed");
569
570 assert_eq!(
572 response.response.answers(),
573 vec!["Generated title".to_string()]
574 );
575 }
576
577 #[tokio::test]
578 async fn test_submit_one_shot_with_backend_rejects_plain_text_utility_output() {
581 let temp_directory = tempdir().expect("failed to create temp dir");
583 let mut backend = MockAgentBackend::new();
584 backend
585 .expect_build_command()
586 .times(2)
587 .returning(|request| {
588 assert!(matches!(
589 request.request_kind,
590 AgentRequestKind::UtilityPrompt
591 ));
592
593 Ok(mock_shell_command("plain text", "", 0))
594 });
595
596 let error = submit_one_shot_with_backend(
598 &backend,
599 OneShotRequest {
600 agent_kind: AgentKind::Codex,
601 child_pid: None,
602 folder: temp_directory.path().to_path_buf(),
603 model: AgentModel::Gpt55,
604 prompt: "Generate title".to_string(),
605 request_kind: AgentRequestKind::UtilityPrompt,
606 reasoning_level: ReasoningLevel::default(),
607 },
608 )
609 .await
610 .expect_err("plain-text utility output should fail");
611
612 assert!(error.contains("did not match the required JSON schema"));
614 assert!(error.contains("debug_details:"));
615 assert!(error.contains("direct_json_error_location: line 1, column 1"));
616 assert!(error.contains("response:\nplain text"));
617 }
618
619 #[tokio::test]
620 async fn test_submit_one_shot_with_backend_rejects_wrapped_plain_text_utility_output() {
623 let temp_directory = tempdir().expect("failed to create temp dir");
625 let mut backend = MockAgentBackend::new();
626 backend
627 .expect_build_command()
628 .times(2)
629 .returning(|request| {
630 assert!(matches!(
631 request.request_kind,
632 AgentRequestKind::UtilityPrompt
633 ));
634
635 Ok(mock_shell_command(
636 r#"{"result":"plain text","usage":{"input_tokens":2,"output_tokens":1}}"#,
637 "",
638 0,
639 ))
640 });
641
642 let error = submit_one_shot_with_backend(
644 &backend,
645 OneShotRequest {
646 agent_kind: AgentKind::Claude,
647 child_pid: None,
648 folder: temp_directory.path().to_path_buf(),
649 model: AgentModel::ClaudeSonnet5,
650 prompt: "Generate title".to_string(),
651 request_kind: AgentRequestKind::UtilityPrompt,
652 reasoning_level: ReasoningLevel::default(),
653 },
654 )
655 .await
656 .expect_err("wrapped plain-text utility output should fail");
657
658 assert!(error.contains("did not match the required JSON schema"));
661 assert!(error.contains("direct_json_error:"));
662 assert!(error.contains("response:\nplain text"));
663 }
664
665 #[tokio::test]
666 async fn test_submit_one_shot_with_backend_recovers_wrapped_protocol_output() {
669 let temp_directory = tempdir().expect("failed to create temp dir");
671 let mut backend = MockAgentBackend::new();
672 backend
673 .expect_build_command()
674 .times(1)
675 .returning(|request| {
676 assert!(matches!(
677 request.request_kind,
678 AgentRequestKind::UtilityPrompt
679 ));
680 assert_eq!(request.prompt, "Generate title");
681
682 Ok(mock_shell_command(
683 concat!(
684 "Now I have full context.\n",
685 r#"{"answer":"Generated title","questions":[],"summary":null}"#
686 ),
687 "",
688 0,
689 ))
690 });
691
692 let response = submit_one_shot_with_backend(
694 &backend,
695 OneShotRequest {
696 agent_kind: AgentKind::Claude,
697 child_pid: None,
698 folder: temp_directory.path().to_path_buf(),
699 model: AgentModel::ClaudeSonnet5,
700 prompt: "Generate title".to_string(),
701 request_kind: AgentRequestKind::UtilityPrompt,
702 reasoning_level: ReasoningLevel::default(),
703 },
704 )
705 .await
706 .expect("wrapped protocol output should succeed");
707
708 assert_eq!(
710 response.response.answers(),
711 vec!["Generated title".to_string()]
712 );
713 }
714
715 #[tokio::test]
716 async fn test_submit_one_shot_with_backend_recovers_via_protocol_repair() {
719 let temp_directory = tempdir().expect("failed to create temp dir");
721 let call_counter = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
722 let mut backend = MockAgentBackend::new();
723 backend.expect_build_command().times(2).returning({
724 let counter = std::sync::Arc::clone(&call_counter);
725
726 move |_| {
727 let call_number = counter.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
728
729 if call_number == 0 {
730 Ok(mock_shell_command("plain text", "", 0))
731 } else {
732 Ok(mock_shell_command(
733 r#"{"answer":"Repaired title","questions":[],"summary":null}"#,
734 "",
735 0,
736 ))
737 }
738 }
739 });
740
741 let response = submit_one_shot_with_backend(
743 &backend,
744 OneShotRequest {
745 agent_kind: AgentKind::Codex,
746 child_pid: None,
747 folder: temp_directory.path().to_path_buf(),
748 model: AgentModel::Gpt55,
749 prompt: "Generate title".to_string(),
750 request_kind: AgentRequestKind::UtilityPrompt,
751 reasoning_level: ReasoningLevel::default(),
752 },
753 )
754 .await
755 .expect("repair retry should succeed");
756
757 assert_eq!(
759 response.response.answers(),
760 vec!["Repaired title".to_string()]
761 );
762 }
763
764 #[tokio::test]
765 async fn test_submit_one_shot_with_backend_rejects_blank_utility_output() {
768 let temp_directory = tempdir().expect("failed to create temp dir");
770 let mut backend = MockAgentBackend::new();
771 backend.expect_build_command().returning(|request| {
772 assert!(matches!(
773 request.request_kind,
774 AgentRequestKind::UtilityPrompt
775 ));
776
777 Ok(mock_shell_command(" ", "", 0))
778 });
779
780 let error = submit_one_shot_with_backend(
782 &backend,
783 OneShotRequest {
784 agent_kind: AgentKind::Codex,
785 child_pid: None,
786 folder: temp_directory.path().to_path_buf(),
787 model: AgentModel::Gpt55,
788 prompt: "Generate title".to_string(),
789 request_kind: AgentRequestKind::UtilityPrompt,
790 reasoning_level: ReasoningLevel::default(),
791 },
792 )
793 .await
794 .expect_err("blank utility output should fail");
795
796 assert!(error.contains("did not match the required JSON schema"));
798 assert!(error.contains("trimmed_len: 0 chars"));
799 assert!(error.contains("response:\n"));
800 }
801
802 #[tokio::test]
803 async fn test_submit_one_shot_with_backend_writes_large_stdin_concurrently() {
806 let temp_directory = tempdir().expect("failed to create temp dir");
808 let large_prompt = "x".repeat(512 * 1024);
809 let mut backend = MockAgentBackend::new();
810 backend.expect_build_command().returning(|_| {
811 let mut command = Command::new("sh");
812 command.arg("-c").arg(
813 "printf 'warming up\\n' >&2; sleep 0.1; cat >/dev/null; printf '%s' \
814 '{\"answer\":\"done\",\"questions\":[],\"summary\":null}'",
815 );
816 command.stdout(std::process::Stdio::piped());
817 command.stderr(std::process::Stdio::piped());
818
819 Ok(command)
820 });
821
822 let response = tokio::time::timeout(
824 Duration::from_secs(5),
825 submit_one_shot_with_backend(
826 &backend,
827 OneShotRequest {
828 agent_kind: AgentKind::Claude,
829 child_pid: None,
830 folder: temp_directory.path().to_path_buf(),
831 model: AgentModel::ClaudeSonnet5,
832 prompt: large_prompt.clone(),
833 request_kind: AgentRequestKind::UtilityPrompt,
834 reasoning_level: ReasoningLevel::default(),
835 },
836 ),
837 )
838 .await
839 .expect("one-shot prompt should not deadlock")
840 .expect("one-shot prompt should succeed");
841
842 assert_eq!(response.response.answers(), vec!["done".to_string()]);
844 }
845
846 #[tokio::test]
847 async fn test_submit_one_shot_with_backend_writes_prompt_to_stdin() {
850 let temp_directory = tempdir().expect("failed to create temp dir");
852 let capture_path = temp_directory.path().join("stdin.txt");
853 let mut backend = MockAgentBackend::new();
854 backend.expect_build_command().returning({
855 let capture_path = capture_path.clone();
856
857 move |_| Ok(stdin_capture_shell_command(&capture_path))
858 });
859
860 let response = submit_one_shot_with_backend(
862 &backend,
863 OneShotRequest {
864 agent_kind: AgentKind::Claude,
865 child_pid: None,
866 folder: temp_directory.path().to_path_buf(),
867 model: AgentModel::ClaudeSonnet5,
868 prompt: "Generate title".to_string(),
869 request_kind: AgentRequestKind::UtilityPrompt,
870 reasoning_level: ReasoningLevel::default(),
871 },
872 )
873 .await
874 .expect("one-shot prompt should succeed");
875 let captured_prompt =
876 std::fs::read_to_string(&capture_path).expect("captured stdin payload should exist");
877
878 assert_eq!(response.response.answers(), vec!["captured".to_string()]);
880 assert!(captured_prompt.contains("Structured response protocol:"));
881 assert!(captured_prompt.contains("Generate title"));
882 }
883
884 #[tokio::test]
885 async fn test_submit_one_shot_with_backend_preserves_exit_error_after_broken_pipe() {
888 let temp_directory = tempdir().expect("failed to create temp dir");
890 let large_prompt = "x".repeat(512 * 1024);
891 let mut backend = MockAgentBackend::new();
892 backend.expect_build_command().returning(|_| {
893 let mut command = Command::new("sh");
894 command.arg("-c").arg("printf 'auth failed' >&2; exit 7");
895 command.stdout(std::process::Stdio::piped());
896 command.stderr(std::process::Stdio::piped());
897
898 Ok(command)
899 });
900
901 let error = submit_one_shot_with_backend(
903 &backend,
904 OneShotRequest {
905 agent_kind: AgentKind::Claude,
906 child_pid: None,
907 folder: temp_directory.path().to_path_buf(),
908 model: AgentModel::ClaudeSonnet5,
909 prompt: large_prompt,
910 request_kind: AgentRequestKind::UtilityPrompt,
911 reasoning_level: ReasoningLevel::default(),
912 },
913 )
914 .await
915 .expect_err("one-shot prompt should surface the child exit");
916
917 assert!(error.contains("exit code 7"), "error was: {error}");
919 assert!(error.contains("auth failed"), "error was: {error}");
920 assert!(
921 !error.contains("stdin payload"),
922 "stdin write error should not mask child failure: {error}"
923 );
924 }
925
926 #[tokio::test]
927 async fn test_submit_one_shot_with_backend_surfaces_claude_auth_guidance() {
930 let temp_directory = tempdir().expect("failed to create temp dir");
932 let mut backend = MockAgentBackend::new();
933 backend.expect_build_command().returning(|_| {
934 Ok(mock_shell_command(
935 r#"{"type":"error","error":{"type":"authentication_error","message":"OAuth token has expired. Please obtain a new token or refresh your existing token."}}"#,
936 "",
937 1,
938 ))
939 });
940
941 let error = submit_one_shot_with_backend(
943 &backend,
944 OneShotRequest {
945 agent_kind: AgentKind::Claude,
946 child_pid: None,
947 folder: temp_directory.path().to_path_buf(),
948 model: AgentModel::ClaudeSonnet5,
949 prompt: "Generate title".to_string(),
950 request_kind: AgentRequestKind::UtilityPrompt,
951 reasoning_level: ReasoningLevel::default(),
952 },
953 )
954 .await
955 .expect_err("expired Claude auth should fail");
956
957 assert!(
959 error.contains("One-shot agent command failed because Claude authentication expired")
960 );
961 assert!(error.contains("`claude auth login`"));
962 assert!(error.contains("`claude auth status`"));
963 }
964
965 #[tokio::test]
966 async fn test_child_pid_guard_tracks_and_clears_child_process() {
969 let child_pid = Arc::new(Mutex::new(None));
971 let mut child = tokio::process::Command::new("sh")
972 .arg("-c")
973 .arg("exit 0")
974 .spawn()
975 .expect("test child should spawn");
976 let expected_pid = child.id().expect("running child should expose its pid");
977 let mut guard = ChildPidGuard::new(Some(Arc::clone(&child_pid)));
978
979 guard.update_from_child(&child);
981 let tracked_pid = *child_pid.lock().expect("child pid lock should succeed");
982 child.wait().await.expect("test child should exit");
983 guard.update_from_child(&child);
984 drop(guard);
985 let cleared_pid = *child_pid.lock().expect("child pid lock should succeed");
986
987 assert_eq!(tracked_pid, Some(expected_pid));
989 assert_eq!(cleared_pid, None);
990 }
991
992 #[tokio::test]
993 async fn test_submit_one_shot_with_app_server_client_returns_protocol_response() {
996 let temp_directory = tempdir().expect("failed to create temp dir");
998 let mut app_server_client = MockAppServerClient::new();
999 app_server_client
1000 .expect_run_turn()
1001 .times(1)
1002 .returning(|request, _| {
1003 assert_eq!(request.model, AgentModel::Gpt55.as_str());
1004 assert!(matches!(
1005 request.request_kind,
1006 AgentRequestKind::UtilityPrompt
1007 ));
1008 assert_eq!(request.prompt.text, "Generate title");
1009
1010 Box::pin(async {
1011 Ok(AppServerTurnResponse {
1012 assistant_message:
1013 r#"{"answer":"Generated title","questions":[],"summary":null}"#
1014 .to_string(),
1015 context_reset: false,
1016 input_tokens: 11,
1017 output_tokens: 7,
1018 pid: Some(42),
1019 provider_conversation_id: Some("thread-1".to_string()),
1020 })
1021 })
1022 });
1023 app_server_client
1024 .expect_shutdown_session()
1025 .times(1)
1026 .returning(|_| Box::pin(async {}));
1027
1028 let response = submit_one_shot_with_app_server_client(
1030 &app_server_client,
1031 OneShotRequest {
1032 agent_kind: AgentKind::Codex,
1033 child_pid: None,
1034 folder: temp_directory.path().to_path_buf(),
1035 model: AgentModel::Gpt55,
1036 prompt: "Generate title".to_string(),
1037 request_kind: AgentRequestKind::UtilityPrompt,
1038 reasoning_level: ReasoningLevel::default(),
1039 },
1040 )
1041 .await
1042 .expect("one-shot prompt should succeed");
1043
1044 assert_eq!(
1046 response.response.answers(),
1047 vec!["Generated title".to_string()]
1048 );
1049 assert_eq!(response.stats.input_tokens, 11);
1050 assert_eq!(response.stats.output_tokens, 7);
1051 }
1052
1053 #[tokio::test]
1054 async fn test_submit_one_shot_with_app_server_client_clears_pid_after_turn_failure() {
1057 let temp_directory = tempdir().expect("failed to create temp dir");
1059 let child_pid = Arc::new(Mutex::new(Some(42)));
1060 let mut app_server_client = MockAppServerClient::new();
1061 app_server_client
1062 .expect_run_turn()
1063 .times(1)
1064 .returning(|_, _| {
1065 Box::pin(async {
1066 Err(AppServerError::Provider(
1067 "app-server turn failed".to_string(),
1068 ))
1069 })
1070 });
1071 app_server_client
1072 .expect_shutdown_session()
1073 .times(1)
1074 .returning(|_| Box::pin(async {}));
1075
1076 let error = submit_one_shot_with_app_server_client(
1078 &app_server_client,
1079 OneShotRequest {
1080 agent_kind: AgentKind::Codex,
1081 child_pid: Some(Arc::clone(&child_pid)),
1082 folder: temp_directory.path().to_path_buf(),
1083 model: AgentModel::Gpt55,
1084 prompt: "Generate title".to_string(),
1085 request_kind: AgentRequestKind::UtilityPrompt,
1086 reasoning_level: ReasoningLevel::default(),
1087 },
1088 )
1089 .await
1090 .expect_err("app-server turn failure should surface");
1091
1092 assert!(error.contains("app-server turn failed"));
1094 assert_eq!(
1095 *child_pid.lock().expect("child pid lock should succeed"),
1096 None
1097 );
1098 }
1099
1100 #[tokio::test]
1101 async fn test_submit_one_shot_with_app_server_client_rejects_plain_text_utility_output() {
1105 let temp_directory = tempdir().expect("failed to create temp dir");
1107 let mut app_server_client = MockAppServerClient::new();
1108 app_server_client
1109 .expect_run_turn()
1110 .times(2)
1111 .returning(|request, _| {
1112 assert_eq!(request.model, AgentModel::Gpt55.as_str());
1113
1114 Box::pin(async {
1115 Ok(AppServerTurnResponse {
1116 assistant_message: "plain text".to_string(),
1117 context_reset: false,
1118 input_tokens: 2,
1119 output_tokens: 1,
1120 pid: None,
1121 provider_conversation_id: None,
1122 })
1123 })
1124 });
1125 app_server_client
1126 .expect_shutdown_session()
1127 .times(1)
1128 .returning(|_| Box::pin(async {}));
1129
1130 let error = submit_one_shot_with_app_server_client(
1132 &app_server_client,
1133 OneShotRequest {
1134 agent_kind: AgentKind::Codex,
1135 child_pid: None,
1136 folder: temp_directory.path().to_path_buf(),
1137 model: AgentModel::Gpt55,
1138 prompt: "Generate title".to_string(),
1139 request_kind: AgentRequestKind::UtilityPrompt,
1140 reasoning_level: ReasoningLevel::default(),
1141 },
1142 )
1143 .await
1144 .expect_err("plain-text utility output should fail");
1145
1146 assert!(error.contains("did not match the required JSON schema"));
1148 assert!(error.contains("debug_details:"));
1149 assert!(error.contains("response:\nplain text"));
1150 }
1151
1152 #[tokio::test]
1153 async fn test_submit_one_shot_with_app_server_client_rejects_plain_text_non_utility_output() {
1157 let temp_directory = tempdir().expect("failed to create temp dir");
1159 let mut app_server_client = MockAppServerClient::new();
1160 app_server_client
1161 .expect_run_turn()
1162 .times(2)
1163 .returning(|request, _| {
1164 assert!(matches!(
1165 request.request_kind,
1166 AgentRequestKind::SessionStart
1167 ));
1168
1169 Box::pin(async {
1170 Ok(AppServerTurnResponse {
1171 assistant_message: "plain text".to_string(),
1172 context_reset: false,
1173 input_tokens: 2,
1174 output_tokens: 1,
1175 pid: None,
1176 provider_conversation_id: None,
1177 })
1178 })
1179 });
1180 app_server_client
1181 .expect_shutdown_session()
1182 .times(1)
1183 .returning(|_| Box::pin(async {}));
1184
1185 let error = submit_one_shot_with_app_server_client(
1187 &app_server_client,
1188 OneShotRequest {
1189 agent_kind: AgentKind::Codex,
1190 child_pid: None,
1191 folder: temp_directory.path().to_path_buf(),
1192 model: AgentModel::Gpt55,
1193 prompt: "Generate title".to_string(),
1194 request_kind: AgentRequestKind::SessionStart,
1195 reasoning_level: ReasoningLevel::default(),
1196 },
1197 )
1198 .await
1199 .expect_err("invalid non-utility output should fail");
1200
1201 assert!(error.contains("did not match the required JSON schema"));
1203 assert!(error.contains("debug_details:"));
1204 assert!(error.contains("response:\nplain text"));
1205 }
1206}