Skip to main content

ag_agent/agent/
submission.rs

1//! One-shot agent prompt execution helpers.
2//!
3//! These helpers run isolated utility prompts outside the long-lived session
4//! turn flow. They require the shared structured response protocol on every
5//! transport so one-shot callers enforce the same schema contract as normal
6//! session turns.
7
8use 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/// Input payload for one isolated prompt that prefers structured protocol
29/// output.
30#[derive(Clone, Debug)]
31pub struct OneShotRequest {
32    /// Provider backend used for command construction, stdin shaping, and
33    /// response parsing.
34    pub agent_kind: AgentKind,
35    /// Optional PID slot used by cancel/stop flows to terminate the spawned
36    /// subprocess while a one-shot prompt is running.
37    pub child_pid: Option<Arc<Mutex<Option<u32>>>>,
38    /// Working directory where the prompt command runs.
39    pub folder: PathBuf,
40    /// Provider-specific model used for command construction and parsing.
41    pub model: AgentModel,
42    /// Prompt text submitted to the agent.
43    pub prompt: String,
44    /// Canonical request kind for this isolated prompt.
45    pub request_kind: AgentRequestKind,
46    /// Reasoning effort preference for the one-shot prompt.
47    pub reasoning_level: ReasoningLevel,
48}
49
50/// Parsed result returned by one isolated prompt execution.
51#[derive(Clone, Debug, PartialEq)]
52pub struct OneShotSubmission {
53    /// Structured protocol response parsed from the final successful attempt.
54    pub response: AgentResponse,
55    /// Aggregated token usage for the one-shot prompt execution.
56    pub stats: SessionStats,
57}
58
59/// Typed failure returned by [`OneShotClient`] submissions.
60///
61/// The concrete transport, protocol-repair, and provider diagnostics remain
62/// available through [`std::fmt::Display`] without exposing transport-specific
63/// variants to callers.
64#[derive(Clone, Debug, Eq, PartialEq, thiserror::Error)]
65#[error("{message}")]
66pub struct OneShotError {
67    message: String,
68}
69
70impl OneShotError {
71    /// Creates an error from one already formatted submission diagnostic.
72    pub fn new(message: impl Into<String>) -> Self {
73        Self {
74            message: message.into(),
75        }
76    }
77}
78
79/// Provider-neutral boundary for isolated structured agent prompts.
80///
81/// Implementations own transport selection, protocol repair, temporary
82/// app-server lifecycle, and usage aggregation so callers submit one request
83/// without selecting a CLI or app-server execution helper.
84#[cfg_attr(any(test, feature = "test-utils"), mockall::automock)]
85#[async_trait]
86pub trait OneShotClient: Send + Sync {
87    /// Executes one isolated prompt and returns its parsed response and usage.
88    async fn submit(&self, request: OneShotRequest) -> Result<OneShotSubmission, OneShotError>;
89}
90
91/// Production [`OneShotClient`] that routes through the selected provider.
92pub struct RealOneShotClient {
93    app_server_client_override: Option<Arc<dyn AppServerClient>>,
94}
95
96impl RealOneShotClient {
97    /// Creates a client with an optional app-server override.
98    ///
99    /// Production passes `None` so each provider supplies its native client;
100    /// deterministic environments may inject a shared app-server boundary.
101    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
119/// Executes one isolated prompt and returns the parsed response plus
120/// aggregated usage statistics, optionally overriding the backend-owned
121/// app-server client.
122///
123/// # Errors
124/// Returns an error when command construction fails, process execution fails,
125/// or the final output is empty or otherwise unusable.
126async 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
149/// Executes one isolated prompt through the shared app-server transport.
150///
151/// The temporary app-server session is shut down after the utility prompt
152/// finishes so one-shot helpers do not keep a provider runtime alive after the
153/// result has been parsed.
154///
155/// # Errors
156/// Returns an error when app-server turn execution fails or the final output
157/// is empty or otherwise unusable.
158async 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
229/// Executes one isolated prompt using the provided backend.
230///
231/// This shared helper keeps process execution behind the existing
232/// `AgentBackend` trait boundary so production callers and tests can reuse
233/// the same one-shot parsing path.
234///
235/// # Errors
236/// Returns an error when command construction fails, process execution fails,
237/// or the final output is empty or otherwise unusable.
238async 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
276/// Parses one one-shot response strictly against the shared protocol schema.
277///
278/// # Errors
279/// Returns an error when the response is empty or not valid protocol JSON. The
280/// error carries the parse reason and derived diagnostics only, never the
281/// provider payload itself.
282fn 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
292/// Attempts one protocol-repair retry through the app-server transport for
293/// a one-shot prompt whose initial response failed schema validation.
294///
295/// The repair prompt is sent as a follow-up turn on the same session so the
296/// agent retains the original conversation context. The initial turn's
297/// `provider_conversation_id` is threaded through so providers that depend
298/// on conversation state can continue the same thread.
299///
300/// Returns the parsed response together with the repair turn's token usage
301/// so the caller can aggregate stats across both attempts.
302///
303/// # Errors
304/// Returns the combined original and repair error when the retry fails.
305async 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
349/// Runs one one-shot backend command and returns the parsed provider content.
350///
351/// The spawned child is configured with `kill_on_drop(true)` so timeout-driven
352/// callers do not leave orphaned agent CLI processes behind when the future is
353/// canceled before completion.
354///
355/// # Errors
356/// Returns an error when the command cannot be built, run, or exits
357/// unsuccessfully.
358async 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
403/// Preserves the established one-shot context around shared execution errors.
404fn 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
416/// Formats one non-zero one-shot command exit into a user-facing error.
417fn 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
432/// Clears the shared one-shot child PID slot when one exists.
433fn 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
443/// Bridges shared CLI PID observations into the one-shot cancellation slot.
444struct 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    /// Builds one shell command that emits controlled stdout/stderr and exits.
475    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    /// Builds one shell command that captures stdin before returning JSON.
491    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        // Arrange
507        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        // Act
516        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
521        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        // Arrange
538        let child_pid = Arc::new(Mutex::new(None));
539        let observer = OneShotCliObserver {
540            child_pid: Some(Arc::clone(&child_pid)),
541        };
542
543        // Act
544        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
551        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        // Arrange
558        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        // Act
568        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
584        assert_eq!(error, "One-shot agent command was interrupted");
585    }
586
587    #[tokio::test]
588    /// Verifies one-shot execution returns the parsed structured answer.
589    async fn test_submit_one_shot_with_backend_returns_protocol_response() {
590        // Arrange
591        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        // Act
608        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
624        assert_eq!(
625            response.response.answers(),
626            vec!["Generated title".to_string()]
627        );
628    }
629
630    #[tokio::test]
631    /// Verifies one-shot execution rejects plain-text utility output after
632    /// both the original parse and the protocol-repair retry fail.
633    async fn test_submit_one_shot_with_backend_rejects_plain_text_utility_output() {
634        // Arrange
635        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        // Act
650        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
666        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    /// Verifies one-shot execution rejects wrapped non-schema utility output
674    /// after both the original parse and the protocol-repair retry fail.
675    async fn test_submit_one_shot_with_backend_rejects_wrapped_plain_text_utility_output() {
676        // Arrange
677        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        // Act
696        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 — the provider parser extracts "plain text" from the
712        // `result` wrapper, so the protocol parser sees raw text, not JSON keys.
713        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    /// Verifies one-shot execution recovers a trailing protocol payload when
720    /// the provider prepends extra prose before the final JSON object.
721    async fn test_submit_one_shot_with_backend_recovers_wrapped_protocol_output() {
722        // Arrange
723        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        // Act
746        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
762        assert_eq!(
763            response.response.answers(),
764            vec!["Generated title".to_string()]
765        );
766    }
767
768    #[tokio::test]
769    /// Verifies one-shot execution recovers valid output when the initial
770    /// parse fails but the protocol-repair retry returns valid protocol JSON.
771    async fn test_submit_one_shot_with_backend_recovers_via_protocol_repair() {
772        // Arrange
773        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        // Act
795        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
811        assert_eq!(
812            response.response.answers(),
813            vec!["Repaired title".to_string()]
814        );
815    }
816
817    #[tokio::test]
818    /// Verifies one-shot execution still rejects blank utility responses
819    /// after both the original parse and the protocol-repair retry fail.
820    async fn test_submit_one_shot_with_backend_rejects_blank_utility_output() {
821        // Arrange
822        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        // Act
834        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
850        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    /// Verifies one-shot execution does not deadlock when the child delays
857    /// reading stdin until after it emits early stderr output.
858    async fn test_submit_one_shot_with_backend_writes_large_stdin_concurrently() {
859        // Arrange
860        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        // Act
876        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
896        assert_eq!(response.response.answers(), vec!["done".to_string()]);
897    }
898
899    #[tokio::test]
900    /// Verifies one-shot execution streams Claude prompts through stdin so
901    /// large review requests avoid argv length limits.
902    async fn test_submit_one_shot_with_backend_writes_prompt_to_stdin() {
903        // Arrange
904        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        // Act
914        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
932        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    /// Verifies a broken stdin pipe does not hide the child exit status or
939    /// stderr when the backend exits before reading the full prompt.
940    async fn test_submit_one_shot_with_backend_preserves_exit_error_after_broken_pipe() {
941        // Arrange
942        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        // Act
955        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
971        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    /// Verifies Claude authentication failures return actionable re-login
981    /// guidance instead of raw transport output.
982    async fn test_submit_one_shot_with_backend_surfaces_claude_auth_guidance() {
983        // Arrange
984        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        // Act
995        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
1011        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    /// Verifies app-server-backed one-shot execution returns the parsed
1020    /// structured answer and usage totals.
1021    async fn test_submit_one_shot_with_app_server_client_returns_protocol_response() {
1022        // Arrange
1023        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        // Act
1055        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
1071        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    /// Verifies app-server turn failures shut down the temporary session and
1081    /// clear the caller's shared child-process slot.
1082    async fn test_submit_one_shot_with_app_server_client_clears_pid_after_turn_failure() {
1083        // Arrange
1084        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        // Act
1103        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
1119        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    /// Verifies app-server-backed one-shot execution rejects plain-text
1128    /// utility output after both the original parse and the protocol-repair
1129    /// retry fail.
1130    async fn test_submit_one_shot_with_app_server_client_rejects_plain_text_utility_output() {
1131        // Arrange
1132        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        // Act
1157        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
1173        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    /// Verifies app-server-backed non-utility one-shot execution still
1180    /// rejects plain-text output after both the original parse and the
1181    /// protocol-repair retry fail.
1182    async fn test_submit_one_shot_with_app_server_client_rejects_plain_text_non_utility_output() {
1183        // Arrange
1184        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        // Act
1212        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
1228        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}