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::{SessionDiffState, 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 personality: crate::channel::PersonalityPrompt::default(),
172 prompt: ag_protocol::TurnPrompt::from_agent_data(request.prompt.clone()),
173 request_kind: request.request_kind.clone(),
174 replay_transcript: None,
175 provider_conversation_id: None,
176 persisted_instruction_conversation_id: None,
177 reasoning_level: request.reasoning_level,
178 session_id: session_id.clone(),
179 };
180
181 let turn_result = app_server_client.run_turn(turn_request, stream_tx).await;
182
183 let child_pid = request.child_pid.as_ref().map(Arc::clone);
184
185 let turn_result = match turn_result {
186 Ok(result) => result,
187 Err(error) => {
188 app_server_client.shutdown_session(session_id).await;
189 clear_child_pid_slot(child_pid.as_deref());
190
191 return Err(format!(
192 "Failed to execute one-shot app-server turn: {error}"
193 ));
194 }
195 };
196
197 let parse_result = match parse_one_shot_response(&turn_result.assistant_message) {
198 Ok(response) => Ok((response, 0, 0)),
199 Err(parse_error) => {
200 attempt_one_shot_app_server_repair(
201 app_server_client,
202 &parse_error,
203 &turn_result.assistant_message,
204 request,
205 &session_id,
206 turn_result.provider_conversation_id.as_deref(),
207 )
208 .await
209 }
210 };
211
212 app_server_client.shutdown_session(session_id).await;
213 clear_child_pid_slot(child_pid.as_deref());
214
215 let (response, repair_input_tokens, repair_output_tokens) = parse_result?;
216
217 Ok(OneShotSubmission {
218 response,
219 stats: SessionStats {
220 added_lines: 0,
221 deleted_lines: 0,
222 diff_state: SessionDiffState::Unknown,
223 input_tokens: turn_result.input_tokens + repair_input_tokens,
224 output_tokens: turn_result.output_tokens + repair_output_tokens,
225 },
226 })
227}
228
229async fn submit_one_shot_with_backend(
239 backend: &dyn AgentBackend,
240 request: OneShotRequest,
241) -> Result<OneShotSubmission, String> {
242 let parsed_response =
243 execute_one_shot_command(backend, &request.prompt, request.clone()).await?;
244 let (agent_response, repair_stats) = match parse_one_shot_response(&parsed_response.content) {
245 Ok(response) => (response, None),
246 Err(parse_error) => {
247 let repair_prompt =
248 build_protocol_repair_prompt(&parse_error, &parsed_response.content);
249 let repair_response = execute_one_shot_command(backend, &repair_prompt, request)
250 .await
251 .map_err(|error| format!("{parse_error}\nrepair transport failed: {error}"))?;
252
253 let response = parse_one_shot_response(&repair_response.content).map_err(|error| {
254 format!(
255 "{parse_error}\nrepair retry also failed: {error}\nrepair_response:\n{}",
256 repair_response.content
257 )
258 })?;
259
260 (response, Some(repair_response.stats))
261 }
262 };
263
264 let mut stats = parsed_response.stats;
265 if let Some(repair) = repair_stats {
266 stats.input_tokens += repair.input_tokens;
267 stats.output_tokens += repair.output_tokens;
268 }
269
270 Ok(OneShotSubmission {
271 response: agent_response,
272 stats,
273 })
274}
275
276fn parse_one_shot_response(content: &str) -> Result<AgentResponse, String> {
283 parse_agent_response_strict(content).map_err(|error| {
284 format!(
285 "One-shot agent output did not match the required JSON schema: \
286 {error}\ndebug_details:\n{}",
287 format_protocol_parse_debug_details(content)
288 )
289 })
290}
291
292async fn attempt_one_shot_app_server_repair(
306 app_server_client: &dyn AppServerClient,
307 parse_error: &str,
308 malformed_response: &str,
309 request: OneShotRequest,
310 session_id: &str,
311 provider_conversation_id: Option<&str>,
312) -> Result<(AgentResponse, u64, u64), String> {
313 let repair_prompt = build_protocol_repair_prompt(parse_error, malformed_response);
314
315 let (repair_stream_tx, _repair_stream_rx) = tokio::sync::mpsc::unbounded_channel();
316 let repair_turn_request = AppServerTurnRequest {
317 folder: request.folder,
318 live_transcript: None,
319 main_checkout_root: None,
320 model: request.model.provider_model_str().to_string(),
321 personality: crate::channel::PersonalityPrompt::default(),
322 prompt: ag_protocol::TurnPrompt::from_agent_data(repair_prompt),
323 request_kind: request.request_kind,
324 replay_transcript: None,
325 provider_conversation_id: provider_conversation_id.map(String::from),
326 persisted_instruction_conversation_id: None,
327 reasoning_level: request.reasoning_level,
328 session_id: session_id.to_string(),
329 };
330 let repair_result = app_server_client
331 .run_turn(repair_turn_request, repair_stream_tx)
332 .await
333 .map_err(|error| format!("{parse_error}\nrepair transport failed: {error}"))?;
334
335 let response = parse_one_shot_response(&repair_result.assistant_message).map_err(|error| {
336 format!(
337 "{parse_error}\nrepair retry also failed: {error}\nrepair_response:\n{}",
338 repair_result.assistant_message
339 )
340 })?;
341
342 Ok((
343 response,
344 repair_result.input_tokens,
345 repair_result.output_tokens,
346 ))
347}
348
349async fn execute_one_shot_command(
359 backend: &dyn AgentBackend,
360 prompt: &str,
361 request: OneShotRequest,
362) -> Result<ParsedResponse, String> {
363 let prompt_payload = ag_protocol::TurnPrompt::from_agent_data(prompt.to_string());
364 let build_request = BuildCommandRequest {
365 attachments: &prompt_payload.attachments,
366 folder: &request.folder,
367 main_checkout_root: None,
368 replay_transcript: None,
369 model: request.model.provider_model_str(),
370 personality_prompt: None,
371 prompt,
372 reasoning_level: request.reasoning_level,
373 request_kind: &request.request_kind,
374 };
375 let observer = OneShotCliObserver {
376 child_pid: request.child_pid,
377 };
378 let output =
379 execution::execute_cli_command(backend, request.agent_kind, build_request, &observer, None)
380 .await
381 .map_err(format_one_shot_execution_error)?;
382
383 match output.exit_status {
384 CliExitStatus::Signaled(_) => {
385 return Err("One-shot agent command was interrupted".to_string());
386 }
387 CliExitStatus::NonZero(exit_code) => {
388 return Err(format_one_shot_exit_error(
389 request.agent_kind,
390 exit_code,
391 &output.stdout,
392 &output.stderr,
393 ));
394 }
395 CliExitStatus::Success => {}
396 }
397
398 let parsed_response = parse_response(request.agent_kind, &output.stdout, &output.stderr);
399
400 Ok(parsed_response)
401}
402
403fn format_one_shot_execution_error(error: CliExecutionError) -> String {
405 match error {
406 CliExecutionError::CommandBuild(error) => {
407 format!("Failed to build one-shot agent command: {error}")
408 }
409 CliExecutionError::StdinBuild(error) => {
410 format!("Failed to build one-shot agent stdin payload: {error}")
411 }
412 error => format!("Failed to execute one-shot agent command: {error}"),
413 }
414}
415
416fn format_one_shot_exit_error(
418 agent_kind: AgentKind,
419 exit_code: Option<i32>,
420 stdout: &str,
421 stderr: &str,
422) -> String {
423 error::format_agent_cli_exit_error(
424 agent_kind,
425 "One-shot agent command",
426 exit_code,
427 stdout,
428 stderr,
429 )
430}
431
432fn clear_child_pid_slot(child_pid: Option<&Mutex<Option<u32>>>) {
434 let Some(child_pid) = child_pid else {
435 return;
436 };
437
438 if let Ok(mut guard) = child_pid.lock() {
439 *guard = None;
440 }
441}
442
443struct OneShotCliObserver {
445 child_pid: Option<Arc<Mutex<Option<u32>>>>,
446}
447
448impl CliExecutionObserver for OneShotCliObserver {
449 fn pid_updated(&self, active_child_pid: Option<u32>) {
450 let Some(child_pid_slot) = self.child_pid.as_deref() else {
451 return;
452 };
453
454 if let Ok(mut guard) = child_pid_slot.lock() {
455 *guard = active_child_pid;
456 }
457 }
458
459 fn stdout_line(&self, _line: &str) {}
460}
461
462#[cfg(test)]
463mod tests {
464 use std::path::Path;
465 use std::process::Command;
466 use std::time::Duration;
467
468 use tempfile::tempdir;
469
470 use super::*;
471 use crate::MockAgentBackend;
472 use crate::app_server::{AppServerError, AppServerTurnResponse, MockAppServerClient};
473
474 fn mock_shell_command(stdout: &str, stderr: &str, exit_code: i32) -> Command {
476 let mut command = Command::new("sh");
477 command.arg("-c").arg(
478 "printf '%s' \"$ONE_SHOT_STDOUT\"; printf '%s' \"$ONE_SHOT_STDERR\" >&2; exit \
479 \"$ONE_SHOT_EXIT\"",
480 );
481 command.env("ONE_SHOT_STDOUT", stdout);
482 command.env("ONE_SHOT_STDERR", stderr);
483 command.env("ONE_SHOT_EXIT", exit_code.to_string());
484 command.stdout(std::process::Stdio::piped());
485 command.stderr(std::process::Stdio::piped());
486
487 command
488 }
489
490 fn stdin_capture_shell_command(capture_path: &Path) -> Command {
492 let mut command = Command::new("sh");
493 command.arg("-c").arg(
494 "cat > \"$ONE_SHOT_CAPTURE_PATH\"; printf '%s' \
495 '{\"answer\":\"captured\",\"questions\":[],\"summary\":null}'",
496 );
497 command.env("ONE_SHOT_CAPTURE_PATH", capture_path);
498 command.stdout(std::process::Stdio::piped());
499 command.stderr(std::process::Stdio::piped());
500
501 command
502 }
503
504 #[test]
505 fn test_format_one_shot_execution_error_preserves_build_context() {
506 let command_error = CliExecutionError::CommandBuild(
508 crate::agent::AgentBackendError::CommandBuild("command".to_string()),
509 );
510 let stdin_error = CliExecutionError::StdinBuild(
511 crate::agent::AgentBackendError::CommandBuild("stdin".to_string()),
512 );
513 let execution_error = CliExecutionError::StdinWrite("write".to_string());
514
515 let command_message = format_one_shot_execution_error(command_error);
517 let stdin_message = format_one_shot_execution_error(stdin_error);
518 let execution_message = format_one_shot_execution_error(execution_error);
519
520 assert_eq!(
522 command_message,
523 "Failed to build one-shot agent command: command"
524 );
525 assert_eq!(
526 stdin_message,
527 "Failed to build one-shot agent stdin payload: stdin"
528 );
529 assert_eq!(
530 execution_message,
531 "Failed to execute one-shot agent command: stdin delivery failed: write"
532 );
533 }
534
535 #[test]
536 fn test_one_shot_cli_observer_updates_child_pid_slot() {
537 let child_pid = Arc::new(Mutex::new(None));
539 let observer = OneShotCliObserver {
540 child_pid: Some(Arc::clone(&child_pid)),
541 };
542
543 observer.pid_updated(Some(42));
545 let active_pid = *child_pid.lock().expect("PID lock should be available");
546 observer.stdout_line("collected output");
547 observer.pid_updated(None);
548 let cleared_pid = *child_pid.lock().expect("PID lock should be available");
549
550 assert_eq!(active_pid, Some(42));
552 assert_eq!(cleared_pid, None);
553 }
554
555 #[tokio::test]
556 async fn test_submit_one_shot_with_backend_reports_signal_interruption() {
557 let temp_directory = tempdir().expect("failed to create temp dir");
559 let mut backend = MockAgentBackend::new();
560 backend.expect_build_command().returning(|_| {
561 let mut command = Command::new("sh");
562 command.arg("-c").arg("kill -9 $$");
563
564 Ok(command)
565 });
566
567 let error = submit_one_shot_with_backend(
569 &backend,
570 OneShotRequest {
571 agent_kind: AgentKind::Codex,
572 child_pid: None,
573 folder: temp_directory.path().to_path_buf(),
574 model: AgentModel::Gpt55,
575 prompt: "Generate title".to_string(),
576 request_kind: AgentRequestKind::UtilityPrompt,
577 reasoning_level: ReasoningLevel::default(),
578 },
579 )
580 .await
581 .expect_err("signal termination should interrupt the one-shot command");
582
583 assert_eq!(error, "One-shot agent command was interrupted");
585 }
586
587 #[tokio::test]
588 async fn test_submit_one_shot_with_backend_returns_protocol_response() {
590 let temp_directory = tempdir().expect("failed to create temp dir");
592 let mut backend = MockAgentBackend::new();
593 backend.expect_build_command().returning(|request| {
594 assert!(matches!(
595 request.request_kind,
596 AgentRequestKind::UtilityPrompt
597 ));
598 assert_eq!(request.prompt, "Generate title");
599
600 Ok(mock_shell_command(
601 r#"{"answer":"Generated title","questions":[],"summary":null}"#,
602 "",
603 0,
604 ))
605 });
606
607 let response = submit_one_shot_with_backend(
609 &backend,
610 OneShotRequest {
611 agent_kind: AgentKind::Claude,
612 child_pid: None,
613 folder: temp_directory.path().to_path_buf(),
614 model: AgentModel::ClaudeSonnet5,
615 prompt: "Generate title".to_string(),
616 request_kind: AgentRequestKind::UtilityPrompt,
617 reasoning_level: ReasoningLevel::default(),
618 },
619 )
620 .await
621 .expect("one-shot prompt should succeed");
622
623 assert_eq!(
625 response.response.answers(),
626 vec!["Generated title".to_string()]
627 );
628 }
629
630 #[tokio::test]
631 async fn test_submit_one_shot_with_backend_rejects_plain_text_utility_output() {
634 let temp_directory = tempdir().expect("failed to create temp dir");
636 let mut backend = MockAgentBackend::new();
637 backend
638 .expect_build_command()
639 .times(2)
640 .returning(|request| {
641 assert!(matches!(
642 request.request_kind,
643 AgentRequestKind::UtilityPrompt
644 ));
645
646 Ok(mock_shell_command("plain text", "", 0))
647 });
648
649 let error = submit_one_shot_with_backend(
651 &backend,
652 OneShotRequest {
653 agent_kind: AgentKind::Codex,
654 child_pid: None,
655 folder: temp_directory.path().to_path_buf(),
656 model: AgentModel::Gpt55,
657 prompt: "Generate title".to_string(),
658 request_kind: AgentRequestKind::UtilityPrompt,
659 reasoning_level: ReasoningLevel::default(),
660 },
661 )
662 .await
663 .expect_err("plain-text utility output should fail");
664
665 assert!(error.contains("did not match the required JSON schema"));
667 assert!(error.contains("debug_details:"));
668 assert!(error.contains("direct_json_error_location: line 1, column 1"));
669 assert!(error.contains("response:\nplain text"));
670 }
671
672 #[tokio::test]
673 async fn test_submit_one_shot_with_backend_rejects_wrapped_plain_text_utility_output() {
676 let temp_directory = tempdir().expect("failed to create temp dir");
678 let mut backend = MockAgentBackend::new();
679 backend
680 .expect_build_command()
681 .times(2)
682 .returning(|request| {
683 assert!(matches!(
684 request.request_kind,
685 AgentRequestKind::UtilityPrompt
686 ));
687
688 Ok(mock_shell_command(
689 r#"{"result":"plain text","usage":{"input_tokens":2,"output_tokens":1}}"#,
690 "",
691 0,
692 ))
693 });
694
695 let error = submit_one_shot_with_backend(
697 &backend,
698 OneShotRequest {
699 agent_kind: AgentKind::Claude,
700 child_pid: None,
701 folder: temp_directory.path().to_path_buf(),
702 model: AgentModel::ClaudeSonnet5,
703 prompt: "Generate title".to_string(),
704 request_kind: AgentRequestKind::UtilityPrompt,
705 reasoning_level: ReasoningLevel::default(),
706 },
707 )
708 .await
709 .expect_err("wrapped plain-text utility output should fail");
710
711 assert!(error.contains("did not match the required JSON schema"));
714 assert!(error.contains("direct_json_error:"));
715 assert!(error.contains("response:\nplain text"));
716 }
717
718 #[tokio::test]
719 async fn test_submit_one_shot_with_backend_recovers_wrapped_protocol_output() {
722 let temp_directory = tempdir().expect("failed to create temp dir");
724 let mut backend = MockAgentBackend::new();
725 backend
726 .expect_build_command()
727 .times(1)
728 .returning(|request| {
729 assert!(matches!(
730 request.request_kind,
731 AgentRequestKind::UtilityPrompt
732 ));
733 assert_eq!(request.prompt, "Generate title");
734
735 Ok(mock_shell_command(
736 concat!(
737 "Now I have full context.\n",
738 r#"{"answer":"Generated title","questions":[],"summary":null}"#
739 ),
740 "",
741 0,
742 ))
743 });
744
745 let response = submit_one_shot_with_backend(
747 &backend,
748 OneShotRequest {
749 agent_kind: AgentKind::Claude,
750 child_pid: None,
751 folder: temp_directory.path().to_path_buf(),
752 model: AgentModel::ClaudeSonnet5,
753 prompt: "Generate title".to_string(),
754 request_kind: AgentRequestKind::UtilityPrompt,
755 reasoning_level: ReasoningLevel::default(),
756 },
757 )
758 .await
759 .expect("wrapped protocol output should succeed");
760
761 assert_eq!(
763 response.response.answers(),
764 vec!["Generated title".to_string()]
765 );
766 }
767
768 #[tokio::test]
769 async fn test_submit_one_shot_with_backend_recovers_via_protocol_repair() {
772 let temp_directory = tempdir().expect("failed to create temp dir");
774 let call_counter = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
775 let mut backend = MockAgentBackend::new();
776 backend.expect_build_command().times(2).returning({
777 let counter = std::sync::Arc::clone(&call_counter);
778
779 move |_| {
780 let call_number = counter.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
781
782 if call_number == 0 {
783 Ok(mock_shell_command("plain text", "", 0))
784 } else {
785 Ok(mock_shell_command(
786 r#"{"answer":"Repaired title","questions":[],"summary":null}"#,
787 "",
788 0,
789 ))
790 }
791 }
792 });
793
794 let response = submit_one_shot_with_backend(
796 &backend,
797 OneShotRequest {
798 agent_kind: AgentKind::Codex,
799 child_pid: None,
800 folder: temp_directory.path().to_path_buf(),
801 model: AgentModel::Gpt55,
802 prompt: "Generate title".to_string(),
803 request_kind: AgentRequestKind::UtilityPrompt,
804 reasoning_level: ReasoningLevel::default(),
805 },
806 )
807 .await
808 .expect("repair retry should succeed");
809
810 assert_eq!(
812 response.response.answers(),
813 vec!["Repaired title".to_string()]
814 );
815 }
816
817 #[tokio::test]
818 async fn test_submit_one_shot_with_backend_rejects_blank_utility_output() {
821 let temp_directory = tempdir().expect("failed to create temp dir");
823 let mut backend = MockAgentBackend::new();
824 backend.expect_build_command().returning(|request| {
825 assert!(matches!(
826 request.request_kind,
827 AgentRequestKind::UtilityPrompt
828 ));
829
830 Ok(mock_shell_command(" ", "", 0))
831 });
832
833 let error = submit_one_shot_with_backend(
835 &backend,
836 OneShotRequest {
837 agent_kind: AgentKind::Codex,
838 child_pid: None,
839 folder: temp_directory.path().to_path_buf(),
840 model: AgentModel::Gpt55,
841 prompt: "Generate title".to_string(),
842 request_kind: AgentRequestKind::UtilityPrompt,
843 reasoning_level: ReasoningLevel::default(),
844 },
845 )
846 .await
847 .expect_err("blank utility output should fail");
848
849 assert!(error.contains("did not match the required JSON schema"));
851 assert!(error.contains("trimmed_len: 0 chars"));
852 assert!(error.contains("response:\n"));
853 }
854
855 #[tokio::test]
856 async fn test_submit_one_shot_with_backend_writes_large_stdin_concurrently() {
859 let temp_directory = tempdir().expect("failed to create temp dir");
861 let large_prompt = "x".repeat(512 * 1024);
862 let mut backend = MockAgentBackend::new();
863 backend.expect_build_command().returning(|_| {
864 let mut command = Command::new("sh");
865 command.arg("-c").arg(
866 "printf 'warming up\\n' >&2; sleep 0.1; cat >/dev/null; printf '%s' \
867 '{\"answer\":\"done\",\"questions\":[],\"summary\":null}'",
868 );
869 command.stdout(std::process::Stdio::piped());
870 command.stderr(std::process::Stdio::piped());
871
872 Ok(command)
873 });
874
875 let response = tokio::time::timeout(
877 Duration::from_secs(5),
878 submit_one_shot_with_backend(
879 &backend,
880 OneShotRequest {
881 agent_kind: AgentKind::Claude,
882 child_pid: None,
883 folder: temp_directory.path().to_path_buf(),
884 model: AgentModel::ClaudeSonnet5,
885 prompt: large_prompt.clone(),
886 request_kind: AgentRequestKind::UtilityPrompt,
887 reasoning_level: ReasoningLevel::default(),
888 },
889 ),
890 )
891 .await
892 .expect("one-shot prompt should not deadlock")
893 .expect("one-shot prompt should succeed");
894
895 assert_eq!(response.response.answers(), vec!["done".to_string()]);
897 }
898
899 #[tokio::test]
900 async fn test_submit_one_shot_with_backend_writes_prompt_to_stdin() {
903 let temp_directory = tempdir().expect("failed to create temp dir");
905 let capture_path = temp_directory.path().join("stdin.txt");
906 let mut backend = MockAgentBackend::new();
907 backend.expect_build_command().returning({
908 let capture_path = capture_path.clone();
909
910 move |_| Ok(stdin_capture_shell_command(&capture_path))
911 });
912
913 let response = submit_one_shot_with_backend(
915 &backend,
916 OneShotRequest {
917 agent_kind: AgentKind::Claude,
918 child_pid: None,
919 folder: temp_directory.path().to_path_buf(),
920 model: AgentModel::ClaudeSonnet5,
921 prompt: "Generate title".to_string(),
922 request_kind: AgentRequestKind::UtilityPrompt,
923 reasoning_level: ReasoningLevel::default(),
924 },
925 )
926 .await
927 .expect("one-shot prompt should succeed");
928 let captured_prompt =
929 std::fs::read_to_string(&capture_path).expect("captured stdin payload should exist");
930
931 assert_eq!(response.response.answers(), vec!["captured".to_string()]);
933 assert!(captured_prompt.contains("Structured response protocol:"));
934 assert!(captured_prompt.contains("Generate title"));
935 }
936
937 #[tokio::test]
938 async fn test_submit_one_shot_with_backend_preserves_exit_error_after_broken_pipe() {
941 let temp_directory = tempdir().expect("failed to create temp dir");
943 let large_prompt = "x".repeat(512 * 1024);
944 let mut backend = MockAgentBackend::new();
945 backend.expect_build_command().returning(|_| {
946 let mut command = Command::new("sh");
947 command.arg("-c").arg("printf 'auth failed' >&2; exit 7");
948 command.stdout(std::process::Stdio::piped());
949 command.stderr(std::process::Stdio::piped());
950
951 Ok(command)
952 });
953
954 let error = submit_one_shot_with_backend(
956 &backend,
957 OneShotRequest {
958 agent_kind: AgentKind::Claude,
959 child_pid: None,
960 folder: temp_directory.path().to_path_buf(),
961 model: AgentModel::ClaudeSonnet5,
962 prompt: large_prompt,
963 request_kind: AgentRequestKind::UtilityPrompt,
964 reasoning_level: ReasoningLevel::default(),
965 },
966 )
967 .await
968 .expect_err("one-shot prompt should surface the child exit");
969
970 assert!(error.contains("exit code 7"), "error was: {error}");
972 assert!(error.contains("auth failed"), "error was: {error}");
973 assert!(
974 !error.contains("stdin payload"),
975 "stdin write error should not mask child failure: {error}"
976 );
977 }
978
979 #[tokio::test]
980 async fn test_submit_one_shot_with_backend_surfaces_claude_auth_guidance() {
983 let temp_directory = tempdir().expect("failed to create temp dir");
985 let mut backend = MockAgentBackend::new();
986 backend.expect_build_command().returning(|_| {
987 Ok(mock_shell_command(
988 r#"{"type":"error","error":{"type":"authentication_error","message":"OAuth token has expired. Please obtain a new token or refresh your existing token."}}"#,
989 "",
990 1,
991 ))
992 });
993
994 let error = submit_one_shot_with_backend(
996 &backend,
997 OneShotRequest {
998 agent_kind: AgentKind::Claude,
999 child_pid: None,
1000 folder: temp_directory.path().to_path_buf(),
1001 model: AgentModel::ClaudeSonnet5,
1002 prompt: "Generate title".to_string(),
1003 request_kind: AgentRequestKind::UtilityPrompt,
1004 reasoning_level: ReasoningLevel::default(),
1005 },
1006 )
1007 .await
1008 .expect_err("expired Claude auth should fail");
1009
1010 assert!(
1012 error.contains("One-shot agent command failed because Claude authentication expired")
1013 );
1014 assert!(error.contains("`claude auth login`"));
1015 assert!(error.contains("`claude auth status`"));
1016 }
1017
1018 #[tokio::test]
1019 async fn test_submit_one_shot_with_app_server_client_returns_protocol_response() {
1022 let temp_directory = tempdir().expect("failed to create temp dir");
1024 let mut app_server_client = MockAppServerClient::new();
1025 app_server_client
1026 .expect_run_turn()
1027 .times(1)
1028 .returning(|request, _| {
1029 assert_eq!(request.model, AgentModel::Gpt55.as_str());
1030 assert!(matches!(
1031 request.request_kind,
1032 AgentRequestKind::UtilityPrompt
1033 ));
1034 assert_eq!(request.prompt.text, "Generate title");
1035
1036 Box::pin(async {
1037 Ok(AppServerTurnResponse {
1038 assistant_message:
1039 r#"{"answer":"Generated title","questions":[],"summary":null}"#
1040 .to_string(),
1041 context_reset: false,
1042 input_tokens: 11,
1043 output_tokens: 7,
1044 pid: Some(42),
1045 provider_conversation_id: Some("thread-1".to_string()),
1046 })
1047 })
1048 });
1049 app_server_client
1050 .expect_shutdown_session()
1051 .times(1)
1052 .returning(|_| Box::pin(async {}));
1053
1054 let response = submit_one_shot_with_app_server_client(
1056 &app_server_client,
1057 OneShotRequest {
1058 agent_kind: AgentKind::Codex,
1059 child_pid: None,
1060 folder: temp_directory.path().to_path_buf(),
1061 model: AgentModel::Gpt55,
1062 prompt: "Generate title".to_string(),
1063 request_kind: AgentRequestKind::UtilityPrompt,
1064 reasoning_level: ReasoningLevel::default(),
1065 },
1066 )
1067 .await
1068 .expect("one-shot prompt should succeed");
1069
1070 assert_eq!(
1072 response.response.answers(),
1073 vec!["Generated title".to_string()]
1074 );
1075 assert_eq!(response.stats.input_tokens, 11);
1076 assert_eq!(response.stats.output_tokens, 7);
1077 }
1078
1079 #[tokio::test]
1080 async fn test_submit_one_shot_with_app_server_client_clears_pid_after_turn_failure() {
1083 let temp_directory = tempdir().expect("failed to create temp dir");
1085 let child_pid = Arc::new(Mutex::new(Some(42)));
1086 let mut app_server_client = MockAppServerClient::new();
1087 app_server_client
1088 .expect_run_turn()
1089 .times(1)
1090 .returning(|_, _| {
1091 Box::pin(async {
1092 Err(AppServerError::Provider(
1093 "app-server turn failed".to_string(),
1094 ))
1095 })
1096 });
1097 app_server_client
1098 .expect_shutdown_session()
1099 .times(1)
1100 .returning(|_| Box::pin(async {}));
1101
1102 let error = submit_one_shot_with_app_server_client(
1104 &app_server_client,
1105 OneShotRequest {
1106 agent_kind: AgentKind::Codex,
1107 child_pid: Some(Arc::clone(&child_pid)),
1108 folder: temp_directory.path().to_path_buf(),
1109 model: AgentModel::Gpt55,
1110 prompt: "Generate title".to_string(),
1111 request_kind: AgentRequestKind::UtilityPrompt,
1112 reasoning_level: ReasoningLevel::default(),
1113 },
1114 )
1115 .await
1116 .expect_err("app-server turn failure should surface");
1117
1118 assert!(error.contains("app-server turn failed"));
1120 assert_eq!(
1121 *child_pid.lock().expect("child pid lock should succeed"),
1122 None
1123 );
1124 }
1125
1126 #[tokio::test]
1127 async fn test_submit_one_shot_with_app_server_client_rejects_plain_text_utility_output() {
1131 let temp_directory = tempdir().expect("failed to create temp dir");
1133 let mut app_server_client = MockAppServerClient::new();
1134 app_server_client
1135 .expect_run_turn()
1136 .times(2)
1137 .returning(|request, _| {
1138 assert_eq!(request.model, AgentModel::Gpt55.as_str());
1139
1140 Box::pin(async {
1141 Ok(AppServerTurnResponse {
1142 assistant_message: "plain text".to_string(),
1143 context_reset: false,
1144 input_tokens: 2,
1145 output_tokens: 1,
1146 pid: None,
1147 provider_conversation_id: None,
1148 })
1149 })
1150 });
1151 app_server_client
1152 .expect_shutdown_session()
1153 .times(1)
1154 .returning(|_| Box::pin(async {}));
1155
1156 let error = submit_one_shot_with_app_server_client(
1158 &app_server_client,
1159 OneShotRequest {
1160 agent_kind: AgentKind::Codex,
1161 child_pid: None,
1162 folder: temp_directory.path().to_path_buf(),
1163 model: AgentModel::Gpt55,
1164 prompt: "Generate title".to_string(),
1165 request_kind: AgentRequestKind::UtilityPrompt,
1166 reasoning_level: ReasoningLevel::default(),
1167 },
1168 )
1169 .await
1170 .expect_err("plain-text utility output should fail");
1171
1172 assert!(error.contains("did not match the required JSON schema"));
1174 assert!(error.contains("debug_details:"));
1175 assert!(error.contains("response:\nplain text"));
1176 }
1177
1178 #[tokio::test]
1179 async fn test_submit_one_shot_with_app_server_client_rejects_plain_text_non_utility_output() {
1183 let temp_directory = tempdir().expect("failed to create temp dir");
1185 let mut app_server_client = MockAppServerClient::new();
1186 app_server_client
1187 .expect_run_turn()
1188 .times(2)
1189 .returning(|request, _| {
1190 assert!(matches!(
1191 request.request_kind,
1192 AgentRequestKind::SessionStart
1193 ));
1194
1195 Box::pin(async {
1196 Ok(AppServerTurnResponse {
1197 assistant_message: "plain text".to_string(),
1198 context_reset: false,
1199 input_tokens: 2,
1200 output_tokens: 1,
1201 pid: None,
1202 provider_conversation_id: None,
1203 })
1204 })
1205 });
1206 app_server_client
1207 .expect_shutdown_session()
1208 .times(1)
1209 .returning(|_| Box::pin(async {}));
1210
1211 let error = submit_one_shot_with_app_server_client(
1213 &app_server_client,
1214 OneShotRequest {
1215 agent_kind: AgentKind::Codex,
1216 child_pid: None,
1217 folder: temp_directory.path().to_path_buf(),
1218 model: AgentModel::Gpt55,
1219 prompt: "Generate title".to_string(),
1220 request_kind: AgentRequestKind::SessionStart,
1221 reasoning_level: ReasoningLevel::default(),
1222 },
1223 )
1224 .await
1225 .expect_err("invalid non-utility output should fail");
1226
1227 assert!(error.contains("did not match the required JSON schema"));
1229 assert!(error.contains("debug_details:"));
1230 assert!(error.contains("response:\nplain text"));
1231 }
1232}