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