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