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