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