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::permission::PermissionMode;
27use crate::model::session::{SessionDiffState, SessionStats};
28
29/// Input payload for one isolated prompt that prefers structured protocol
30/// output.
31#[derive(Clone, Debug)]
32pub struct OneShotRequest {
33    /// Provider backend used for command construction, stdin shaping, and
34    /// response parsing.
35    pub agent_kind: AgentKind,
36    /// Optional PID slot used by cancel/stop flows to terminate the spawned
37    /// subprocess while a one-shot prompt is running.
38    pub child_pid: Option<Arc<Mutex<Option<u32>>>>,
39    /// Working directory where the prompt command runs.
40    pub folder: PathBuf,
41    /// Provider-specific model used for command construction and parsing.
42    pub model: AgentModel,
43    /// Filesystem and command permission policy for this isolated prompt.
44    pub permission_mode: PermissionMode,
45    /// Prompt text submitted to the agent.
46    pub prompt: String,
47    /// Canonical request kind for this isolated prompt.
48    pub request_kind: AgentRequestKind,
49    /// Reasoning effort preference for the one-shot prompt.
50    pub reasoning_level: ReasoningLevel,
51}
52
53/// Parsed result returned by one isolated prompt execution.
54#[derive(Clone, Debug, PartialEq)]
55pub struct OneShotSubmission {
56    /// Structured protocol response parsed from the final successful attempt.
57    pub response: AgentResponse,
58    /// Aggregated token usage for the one-shot prompt execution.
59    pub stats: SessionStats,
60}
61
62/// Typed failure returned by [`OneShotClient`] submissions.
63///
64/// The concrete transport, protocol-repair, and provider diagnostics remain
65/// available through [`std::fmt::Display`] without exposing transport-specific
66/// variants to callers.
67#[derive(Clone, Debug, Eq, PartialEq, thiserror::Error)]
68#[error("{message}")]
69pub struct OneShotError {
70    message: String,
71}
72
73impl OneShotError {
74    /// Creates an error from one already formatted submission diagnostic.
75    pub fn new(message: impl Into<String>) -> Self {
76        Self {
77            message: message.into(),
78        }
79    }
80}
81
82/// Provider-neutral boundary for isolated structured agent prompts.
83///
84/// Implementations own transport selection, protocol repair, temporary
85/// app-server lifecycle, and usage aggregation so callers submit one request
86/// without selecting a CLI or app-server execution helper.
87#[cfg_attr(any(test, feature = "test-utils"), mockall::automock)]
88#[async_trait]
89pub trait OneShotClient: Send + Sync {
90    /// Executes one isolated prompt and returns its parsed response and usage.
91    async fn submit(&self, request: OneShotRequest) -> Result<OneShotSubmission, OneShotError>;
92}
93
94/// Production [`OneShotClient`] that routes through the selected provider.
95pub struct RealOneShotClient {
96    app_server_client_override: Option<Arc<dyn AppServerClient>>,
97}
98
99impl RealOneShotClient {
100    /// Creates a client with an optional app-server override.
101    ///
102    /// Production passes `None` so each provider supplies its native client;
103    /// deterministic environments may inject a shared app-server boundary.
104    pub fn new(app_server_client_override: Option<Arc<dyn AppServerClient>>) -> Self {
105        Self {
106            app_server_client_override,
107        }
108    }
109}
110
111#[async_trait]
112impl OneShotClient for RealOneShotClient {
113    async fn submit(&self, request: OneShotRequest) -> Result<OneShotSubmission, OneShotError> {
114        let app_server_client_override = self.app_server_client_override.as_ref().map(Arc::clone);
115
116        submit_one_shot_with_stats_and_app_server_client(request, app_server_client_override)
117            .await
118            .map_err(OneShotError::new)
119    }
120}
121
122/// Executes one isolated prompt and returns the parsed response plus
123/// aggregated usage statistics, optionally overriding the backend-owned
124/// app-server client.
125///
126/// # Errors
127/// Returns an error when command construction fails, process execution fails,
128/// or the final output is empty or otherwise unusable.
129async fn submit_one_shot_with_stats_and_app_server_client(
130    request: OneShotRequest,
131    app_server_client_override: Option<Arc<dyn AppServerClient>>,
132) -> Result<OneShotSubmission, String> {
133    let backend = create_backend(request.agent_kind);
134
135    if transport_mode(request.agent_kind).uses_app_server() {
136        let app_server_client =
137            create_app_server_client(request.agent_kind, app_server_client_override).ok_or_else(
138                || {
139                    format!(
140                        "{} provider did not provide an app-server client",
141                        request.agent_kind
142                    )
143                },
144            )?;
145
146        return submit_one_shot_with_app_server_client(app_server_client.as_ref(), request).await;
147    }
148
149    submit_one_shot_with_backend(backend.as_ref(), request).await
150}
151
152/// Executes one isolated prompt through the shared app-server transport.
153///
154/// The temporary app-server session is shut down after the utility prompt
155/// finishes so one-shot helpers do not keep a provider runtime alive after the
156/// result has been parsed.
157///
158/// # Errors
159/// Returns an error when app-server turn execution fails or the final output
160/// is empty or otherwise unusable.
161async fn submit_one_shot_with_app_server_client(
162    app_server_client: &dyn AppServerClient,
163    request: OneShotRequest,
164) -> Result<OneShotSubmission, String> {
165    clear_child_pid_slot(request.child_pid.as_deref());
166
167    let session_id = format!("one-shot-{}", uuid::Uuid::new_v4());
168    let (stream_tx, _stream_rx) = tokio::sync::mpsc::unbounded_channel();
169    let turn_request = AppServerTurnRequest {
170        folder: request.folder.clone(),
171        live_transcript: None,
172        main_checkout_root: None,
173        model: request.model.provider_model_str().to_string(),
174        permission_mode: request.permission_mode,
175        personality: crate::channel::PersonalityPrompt::default(),
176        prompt: ag_protocol::TurnPrompt::from_agent_data(request.prompt.clone()),
177        request_kind: request.request_kind.clone(),
178        replay_transcript: None,
179        provider_conversation_id: None,
180        persisted_instruction_conversation_id: None,
181        reasoning_level: request.reasoning_level,
182        session_id: session_id.clone(),
183        speed_mode: crate::model::session::SpeedMode::default(),
184    };
185
186    let turn_result = app_server_client.run_turn(turn_request, stream_tx).await;
187
188    let child_pid = request.child_pid.as_ref().map(Arc::clone);
189
190    let turn_result = match turn_result {
191        Ok(result) => result,
192        Err(error) => {
193            app_server_client.shutdown_session(session_id).await;
194            clear_child_pid_slot(child_pid.as_deref());
195
196            return Err(format!(
197                "Failed to execute one-shot app-server turn: {error}"
198            ));
199        }
200    };
201
202    let parse_result = match parse_one_shot_response(&turn_result.assistant_message) {
203        Ok(response) => Ok((response, 0, 0)),
204        Err(parse_error) => {
205            attempt_one_shot_app_server_repair(
206                app_server_client,
207                &parse_error,
208                &turn_result.assistant_message,
209                request,
210                &session_id,
211                turn_result.provider_conversation_id.as_deref(),
212            )
213            .await
214        }
215    };
216
217    app_server_client.shutdown_session(session_id).await;
218    clear_child_pid_slot(child_pid.as_deref());
219
220    let (response, repair_input_tokens, repair_output_tokens) = parse_result?;
221
222    Ok(OneShotSubmission {
223        response,
224        stats: SessionStats {
225            added_lines: 0,
226            deleted_lines: 0,
227            diff_state: SessionDiffState::Unknown,
228            input_tokens: turn_result.input_tokens + repair_input_tokens,
229            output_tokens: turn_result.output_tokens + repair_output_tokens,
230        },
231    })
232}
233
234/// Executes one isolated prompt using the provided backend.
235///
236/// This shared helper keeps process execution behind the existing
237/// `AgentBackend` trait boundary so production callers and tests can reuse
238/// the same one-shot parsing path.
239///
240/// # Errors
241/// Returns an error when command construction fails, process execution fails,
242/// or the final output is empty or otherwise unusable.
243async fn submit_one_shot_with_backend(
244    backend: &dyn AgentBackend,
245    request: OneShotRequest,
246) -> Result<OneShotSubmission, String> {
247    let parsed_response =
248        execute_one_shot_command(backend, &request.prompt, request.clone()).await?;
249    let (agent_response, repair_stats) = match parse_one_shot_response(&parsed_response.content) {
250        Ok(response) => (response, None),
251        Err(parse_error) => {
252            let repair_prompt =
253                build_protocol_repair_prompt(&parse_error, &parsed_response.content);
254            let repair_response = execute_one_shot_command(backend, &repair_prompt, request)
255                .await
256                .map_err(|error| format!("{parse_error}\nrepair transport failed: {error}"))?;
257
258            let response = parse_one_shot_response(&repair_response.content).map_err(|error| {
259                format!(
260                    "{parse_error}\nrepair retry also failed: {error}\nrepair_response:\n{}",
261                    repair_response.content
262                )
263            })?;
264
265            (response, Some(repair_response.stats))
266        }
267    };
268
269    let mut stats = parsed_response.stats;
270    if let Some(repair) = repair_stats {
271        stats.input_tokens += repair.input_tokens;
272        stats.output_tokens += repair.output_tokens;
273    }
274
275    Ok(OneShotSubmission {
276        response: agent_response,
277        stats,
278    })
279}
280
281/// Parses one one-shot response strictly against the shared protocol schema.
282///
283/// # Errors
284/// Returns an error when the response is empty or not valid protocol JSON. The
285/// error carries the parse reason and derived diagnostics only, never the
286/// provider payload itself.
287fn parse_one_shot_response(content: &str) -> Result<AgentResponse, String> {
288    parse_agent_response_strict(content).map_err(|error| {
289        format!(
290            "One-shot agent output did not match the required JSON schema: \
291             {error}\ndebug_details:\n{}",
292            format_protocol_parse_debug_details(content)
293        )
294    })
295}
296
297/// Attempts one protocol-repair retry through the app-server transport for
298/// a one-shot prompt whose initial response failed schema validation.
299///
300/// The repair prompt is sent as a follow-up turn on the same session so the
301/// agent retains the original conversation context. The initial turn's
302/// `provider_conversation_id` is threaded through so providers that depend
303/// on conversation state can continue the same thread.
304///
305/// Returns the parsed response together with the repair turn's token usage
306/// so the caller can aggregate stats across both attempts.
307///
308/// # Errors
309/// Returns the combined original and repair error when the retry fails.
310async fn attempt_one_shot_app_server_repair(
311    app_server_client: &dyn AppServerClient,
312    parse_error: &str,
313    malformed_response: &str,
314    request: OneShotRequest,
315    session_id: &str,
316    provider_conversation_id: Option<&str>,
317) -> Result<(AgentResponse, u64, u64), String> {
318    let repair_prompt = build_protocol_repair_prompt(parse_error, malformed_response);
319
320    let (repair_stream_tx, _repair_stream_rx) = tokio::sync::mpsc::unbounded_channel();
321    let repair_turn_request = AppServerTurnRequest {
322        folder: request.folder,
323        live_transcript: None,
324        main_checkout_root: None,
325        model: request.model.provider_model_str().to_string(),
326        permission_mode: request.permission_mode,
327        personality: crate::channel::PersonalityPrompt::default(),
328        prompt: ag_protocol::TurnPrompt::from_agent_data(repair_prompt),
329        request_kind: request.request_kind,
330        replay_transcript: None,
331        provider_conversation_id: provider_conversation_id.map(String::from),
332        persisted_instruction_conversation_id: None,
333        reasoning_level: request.reasoning_level,
334        session_id: session_id.to_string(),
335        speed_mode: crate::model::session::SpeedMode::default(),
336    };
337    let repair_result = app_server_client
338        .run_turn(repair_turn_request, repair_stream_tx)
339        .await
340        .map_err(|error| format!("{parse_error}\nrepair transport failed: {error}"))?;
341
342    let response = parse_one_shot_response(&repair_result.assistant_message).map_err(|error| {
343        format!(
344            "{parse_error}\nrepair retry also failed: {error}\nrepair_response:\n{}",
345            repair_result.assistant_message
346        )
347    })?;
348
349    Ok((
350        response,
351        repair_result.input_tokens,
352        repair_result.output_tokens,
353    ))
354}
355
356/// Runs one one-shot backend command and returns the parsed provider content.
357///
358/// The spawned child is configured with `kill_on_drop(true)` so timeout-driven
359/// callers do not leave orphaned agent CLI processes behind when the future is
360/// canceled before completion.
361///
362/// # Errors
363/// Returns an error when the command cannot be built, run, or exits
364/// unsuccessfully.
365async fn execute_one_shot_command(
366    backend: &dyn AgentBackend,
367    prompt: &str,
368    request: OneShotRequest,
369) -> Result<ParsedResponse, String> {
370    let prompt_payload = ag_protocol::TurnPrompt::from_agent_data(prompt.to_string());
371    let build_request = BuildCommandRequest {
372        attachments: &prompt_payload.attachments,
373        folder: &request.folder,
374        main_checkout_root: None,
375        replay_transcript: None,
376        model: request.model.provider_model_str(),
377        permission_mode: request.permission_mode,
378        personality_prompt: None,
379        prompt,
380        reasoning_level: request.reasoning_level,
381        request_kind: &request.request_kind,
382        speed_mode: crate::model::session::SpeedMode::default(),
383    };
384    let observer = OneShotCliObserver {
385        child_pid: request.child_pid,
386    };
387    let output =
388        execution::execute_cli_command(backend, request.agent_kind, build_request, &observer, None)
389            .await
390            .map_err(format_one_shot_execution_error)?;
391
392    match output.exit_status {
393        CliExitStatus::Signaled(_) => {
394            return Err("One-shot agent command was interrupted".to_string());
395        }
396        CliExitStatus::NonZero(exit_code) => {
397            return Err(format_one_shot_exit_error(
398                request.agent_kind,
399                exit_code,
400                &output.stdout,
401                &output.stderr,
402            ));
403        }
404        CliExitStatus::Success => {}
405    }
406
407    let parsed_response = parse_response(request.agent_kind, &output.stdout, &output.stderr);
408
409    Ok(parsed_response)
410}
411
412/// Preserves the established one-shot context around shared execution errors.
413fn format_one_shot_execution_error(error: CliExecutionError) -> String {
414    match error {
415        CliExecutionError::CommandBuild(error) => {
416            format!("Failed to build one-shot agent command: {error}")
417        }
418        CliExecutionError::StdinBuild(error) => {
419            format!("Failed to build one-shot agent stdin payload: {error}")
420        }
421        error => format!("Failed to execute one-shot agent command: {error}"),
422    }
423}
424
425/// Formats one non-zero one-shot command exit into a user-facing error.
426fn format_one_shot_exit_error(
427    agent_kind: AgentKind,
428    exit_code: Option<i32>,
429    stdout: &str,
430    stderr: &str,
431) -> String {
432    error::format_agent_cli_exit_error(
433        agent_kind,
434        "One-shot agent command",
435        exit_code,
436        stdout,
437        stderr,
438    )
439}
440
441/// Clears the shared one-shot child PID slot when one exists.
442fn clear_child_pid_slot(child_pid: Option<&Mutex<Option<u32>>>) {
443    let Some(child_pid) = child_pid else {
444        return;
445    };
446
447    if let Ok(mut guard) = child_pid.lock() {
448        *guard = None;
449    }
450}
451
452/// Bridges shared CLI PID observations into the one-shot cancellation slot.
453struct OneShotCliObserver {
454    child_pid: Option<Arc<Mutex<Option<u32>>>>,
455}
456
457impl CliExecutionObserver for OneShotCliObserver {
458    fn pid_updated(&self, active_child_pid: Option<u32>) {
459        let Some(child_pid_slot) = self.child_pid.as_deref() else {
460            return;
461        };
462
463        if let Ok(mut guard) = child_pid_slot.lock() {
464            *guard = active_child_pid;
465        }
466    }
467
468    fn stdout_line(&self, _line: &str) {}
469}
470
471#[cfg(test)]
472mod tests {
473    use std::path::Path;
474    use std::process::Command;
475    use std::time::Duration;
476
477    use tempfile::tempdir;
478
479    use super::*;
480    use crate::MockAgentBackend;
481    use crate::app_server::{AppServerError, AppServerTurnResponse, MockAppServerClient};
482
483    /// Builds one shell command that emits controlled stdout/stderr and exits.
484    fn mock_shell_command(stdout: &str, stderr: &str, exit_code: i32) -> Command {
485        let mut command = Command::new("sh");
486        command.arg("-c").arg(
487            "printf '%s' \"$ONE_SHOT_STDOUT\"; printf '%s' \"$ONE_SHOT_STDERR\" >&2; exit \
488             \"$ONE_SHOT_EXIT\"",
489        );
490        command.env("ONE_SHOT_STDOUT", stdout);
491        command.env("ONE_SHOT_STDERR", stderr);
492        command.env("ONE_SHOT_EXIT", exit_code.to_string());
493        command.stdout(std::process::Stdio::piped());
494        command.stderr(std::process::Stdio::piped());
495
496        command
497    }
498
499    /// Builds one shell command that captures stdin before returning JSON.
500    fn stdin_capture_shell_command(capture_path: &Path) -> Command {
501        let mut command = Command::new("sh");
502        command.arg("-c").arg(
503            "cat > \"$ONE_SHOT_CAPTURE_PATH\"; printf '%s' \
504             '{\"answer\":\"captured\",\"questions\":[],\"summary\":null}'",
505        );
506        command.env("ONE_SHOT_CAPTURE_PATH", capture_path);
507        command.stdout(std::process::Stdio::piped());
508        command.stderr(std::process::Stdio::piped());
509
510        command
511    }
512
513    #[test]
514    fn test_format_one_shot_execution_error_preserves_build_context() {
515        // Arrange
516        let command_error = CliExecutionError::CommandBuild(
517            crate::agent::AgentBackendError::CommandBuild("command".to_string()),
518        );
519        let stdin_error = CliExecutionError::StdinBuild(
520            crate::agent::AgentBackendError::CommandBuild("stdin".to_string()),
521        );
522        let execution_error = CliExecutionError::StdinWrite("write".to_string());
523
524        // Act
525        let command_message = format_one_shot_execution_error(command_error);
526        let stdin_message = format_one_shot_execution_error(stdin_error);
527        let execution_message = format_one_shot_execution_error(execution_error);
528
529        // Assert
530        assert_eq!(
531            command_message,
532            "Failed to build one-shot agent command: command"
533        );
534        assert_eq!(
535            stdin_message,
536            "Failed to build one-shot agent stdin payload: stdin"
537        );
538        assert_eq!(
539            execution_message,
540            "Failed to execute one-shot agent command: stdin delivery failed: write"
541        );
542    }
543
544    #[test]
545    fn test_one_shot_cli_observer_updates_child_pid_slot() {
546        // Arrange
547        let child_pid = Arc::new(Mutex::new(None));
548        let observer = OneShotCliObserver {
549            child_pid: Some(Arc::clone(&child_pid)),
550        };
551
552        // Act
553        observer.pid_updated(Some(42));
554        let active_pid = *child_pid.lock().expect("PID lock should be available");
555        observer.stdout_line("collected output");
556        observer.pid_updated(None);
557        let cleared_pid = *child_pid.lock().expect("PID lock should be available");
558
559        // Assert
560        assert_eq!(active_pid, Some(42));
561        assert_eq!(cleared_pid, None);
562    }
563
564    #[tokio::test]
565    async fn test_submit_one_shot_with_backend_reports_signal_interruption() {
566        // Arrange
567        let temp_directory = tempdir().expect("failed to create temp dir");
568        let mut backend = MockAgentBackend::new();
569        backend.expect_build_command().returning(|_| {
570            let mut command = Command::new("sh");
571            command.arg("-c").arg("kill -9 $$");
572
573            Ok(command)
574        });
575
576        // Act
577        let error = submit_one_shot_with_backend(
578            &backend,
579            OneShotRequest {
580                agent_kind: AgentKind::Codex,
581                child_pid: None,
582                folder: temp_directory.path().to_path_buf(),
583                model: AgentModel::Gpt56Sol,
584                permission_mode: PermissionMode::AutoEdit,
585                prompt: "Generate title".to_string(),
586                request_kind: AgentRequestKind::UtilityPrompt,
587                reasoning_level: ReasoningLevel::default(),
588            },
589        )
590        .await
591        .expect_err("signal termination should interrupt the one-shot command");
592
593        // Assert
594        assert_eq!(error, "One-shot agent command was interrupted");
595    }
596
597    #[tokio::test]
598    /// Verifies one-shot execution returns the parsed structured answer.
599    async fn test_submit_one_shot_with_backend_returns_protocol_response() {
600        // Arrange
601        let temp_directory = tempdir().expect("failed to create temp dir");
602        let mut backend = MockAgentBackend::new();
603        backend.expect_build_command().returning(|request| {
604            assert!(matches!(
605                request.request_kind,
606                AgentRequestKind::UtilityPrompt
607            ));
608            assert_eq!(request.permission_mode, PermissionMode::ReadOnly);
609            assert_eq!(request.prompt, "Generate title");
610
611            Ok(mock_shell_command(
612                r#"{"answer":"Generated title","questions":[],"summary":null}"#,
613                "",
614                0,
615            ))
616        });
617
618        // Act
619        let response = submit_one_shot_with_backend(
620            &backend,
621            OneShotRequest {
622                agent_kind: AgentKind::Claude,
623                child_pid: None,
624                folder: temp_directory.path().to_path_buf(),
625                model: AgentModel::ClaudeSonnet5,
626                permission_mode: PermissionMode::ReadOnly,
627                prompt: "Generate title".to_string(),
628                request_kind: AgentRequestKind::UtilityPrompt,
629                reasoning_level: ReasoningLevel::default(),
630            },
631        )
632        .await
633        .expect("one-shot prompt should succeed");
634
635        // Assert
636        assert_eq!(
637            response.response.answers(),
638            vec!["Generated title".to_string()]
639        );
640    }
641
642    #[tokio::test]
643    /// Verifies one-shot execution rejects plain-text utility output after
644    /// both the original parse and the protocol-repair retry fail.
645    async fn test_submit_one_shot_with_backend_rejects_plain_text_utility_output() {
646        // Arrange
647        let temp_directory = tempdir().expect("failed to create temp dir");
648        let mut backend = MockAgentBackend::new();
649        backend
650            .expect_build_command()
651            .times(2)
652            .returning(|request| {
653                assert!(matches!(
654                    request.request_kind,
655                    AgentRequestKind::UtilityPrompt
656                ));
657
658                Ok(mock_shell_command("plain text", "", 0))
659            });
660
661        // Act
662        let error = submit_one_shot_with_backend(
663            &backend,
664            OneShotRequest {
665                agent_kind: AgentKind::Codex,
666                child_pid: None,
667                folder: temp_directory.path().to_path_buf(),
668                model: AgentModel::Gpt56Sol,
669                permission_mode: PermissionMode::AutoEdit,
670                prompt: "Generate title".to_string(),
671                request_kind: AgentRequestKind::UtilityPrompt,
672                reasoning_level: ReasoningLevel::default(),
673            },
674        )
675        .await
676        .expect_err("plain-text utility output should fail");
677
678        // Assert
679        assert!(error.contains("did not match the required JSON schema"));
680        assert!(error.contains("debug_details:"));
681        assert!(error.contains("direct_json_error_location: line 1, column 1"));
682        assert!(error.contains("response:\nplain text"));
683    }
684
685    #[tokio::test]
686    /// Verifies one-shot execution rejects wrapped non-schema utility output
687    /// after both the original parse and the protocol-repair retry fail.
688    async fn test_submit_one_shot_with_backend_rejects_wrapped_plain_text_utility_output() {
689        // Arrange
690        let temp_directory = tempdir().expect("failed to create temp dir");
691        let mut backend = MockAgentBackend::new();
692        backend
693            .expect_build_command()
694            .times(2)
695            .returning(|request| {
696                assert!(matches!(
697                    request.request_kind,
698                    AgentRequestKind::UtilityPrompt
699                ));
700
701                Ok(mock_shell_command(
702                    r#"{"result":"plain text","usage":{"input_tokens":2,"output_tokens":1}}"#,
703                    "",
704                    0,
705                ))
706            });
707
708        // Act
709        let error = submit_one_shot_with_backend(
710            &backend,
711            OneShotRequest {
712                agent_kind: AgentKind::Claude,
713                child_pid: None,
714                folder: temp_directory.path().to_path_buf(),
715                model: AgentModel::ClaudeSonnet5,
716                permission_mode: PermissionMode::AutoEdit,
717                prompt: "Generate title".to_string(),
718                request_kind: AgentRequestKind::UtilityPrompt,
719                reasoning_level: ReasoningLevel::default(),
720            },
721        )
722        .await
723        .expect_err("wrapped plain-text utility output should fail");
724
725        // Assert — the provider parser extracts "plain text" from the
726        // `result` wrapper, so the protocol parser sees raw text, not JSON keys.
727        assert!(error.contains("did not match the required JSON schema"));
728        assert!(error.contains("direct_json_error:"));
729        assert!(error.contains("response:\nplain text"));
730    }
731
732    #[tokio::test]
733    /// Verifies one-shot execution recovers a trailing protocol payload when
734    /// the provider prepends extra prose before the final JSON object.
735    async fn test_submit_one_shot_with_backend_recovers_wrapped_protocol_output() {
736        // Arrange
737        let temp_directory = tempdir().expect("failed to create temp dir");
738        let mut backend = MockAgentBackend::new();
739        backend
740            .expect_build_command()
741            .times(1)
742            .returning(|request| {
743                assert!(matches!(
744                    request.request_kind,
745                    AgentRequestKind::UtilityPrompt
746                ));
747                assert_eq!(request.prompt, "Generate title");
748
749                Ok(mock_shell_command(
750                    concat!(
751                        "Now I have full context.\n",
752                        r#"{"answer":"Generated title","questions":[],"summary":null}"#
753                    ),
754                    "",
755                    0,
756                ))
757            });
758
759        // Act
760        let response = submit_one_shot_with_backend(
761            &backend,
762            OneShotRequest {
763                agent_kind: AgentKind::Claude,
764                child_pid: None,
765                folder: temp_directory.path().to_path_buf(),
766                model: AgentModel::ClaudeSonnet5,
767                permission_mode: PermissionMode::AutoEdit,
768                prompt: "Generate title".to_string(),
769                request_kind: AgentRequestKind::UtilityPrompt,
770                reasoning_level: ReasoningLevel::default(),
771            },
772        )
773        .await
774        .expect("wrapped protocol output should succeed");
775
776        // Assert
777        assert_eq!(
778            response.response.answers(),
779            vec!["Generated title".to_string()]
780        );
781    }
782
783    #[tokio::test]
784    /// Verifies one-shot execution recovers valid output when the initial
785    /// parse fails but the protocol-repair retry returns valid protocol JSON.
786    async fn test_submit_one_shot_with_backend_recovers_via_protocol_repair() {
787        // Arrange
788        let temp_directory = tempdir().expect("failed to create temp dir");
789        let call_counter = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
790        let mut backend = MockAgentBackend::new();
791        backend.expect_build_command().times(2).returning({
792            let counter = std::sync::Arc::clone(&call_counter);
793
794            move |_| {
795                let call_number = counter.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
796
797                if call_number == 0 {
798                    Ok(mock_shell_command("plain text", "", 0))
799                } else {
800                    Ok(mock_shell_command(
801                        r#"{"answer":"Repaired title","questions":[],"summary":null}"#,
802                        "",
803                        0,
804                    ))
805                }
806            }
807        });
808
809        // Act
810        let response = submit_one_shot_with_backend(
811            &backend,
812            OneShotRequest {
813                agent_kind: AgentKind::Codex,
814                child_pid: None,
815                folder: temp_directory.path().to_path_buf(),
816                model: AgentModel::Gpt56Sol,
817                permission_mode: PermissionMode::AutoEdit,
818                prompt: "Generate title".to_string(),
819                request_kind: AgentRequestKind::UtilityPrompt,
820                reasoning_level: ReasoningLevel::default(),
821            },
822        )
823        .await
824        .expect("repair retry should succeed");
825
826        // Assert
827        assert_eq!(
828            response.response.answers(),
829            vec!["Repaired title".to_string()]
830        );
831    }
832
833    #[tokio::test]
834    /// Verifies one-shot execution still rejects blank utility responses
835    /// after both the original parse and the protocol-repair retry fail.
836    async fn test_submit_one_shot_with_backend_rejects_blank_utility_output() {
837        // Arrange
838        let temp_directory = tempdir().expect("failed to create temp dir");
839        let mut backend = MockAgentBackend::new();
840        backend.expect_build_command().returning(|request| {
841            assert!(matches!(
842                request.request_kind,
843                AgentRequestKind::UtilityPrompt
844            ));
845
846            Ok(mock_shell_command("   ", "", 0))
847        });
848
849        // Act
850        let error = submit_one_shot_with_backend(
851            &backend,
852            OneShotRequest {
853                agent_kind: AgentKind::Codex,
854                child_pid: None,
855                folder: temp_directory.path().to_path_buf(),
856                model: AgentModel::Gpt56Sol,
857                permission_mode: PermissionMode::AutoEdit,
858                prompt: "Generate title".to_string(),
859                request_kind: AgentRequestKind::UtilityPrompt,
860                reasoning_level: ReasoningLevel::default(),
861            },
862        )
863        .await
864        .expect_err("blank utility output should fail");
865
866        // Assert
867        assert!(error.contains("did not match the required JSON schema"));
868        assert!(error.contains("trimmed_len: 0 chars"));
869        assert!(error.contains("response:\n"));
870    }
871
872    #[tokio::test]
873    /// Verifies one-shot execution does not deadlock when the child delays
874    /// reading stdin until after it emits early stderr output.
875    async fn test_submit_one_shot_with_backend_writes_large_stdin_concurrently() {
876        // Arrange
877        let temp_directory = tempdir().expect("failed to create temp dir");
878        let large_prompt = "x".repeat(512 * 1024);
879        let mut backend = MockAgentBackend::new();
880        backend.expect_build_command().returning(|_| {
881            let mut command = Command::new("sh");
882            command.arg("-c").arg(
883                "printf 'warming up\\n' >&2; sleep 0.1; cat >/dev/null; printf '%s' \
884                 '{\"answer\":\"done\",\"questions\":[],\"summary\":null}'",
885            );
886            command.stdout(std::process::Stdio::piped());
887            command.stderr(std::process::Stdio::piped());
888
889            Ok(command)
890        });
891
892        // Act
893        let response = tokio::time::timeout(
894            Duration::from_secs(5),
895            submit_one_shot_with_backend(
896                &backend,
897                OneShotRequest {
898                    agent_kind: AgentKind::Claude,
899                    child_pid: None,
900                    folder: temp_directory.path().to_path_buf(),
901                    model: AgentModel::ClaudeSonnet5,
902                    permission_mode: PermissionMode::AutoEdit,
903                    prompt: large_prompt.clone(),
904                    request_kind: AgentRequestKind::UtilityPrompt,
905                    reasoning_level: ReasoningLevel::default(),
906                },
907            ),
908        )
909        .await
910        .expect("one-shot prompt should not deadlock")
911        .expect("one-shot prompt should succeed");
912
913        // Assert
914        assert_eq!(response.response.answers(), vec!["done".to_string()]);
915    }
916
917    #[tokio::test]
918    /// Verifies one-shot execution streams Claude prompts through stdin so
919    /// large review requests avoid argv length limits.
920    async fn test_submit_one_shot_with_backend_writes_prompt_to_stdin() {
921        // Arrange
922        let temp_directory = tempdir().expect("failed to create temp dir");
923        let capture_path = temp_directory.path().join("stdin.txt");
924        let mut backend = MockAgentBackend::new();
925        backend.expect_build_command().returning({
926            let capture_path = capture_path.clone();
927
928            move |_| Ok(stdin_capture_shell_command(&capture_path))
929        });
930
931        // Act
932        let response = submit_one_shot_with_backend(
933            &backend,
934            OneShotRequest {
935                agent_kind: AgentKind::Claude,
936                child_pid: None,
937                folder: temp_directory.path().to_path_buf(),
938                model: AgentModel::ClaudeSonnet5,
939                permission_mode: PermissionMode::AutoEdit,
940                prompt: "Generate title".to_string(),
941                request_kind: AgentRequestKind::UtilityPrompt,
942                reasoning_level: ReasoningLevel::default(),
943            },
944        )
945        .await
946        .expect("one-shot prompt should succeed");
947        let captured_prompt =
948            std::fs::read_to_string(&capture_path).expect("captured stdin payload should exist");
949
950        // Assert
951        assert_eq!(response.response.answers(), vec!["captured".to_string()]);
952        assert!(captured_prompt.contains("Structured response protocol:"));
953        assert!(captured_prompt.contains("Generate title"));
954    }
955
956    #[tokio::test]
957    /// Verifies a broken stdin pipe does not hide the child exit status or
958    /// stderr when the backend exits before reading the full prompt.
959    async fn test_submit_one_shot_with_backend_preserves_exit_error_after_broken_pipe() {
960        // Arrange
961        let temp_directory = tempdir().expect("failed to create temp dir");
962        let large_prompt = "x".repeat(512 * 1024);
963        let mut backend = MockAgentBackend::new();
964        backend.expect_build_command().returning(|_| {
965            let mut command = Command::new("sh");
966            command.arg("-c").arg("printf 'auth failed' >&2; exit 7");
967            command.stdout(std::process::Stdio::piped());
968            command.stderr(std::process::Stdio::piped());
969
970            Ok(command)
971        });
972
973        // Act
974        let error = submit_one_shot_with_backend(
975            &backend,
976            OneShotRequest {
977                agent_kind: AgentKind::Claude,
978                child_pid: None,
979                folder: temp_directory.path().to_path_buf(),
980                model: AgentModel::ClaudeSonnet5,
981                permission_mode: PermissionMode::AutoEdit,
982                prompt: large_prompt,
983                request_kind: AgentRequestKind::UtilityPrompt,
984                reasoning_level: ReasoningLevel::default(),
985            },
986        )
987        .await
988        .expect_err("one-shot prompt should surface the child exit");
989
990        // Assert
991        assert!(error.contains("exit code 7"), "error was: {error}");
992        assert!(error.contains("auth failed"), "error was: {error}");
993        assert!(
994            !error.contains("stdin payload"),
995            "stdin write error should not mask child failure: {error}"
996        );
997    }
998
999    #[tokio::test]
1000    /// Verifies Claude authentication failures return actionable re-login
1001    /// guidance instead of raw transport output.
1002    async fn test_submit_one_shot_with_backend_surfaces_claude_auth_guidance() {
1003        // Arrange
1004        let temp_directory = tempdir().expect("failed to create temp dir");
1005        let mut backend = MockAgentBackend::new();
1006        backend.expect_build_command().returning(|_| {
1007            Ok(mock_shell_command(
1008                r#"{"type":"error","error":{"type":"authentication_error","message":"OAuth token has expired. Please obtain a new token or refresh your existing token."}}"#,
1009                "",
1010                1,
1011            ))
1012        });
1013
1014        // Act
1015        let error = submit_one_shot_with_backend(
1016            &backend,
1017            OneShotRequest {
1018                agent_kind: AgentKind::Claude,
1019                child_pid: None,
1020                folder: temp_directory.path().to_path_buf(),
1021                model: AgentModel::ClaudeSonnet5,
1022                permission_mode: PermissionMode::AutoEdit,
1023                prompt: "Generate title".to_string(),
1024                request_kind: AgentRequestKind::UtilityPrompt,
1025                reasoning_level: ReasoningLevel::default(),
1026            },
1027        )
1028        .await
1029        .expect_err("expired Claude auth should fail");
1030
1031        // Assert
1032        assert!(
1033            error.contains("One-shot agent command failed because Claude authentication expired")
1034        );
1035        assert!(error.contains("`claude auth login`"));
1036        assert!(error.contains("`claude auth status`"));
1037    }
1038
1039    #[tokio::test]
1040    /// Verifies app-server-backed one-shot execution returns the parsed
1041    /// structured answer and usage totals.
1042    async fn test_submit_one_shot_with_app_server_client_returns_protocol_response() {
1043        // Arrange
1044        let temp_directory = tempdir().expect("failed to create temp dir");
1045        let mut app_server_client = MockAppServerClient::new();
1046        app_server_client
1047            .expect_run_turn()
1048            .times(1)
1049            .returning(|request, _| {
1050                assert_eq!(request.model, AgentModel::Gpt56Sol.as_str());
1051                assert!(matches!(
1052                    request.request_kind,
1053                    AgentRequestKind::UtilityPrompt
1054                ));
1055                assert_eq!(request.permission_mode, PermissionMode::ReadOnly);
1056                assert_eq!(request.prompt.text, "Generate title");
1057
1058                Box::pin(async {
1059                    Ok(AppServerTurnResponse {
1060                        assistant_message:
1061                            r#"{"answer":"Generated title","questions":[],"summary":null}"#
1062                                .to_string(),
1063                        context_reset: false,
1064                        input_tokens: 11,
1065                        output_tokens: 7,
1066                        pid: Some(42),
1067                        provider_conversation_id: Some("thread-1".to_string()),
1068                    })
1069                })
1070            });
1071        app_server_client
1072            .expect_shutdown_session()
1073            .times(1)
1074            .returning(|_| Box::pin(async {}));
1075
1076        // Act
1077        let response = submit_one_shot_with_app_server_client(
1078            &app_server_client,
1079            OneShotRequest {
1080                agent_kind: AgentKind::Codex,
1081                child_pid: None,
1082                folder: temp_directory.path().to_path_buf(),
1083                model: AgentModel::Gpt56Sol,
1084                permission_mode: PermissionMode::ReadOnly,
1085                prompt: "Generate title".to_string(),
1086                request_kind: AgentRequestKind::UtilityPrompt,
1087                reasoning_level: ReasoningLevel::default(),
1088            },
1089        )
1090        .await
1091        .expect("one-shot prompt should succeed");
1092
1093        // Assert
1094        assert_eq!(
1095            response.response.answers(),
1096            vec!["Generated title".to_string()]
1097        );
1098        assert_eq!(response.stats.input_tokens, 11);
1099        assert_eq!(response.stats.output_tokens, 7);
1100    }
1101
1102    #[tokio::test]
1103    /// Verifies app-server turn failures shut down the temporary session and
1104    /// clear the caller's shared child-process slot.
1105    async fn test_submit_one_shot_with_app_server_client_clears_pid_after_turn_failure() {
1106        // Arrange
1107        let temp_directory = tempdir().expect("failed to create temp dir");
1108        let child_pid = Arc::new(Mutex::new(Some(42)));
1109        let mut app_server_client = MockAppServerClient::new();
1110        app_server_client
1111            .expect_run_turn()
1112            .times(1)
1113            .returning(|_, _| {
1114                Box::pin(async {
1115                    Err(AppServerError::Provider(
1116                        "app-server turn failed".to_string(),
1117                    ))
1118                })
1119            });
1120        app_server_client
1121            .expect_shutdown_session()
1122            .times(1)
1123            .returning(|_| Box::pin(async {}));
1124
1125        // Act
1126        let error = submit_one_shot_with_app_server_client(
1127            &app_server_client,
1128            OneShotRequest {
1129                agent_kind: AgentKind::Codex,
1130                child_pid: Some(Arc::clone(&child_pid)),
1131                folder: temp_directory.path().to_path_buf(),
1132                model: AgentModel::Gpt56Sol,
1133                permission_mode: PermissionMode::AutoEdit,
1134                prompt: "Generate title".to_string(),
1135                request_kind: AgentRequestKind::UtilityPrompt,
1136                reasoning_level: ReasoningLevel::default(),
1137            },
1138        )
1139        .await
1140        .expect_err("app-server turn failure should surface");
1141
1142        // Assert
1143        assert!(error.contains("app-server turn failed"));
1144        assert_eq!(
1145            *child_pid.lock().expect("child pid lock should succeed"),
1146            None
1147        );
1148    }
1149
1150    #[tokio::test]
1151    /// Verifies app-server-backed one-shot execution rejects plain-text
1152    /// utility output after both the original parse and the protocol-repair
1153    /// retry fail.
1154    async fn test_submit_one_shot_with_app_server_client_rejects_plain_text_utility_output() {
1155        // Arrange
1156        let temp_directory = tempdir().expect("failed to create temp dir");
1157        let mut app_server_client = MockAppServerClient::new();
1158        app_server_client
1159            .expect_run_turn()
1160            .times(2)
1161            .returning(|request, _| {
1162                assert_eq!(request.model, AgentModel::Gpt56Sol.as_str());
1163                assert_eq!(request.permission_mode, PermissionMode::ReadOnly);
1164
1165                Box::pin(async {
1166                    Ok(AppServerTurnResponse {
1167                        assistant_message: "plain text".to_string(),
1168                        context_reset: false,
1169                        input_tokens: 2,
1170                        output_tokens: 1,
1171                        pid: None,
1172                        provider_conversation_id: None,
1173                    })
1174                })
1175            });
1176        app_server_client
1177            .expect_shutdown_session()
1178            .times(1)
1179            .returning(|_| Box::pin(async {}));
1180
1181        // Act
1182        let error = submit_one_shot_with_app_server_client(
1183            &app_server_client,
1184            OneShotRequest {
1185                agent_kind: AgentKind::Codex,
1186                child_pid: None,
1187                folder: temp_directory.path().to_path_buf(),
1188                model: AgentModel::Gpt56Sol,
1189                permission_mode: PermissionMode::ReadOnly,
1190                prompt: "Generate title".to_string(),
1191                request_kind: AgentRequestKind::UtilityPrompt,
1192                reasoning_level: ReasoningLevel::default(),
1193            },
1194        )
1195        .await
1196        .expect_err("plain-text utility output should fail");
1197
1198        // Assert
1199        assert!(error.contains("did not match the required JSON schema"));
1200        assert!(error.contains("debug_details:"));
1201        assert!(error.contains("response:\nplain text"));
1202    }
1203
1204    #[tokio::test]
1205    /// Verifies app-server-backed non-utility one-shot execution still
1206    /// rejects plain-text output after both the original parse and the
1207    /// protocol-repair retry fail.
1208    async fn test_submit_one_shot_with_app_server_client_rejects_plain_text_non_utility_output() {
1209        // Arrange
1210        let temp_directory = tempdir().expect("failed to create temp dir");
1211        let mut app_server_client = MockAppServerClient::new();
1212        app_server_client
1213            .expect_run_turn()
1214            .times(2)
1215            .returning(|request, _| {
1216                assert!(matches!(
1217                    request.request_kind,
1218                    AgentRequestKind::SessionStart
1219                ));
1220
1221                Box::pin(async {
1222                    Ok(AppServerTurnResponse {
1223                        assistant_message: "plain text".to_string(),
1224                        context_reset: false,
1225                        input_tokens: 2,
1226                        output_tokens: 1,
1227                        pid: None,
1228                        provider_conversation_id: None,
1229                    })
1230                })
1231            });
1232        app_server_client
1233            .expect_shutdown_session()
1234            .times(1)
1235            .returning(|_| Box::pin(async {}));
1236
1237        // Act
1238        let error = submit_one_shot_with_app_server_client(
1239            &app_server_client,
1240            OneShotRequest {
1241                agent_kind: AgentKind::Codex,
1242                child_pid: None,
1243                folder: temp_directory.path().to_path_buf(),
1244                model: AgentModel::Gpt56Sol,
1245                permission_mode: PermissionMode::AutoEdit,
1246                prompt: "Generate title".to_string(),
1247                request_kind: AgentRequestKind::SessionStart,
1248                reasoning_level: ReasoningLevel::default(),
1249            },
1250        )
1251        .await
1252        .expect_err("invalid non-utility output should fail");
1253
1254        // Assert
1255        assert!(error.contains("did not match the required JSON schema"));
1256        assert!(error.contains("debug_details:"));
1257        assert!(error.contains("response:\nplain text"));
1258    }
1259}