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::os::unix::process::ExitStatusExt as _;
9use std::path::PathBuf;
10use std::sync::{Arc, Mutex};
11
12use ag_protocol::{
13    AgentResponse, build_protocol_repair_prompt, format_protocol_parse_debug_details,
14    parse_agent_response_strict,
15};
16use async_trait::async_trait;
17
18use super::backend::{AgentBackend, BuildCommandRequest};
19use super::cli::{error, stdin};
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 command = backend
372        .build_command(build_request)
373        .map_err(|error| format!("Failed to build one-shot agent command: {error}"))?;
374    let stdin_payload = super::build_command_stdin_payload(request.agent_kind, build_request)
375        .map_err(|error| format!("Failed to build one-shot agent stdin payload: {error}"))?;
376    let mut tokio_command = tokio::process::Command::from(command);
377    tokio_command
378        .stdin(if stdin_payload.is_some() {
379            std::process::Stdio::piped()
380        } else {
381            std::process::Stdio::null()
382        })
383        .kill_on_drop(true);
384    let mut pid_guard = ChildPidGuard::new(request.child_pid);
385    let mut child = tokio_command
386        .spawn()
387        .map_err(|error| format!("Failed to execute one-shot agent command: {error}"))?;
388    pid_guard.update_from_child(&child);
389    let stdin_write_task = stdin::spawn_optional_stdin_write(
390        child.stdin.take(),
391        stdin_payload,
392        "one-shot stdin pipe unavailable after spawn",
393        std::convert::identity,
394    );
395    let output = child
396        .wait_with_output()
397        .await
398        .map_err(|error| format!("Failed to execute one-shot agent command: {error}"))?;
399    stdin::await_optional_stdin_write(
400        stdin_write_task,
401        "One-shot stdin write task failed",
402        std::convert::identity,
403    )
404    .await?;
405
406    if output.status.signal().is_some() {
407        return Err("One-shot agent command was interrupted".to_string());
408    }
409
410    let stdout_text = String::from_utf8_lossy(&output.stdout).into_owned();
411    let stderr_text = String::from_utf8_lossy(&output.stderr).into_owned();
412    if !output.status.success() {
413        return Err(format_one_shot_exit_error(
414            request.agent_kind,
415            output.status.code(),
416            &stdout_text,
417            &stderr_text,
418        ));
419    }
420
421    let parsed_response = parse_response(request.agent_kind, &stdout_text, &stderr_text);
422
423    Ok(parsed_response)
424}
425
426/// Formats one non-zero one-shot command exit into a user-facing error.
427fn format_one_shot_exit_error(
428    agent_kind: AgentKind,
429    exit_code: Option<i32>,
430    stdout: &str,
431    stderr: &str,
432) -> String {
433    error::format_agent_cli_exit_error(
434        agent_kind,
435        "One-shot agent command",
436        exit_code,
437        stdout,
438        stderr,
439    )
440}
441
442/// Clears the shared one-shot child PID slot when one exists.
443fn clear_child_pid_slot(child_pid: Option<&Mutex<Option<u32>>>) {
444    let Some(child_pid) = child_pid else {
445        return;
446    };
447
448    if let Ok(mut guard) = child_pid.lock() {
449        *guard = None;
450    }
451}
452
453/// Tracks the active one-shot subprocess identifier for cancel/stop flows.
454struct ChildPidGuard {
455    child_pid: Option<Arc<Mutex<Option<u32>>>>,
456}
457
458impl ChildPidGuard {
459    /// Creates one PID guard for the optional shared child slot.
460    fn new(child_pid: Option<Arc<Mutex<Option<u32>>>>) -> Self {
461        Self { child_pid }
462    }
463
464    /// Copies the spawned child PID into the shared slot when available.
465    fn update_from_child(&mut self, child: &tokio::process::Child) {
466        let Some(pid) = child.id() else {
467            return;
468        };
469
470        let Some(child_pid) = self.child_pid.as_deref() else {
471            return;
472        };
473
474        if let Ok(mut guard) = child_pid.lock() {
475            *guard = Some(pid);
476        }
477    }
478}
479
480impl Drop for ChildPidGuard {
481    fn drop(&mut self) {
482        let Some(child_pid) = self.child_pid.as_deref() else {
483            return;
484        };
485
486        if let Ok(mut guard) = child_pid.lock() {
487            *guard = None;
488        }
489    }
490}
491
492#[cfg(test)]
493mod tests {
494    use std::path::Path;
495    use std::process::Command;
496    use std::time::Duration;
497
498    use tempfile::tempdir;
499
500    use super::*;
501    use crate::MockAgentBackend;
502    use crate::app_server::{AppServerError, AppServerTurnResponse, MockAppServerClient};
503
504    /// Builds one shell command that emits controlled stdout/stderr and exits.
505    fn mock_shell_command(stdout: &str, stderr: &str, exit_code: i32) -> Command {
506        let mut command = Command::new("sh");
507        command.arg("-c").arg(
508            "printf '%s' \"$ONE_SHOT_STDOUT\"; printf '%s' \"$ONE_SHOT_STDERR\" >&2; exit \
509             \"$ONE_SHOT_EXIT\"",
510        );
511        command.env("ONE_SHOT_STDOUT", stdout);
512        command.env("ONE_SHOT_STDERR", stderr);
513        command.env("ONE_SHOT_EXIT", exit_code.to_string());
514        command.stdout(std::process::Stdio::piped());
515        command.stderr(std::process::Stdio::piped());
516
517        command
518    }
519
520    /// Builds one shell command that captures stdin before returning JSON.
521    fn stdin_capture_shell_command(capture_path: &Path) -> Command {
522        let mut command = Command::new("sh");
523        command.arg("-c").arg(
524            "cat > \"$ONE_SHOT_CAPTURE_PATH\"; printf '%s' \
525             '{\"answer\":\"captured\",\"questions\":[],\"summary\":null}'",
526        );
527        command.env("ONE_SHOT_CAPTURE_PATH", capture_path);
528        command.stdout(std::process::Stdio::piped());
529        command.stderr(std::process::Stdio::piped());
530
531        command
532    }
533
534    #[tokio::test]
535    /// Verifies one-shot execution returns the parsed structured answer.
536    async fn test_submit_one_shot_with_backend_returns_protocol_response() {
537        // Arrange
538        let temp_directory = tempdir().expect("failed to create temp dir");
539        let mut backend = MockAgentBackend::new();
540        backend.expect_build_command().returning(|request| {
541            assert!(matches!(
542                request.request_kind,
543                AgentRequestKind::UtilityPrompt
544            ));
545            assert_eq!(request.prompt, "Generate title");
546
547            Ok(mock_shell_command(
548                r#"{"answer":"Generated title","questions":[],"summary":null}"#,
549                "",
550                0,
551            ))
552        });
553
554        // Act
555        let response = submit_one_shot_with_backend(
556            &backend,
557            OneShotRequest {
558                agent_kind: AgentKind::Claude,
559                child_pid: None,
560                folder: temp_directory.path().to_path_buf(),
561                model: AgentModel::ClaudeSonnet5,
562                prompt: "Generate title".to_string(),
563                request_kind: AgentRequestKind::UtilityPrompt,
564                reasoning_level: ReasoningLevel::default(),
565            },
566        )
567        .await
568        .expect("one-shot prompt should succeed");
569
570        // Assert
571        assert_eq!(
572            response.response.answers(),
573            vec!["Generated title".to_string()]
574        );
575    }
576
577    #[tokio::test]
578    /// Verifies one-shot execution rejects plain-text utility output after
579    /// both the original parse and the protocol-repair retry fail.
580    async fn test_submit_one_shot_with_backend_rejects_plain_text_utility_output() {
581        // Arrange
582        let temp_directory = tempdir().expect("failed to create temp dir");
583        let mut backend = MockAgentBackend::new();
584        backend
585            .expect_build_command()
586            .times(2)
587            .returning(|request| {
588                assert!(matches!(
589                    request.request_kind,
590                    AgentRequestKind::UtilityPrompt
591                ));
592
593                Ok(mock_shell_command("plain text", "", 0))
594            });
595
596        // Act
597        let error = submit_one_shot_with_backend(
598            &backend,
599            OneShotRequest {
600                agent_kind: AgentKind::Codex,
601                child_pid: None,
602                folder: temp_directory.path().to_path_buf(),
603                model: AgentModel::Gpt55,
604                prompt: "Generate title".to_string(),
605                request_kind: AgentRequestKind::UtilityPrompt,
606                reasoning_level: ReasoningLevel::default(),
607            },
608        )
609        .await
610        .expect_err("plain-text utility output should fail");
611
612        // Assert
613        assert!(error.contains("did not match the required JSON schema"));
614        assert!(error.contains("debug_details:"));
615        assert!(error.contains("direct_json_error_location: line 1, column 1"));
616        assert!(error.contains("response:\nplain text"));
617    }
618
619    #[tokio::test]
620    /// Verifies one-shot execution rejects wrapped non-schema utility output
621    /// after both the original parse and the protocol-repair retry fail.
622    async fn test_submit_one_shot_with_backend_rejects_wrapped_plain_text_utility_output() {
623        // Arrange
624        let temp_directory = tempdir().expect("failed to create temp dir");
625        let mut backend = MockAgentBackend::new();
626        backend
627            .expect_build_command()
628            .times(2)
629            .returning(|request| {
630                assert!(matches!(
631                    request.request_kind,
632                    AgentRequestKind::UtilityPrompt
633                ));
634
635                Ok(mock_shell_command(
636                    r#"{"result":"plain text","usage":{"input_tokens":2,"output_tokens":1}}"#,
637                    "",
638                    0,
639                ))
640            });
641
642        // Act
643        let error = submit_one_shot_with_backend(
644            &backend,
645            OneShotRequest {
646                agent_kind: AgentKind::Claude,
647                child_pid: None,
648                folder: temp_directory.path().to_path_buf(),
649                model: AgentModel::ClaudeSonnet5,
650                prompt: "Generate title".to_string(),
651                request_kind: AgentRequestKind::UtilityPrompt,
652                reasoning_level: ReasoningLevel::default(),
653            },
654        )
655        .await
656        .expect_err("wrapped plain-text utility output should fail");
657
658        // Assert — the provider parser extracts "plain text" from the
659        // `result` wrapper, so the protocol parser sees raw text, not JSON keys.
660        assert!(error.contains("did not match the required JSON schema"));
661        assert!(error.contains("direct_json_error:"));
662        assert!(error.contains("response:\nplain text"));
663    }
664
665    #[tokio::test]
666    /// Verifies one-shot execution recovers a trailing protocol payload when
667    /// the provider prepends extra prose before the final JSON object.
668    async fn test_submit_one_shot_with_backend_recovers_wrapped_protocol_output() {
669        // Arrange
670        let temp_directory = tempdir().expect("failed to create temp dir");
671        let mut backend = MockAgentBackend::new();
672        backend
673            .expect_build_command()
674            .times(1)
675            .returning(|request| {
676                assert!(matches!(
677                    request.request_kind,
678                    AgentRequestKind::UtilityPrompt
679                ));
680                assert_eq!(request.prompt, "Generate title");
681
682                Ok(mock_shell_command(
683                    concat!(
684                        "Now I have full context.\n",
685                        r#"{"answer":"Generated title","questions":[],"summary":null}"#
686                    ),
687                    "",
688                    0,
689                ))
690            });
691
692        // Act
693        let response = submit_one_shot_with_backend(
694            &backend,
695            OneShotRequest {
696                agent_kind: AgentKind::Claude,
697                child_pid: None,
698                folder: temp_directory.path().to_path_buf(),
699                model: AgentModel::ClaudeSonnet5,
700                prompt: "Generate title".to_string(),
701                request_kind: AgentRequestKind::UtilityPrompt,
702                reasoning_level: ReasoningLevel::default(),
703            },
704        )
705        .await
706        .expect("wrapped protocol output should succeed");
707
708        // Assert
709        assert_eq!(
710            response.response.answers(),
711            vec!["Generated title".to_string()]
712        );
713    }
714
715    #[tokio::test]
716    /// Verifies one-shot execution recovers valid output when the initial
717    /// parse fails but the protocol-repair retry returns valid protocol JSON.
718    async fn test_submit_one_shot_with_backend_recovers_via_protocol_repair() {
719        // Arrange
720        let temp_directory = tempdir().expect("failed to create temp dir");
721        let call_counter = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
722        let mut backend = MockAgentBackend::new();
723        backend.expect_build_command().times(2).returning({
724            let counter = std::sync::Arc::clone(&call_counter);
725
726            move |_| {
727                let call_number = counter.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
728
729                if call_number == 0 {
730                    Ok(mock_shell_command("plain text", "", 0))
731                } else {
732                    Ok(mock_shell_command(
733                        r#"{"answer":"Repaired title","questions":[],"summary":null}"#,
734                        "",
735                        0,
736                    ))
737                }
738            }
739        });
740
741        // Act
742        let response = submit_one_shot_with_backend(
743            &backend,
744            OneShotRequest {
745                agent_kind: AgentKind::Codex,
746                child_pid: None,
747                folder: temp_directory.path().to_path_buf(),
748                model: AgentModel::Gpt55,
749                prompt: "Generate title".to_string(),
750                request_kind: AgentRequestKind::UtilityPrompt,
751                reasoning_level: ReasoningLevel::default(),
752            },
753        )
754        .await
755        .expect("repair retry should succeed");
756
757        // Assert
758        assert_eq!(
759            response.response.answers(),
760            vec!["Repaired title".to_string()]
761        );
762    }
763
764    #[tokio::test]
765    /// Verifies one-shot execution still rejects blank utility responses
766    /// after both the original parse and the protocol-repair retry fail.
767    async fn test_submit_one_shot_with_backend_rejects_blank_utility_output() {
768        // Arrange
769        let temp_directory = tempdir().expect("failed to create temp dir");
770        let mut backend = MockAgentBackend::new();
771        backend.expect_build_command().returning(|request| {
772            assert!(matches!(
773                request.request_kind,
774                AgentRequestKind::UtilityPrompt
775            ));
776
777            Ok(mock_shell_command("   ", "", 0))
778        });
779
780        // Act
781        let error = submit_one_shot_with_backend(
782            &backend,
783            OneShotRequest {
784                agent_kind: AgentKind::Codex,
785                child_pid: None,
786                folder: temp_directory.path().to_path_buf(),
787                model: AgentModel::Gpt55,
788                prompt: "Generate title".to_string(),
789                request_kind: AgentRequestKind::UtilityPrompt,
790                reasoning_level: ReasoningLevel::default(),
791            },
792        )
793        .await
794        .expect_err("blank utility output should fail");
795
796        // Assert
797        assert!(error.contains("did not match the required JSON schema"));
798        assert!(error.contains("trimmed_len: 0 chars"));
799        assert!(error.contains("response:\n"));
800    }
801
802    #[tokio::test]
803    /// Verifies one-shot execution does not deadlock when the child delays
804    /// reading stdin until after it emits early stderr output.
805    async fn test_submit_one_shot_with_backend_writes_large_stdin_concurrently() {
806        // Arrange
807        let temp_directory = tempdir().expect("failed to create temp dir");
808        let large_prompt = "x".repeat(512 * 1024);
809        let mut backend = MockAgentBackend::new();
810        backend.expect_build_command().returning(|_| {
811            let mut command = Command::new("sh");
812            command.arg("-c").arg(
813                "printf 'warming up\\n' >&2; sleep 0.1; cat >/dev/null; printf '%s' \
814                 '{\"answer\":\"done\",\"questions\":[],\"summary\":null}'",
815            );
816            command.stdout(std::process::Stdio::piped());
817            command.stderr(std::process::Stdio::piped());
818
819            Ok(command)
820        });
821
822        // Act
823        let response = tokio::time::timeout(
824            Duration::from_secs(5),
825            submit_one_shot_with_backend(
826                &backend,
827                OneShotRequest {
828                    agent_kind: AgentKind::Claude,
829                    child_pid: None,
830                    folder: temp_directory.path().to_path_buf(),
831                    model: AgentModel::ClaudeSonnet5,
832                    prompt: large_prompt.clone(),
833                    request_kind: AgentRequestKind::UtilityPrompt,
834                    reasoning_level: ReasoningLevel::default(),
835                },
836            ),
837        )
838        .await
839        .expect("one-shot prompt should not deadlock")
840        .expect("one-shot prompt should succeed");
841
842        // Assert
843        assert_eq!(response.response.answers(), vec!["done".to_string()]);
844    }
845
846    #[tokio::test]
847    /// Verifies one-shot execution streams Claude prompts through stdin so
848    /// large review requests avoid argv length limits.
849    async fn test_submit_one_shot_with_backend_writes_prompt_to_stdin() {
850        // Arrange
851        let temp_directory = tempdir().expect("failed to create temp dir");
852        let capture_path = temp_directory.path().join("stdin.txt");
853        let mut backend = MockAgentBackend::new();
854        backend.expect_build_command().returning({
855            let capture_path = capture_path.clone();
856
857            move |_| Ok(stdin_capture_shell_command(&capture_path))
858        });
859
860        // Act
861        let response = submit_one_shot_with_backend(
862            &backend,
863            OneShotRequest {
864                agent_kind: AgentKind::Claude,
865                child_pid: None,
866                folder: temp_directory.path().to_path_buf(),
867                model: AgentModel::ClaudeSonnet5,
868                prompt: "Generate title".to_string(),
869                request_kind: AgentRequestKind::UtilityPrompt,
870                reasoning_level: ReasoningLevel::default(),
871            },
872        )
873        .await
874        .expect("one-shot prompt should succeed");
875        let captured_prompt =
876            std::fs::read_to_string(&capture_path).expect("captured stdin payload should exist");
877
878        // Assert
879        assert_eq!(response.response.answers(), vec!["captured".to_string()]);
880        assert!(captured_prompt.contains("Structured response protocol:"));
881        assert!(captured_prompt.contains("Generate title"));
882    }
883
884    #[tokio::test]
885    /// Verifies a broken stdin pipe does not hide the child exit status or
886    /// stderr when the backend exits before reading the full prompt.
887    async fn test_submit_one_shot_with_backend_preserves_exit_error_after_broken_pipe() {
888        // Arrange
889        let temp_directory = tempdir().expect("failed to create temp dir");
890        let large_prompt = "x".repeat(512 * 1024);
891        let mut backend = MockAgentBackend::new();
892        backend.expect_build_command().returning(|_| {
893            let mut command = Command::new("sh");
894            command.arg("-c").arg("printf 'auth failed' >&2; exit 7");
895            command.stdout(std::process::Stdio::piped());
896            command.stderr(std::process::Stdio::piped());
897
898            Ok(command)
899        });
900
901        // Act
902        let error = submit_one_shot_with_backend(
903            &backend,
904            OneShotRequest {
905                agent_kind: AgentKind::Claude,
906                child_pid: None,
907                folder: temp_directory.path().to_path_buf(),
908                model: AgentModel::ClaudeSonnet5,
909                prompt: large_prompt,
910                request_kind: AgentRequestKind::UtilityPrompt,
911                reasoning_level: ReasoningLevel::default(),
912            },
913        )
914        .await
915        .expect_err("one-shot prompt should surface the child exit");
916
917        // Assert
918        assert!(error.contains("exit code 7"), "error was: {error}");
919        assert!(error.contains("auth failed"), "error was: {error}");
920        assert!(
921            !error.contains("stdin payload"),
922            "stdin write error should not mask child failure: {error}"
923        );
924    }
925
926    #[tokio::test]
927    /// Verifies Claude authentication failures return actionable re-login
928    /// guidance instead of raw transport output.
929    async fn test_submit_one_shot_with_backend_surfaces_claude_auth_guidance() {
930        // Arrange
931        let temp_directory = tempdir().expect("failed to create temp dir");
932        let mut backend = MockAgentBackend::new();
933        backend.expect_build_command().returning(|_| {
934            Ok(mock_shell_command(
935                r#"{"type":"error","error":{"type":"authentication_error","message":"OAuth token has expired. Please obtain a new token or refresh your existing token."}}"#,
936                "",
937                1,
938            ))
939        });
940
941        // Act
942        let error = submit_one_shot_with_backend(
943            &backend,
944            OneShotRequest {
945                agent_kind: AgentKind::Claude,
946                child_pid: None,
947                folder: temp_directory.path().to_path_buf(),
948                model: AgentModel::ClaudeSonnet5,
949                prompt: "Generate title".to_string(),
950                request_kind: AgentRequestKind::UtilityPrompt,
951                reasoning_level: ReasoningLevel::default(),
952            },
953        )
954        .await
955        .expect_err("expired Claude auth should fail");
956
957        // Assert
958        assert!(
959            error.contains("One-shot agent command failed because Claude authentication expired")
960        );
961        assert!(error.contains("`claude auth login`"));
962        assert!(error.contains("`claude auth status`"));
963    }
964
965    #[tokio::test]
966    /// Verifies the child PID guard tracks a running process, ignores a
967    /// completed process without an id, and clears the slot when dropped.
968    async fn test_child_pid_guard_tracks_and_clears_child_process() {
969        // Arrange
970        let child_pid = Arc::new(Mutex::new(None));
971        let mut child = tokio::process::Command::new("sh")
972            .arg("-c")
973            .arg("exit 0")
974            .spawn()
975            .expect("test child should spawn");
976        let expected_pid = child.id().expect("running child should expose its pid");
977        let mut guard = ChildPidGuard::new(Some(Arc::clone(&child_pid)));
978
979        // Act
980        guard.update_from_child(&child);
981        let tracked_pid = *child_pid.lock().expect("child pid lock should succeed");
982        child.wait().await.expect("test child should exit");
983        guard.update_from_child(&child);
984        drop(guard);
985        let cleared_pid = *child_pid.lock().expect("child pid lock should succeed");
986
987        // Assert
988        assert_eq!(tracked_pid, Some(expected_pid));
989        assert_eq!(cleared_pid, None);
990    }
991
992    #[tokio::test]
993    /// Verifies app-server-backed one-shot execution returns the parsed
994    /// structured answer and usage totals.
995    async fn test_submit_one_shot_with_app_server_client_returns_protocol_response() {
996        // Arrange
997        let temp_directory = tempdir().expect("failed to create temp dir");
998        let mut app_server_client = MockAppServerClient::new();
999        app_server_client
1000            .expect_run_turn()
1001            .times(1)
1002            .returning(|request, _| {
1003                assert_eq!(request.model, AgentModel::Gpt55.as_str());
1004                assert!(matches!(
1005                    request.request_kind,
1006                    AgentRequestKind::UtilityPrompt
1007                ));
1008                assert_eq!(request.prompt.text, "Generate title");
1009
1010                Box::pin(async {
1011                    Ok(AppServerTurnResponse {
1012                        assistant_message:
1013                            r#"{"answer":"Generated title","questions":[],"summary":null}"#
1014                                .to_string(),
1015                        context_reset: false,
1016                        input_tokens: 11,
1017                        output_tokens: 7,
1018                        pid: Some(42),
1019                        provider_conversation_id: Some("thread-1".to_string()),
1020                    })
1021                })
1022            });
1023        app_server_client
1024            .expect_shutdown_session()
1025            .times(1)
1026            .returning(|_| Box::pin(async {}));
1027
1028        // Act
1029        let response = submit_one_shot_with_app_server_client(
1030            &app_server_client,
1031            OneShotRequest {
1032                agent_kind: AgentKind::Codex,
1033                child_pid: None,
1034                folder: temp_directory.path().to_path_buf(),
1035                model: AgentModel::Gpt55,
1036                prompt: "Generate title".to_string(),
1037                request_kind: AgentRequestKind::UtilityPrompt,
1038                reasoning_level: ReasoningLevel::default(),
1039            },
1040        )
1041        .await
1042        .expect("one-shot prompt should succeed");
1043
1044        // Assert
1045        assert_eq!(
1046            response.response.answers(),
1047            vec!["Generated title".to_string()]
1048        );
1049        assert_eq!(response.stats.input_tokens, 11);
1050        assert_eq!(response.stats.output_tokens, 7);
1051    }
1052
1053    #[tokio::test]
1054    /// Verifies app-server turn failures shut down the temporary session and
1055    /// clear the caller's shared child-process slot.
1056    async fn test_submit_one_shot_with_app_server_client_clears_pid_after_turn_failure() {
1057        // Arrange
1058        let temp_directory = tempdir().expect("failed to create temp dir");
1059        let child_pid = Arc::new(Mutex::new(Some(42)));
1060        let mut app_server_client = MockAppServerClient::new();
1061        app_server_client
1062            .expect_run_turn()
1063            .times(1)
1064            .returning(|_, _| {
1065                Box::pin(async {
1066                    Err(AppServerError::Provider(
1067                        "app-server turn failed".to_string(),
1068                    ))
1069                })
1070            });
1071        app_server_client
1072            .expect_shutdown_session()
1073            .times(1)
1074            .returning(|_| Box::pin(async {}));
1075
1076        // Act
1077        let error = submit_one_shot_with_app_server_client(
1078            &app_server_client,
1079            OneShotRequest {
1080                agent_kind: AgentKind::Codex,
1081                child_pid: Some(Arc::clone(&child_pid)),
1082                folder: temp_directory.path().to_path_buf(),
1083                model: AgentModel::Gpt55,
1084                prompt: "Generate title".to_string(),
1085                request_kind: AgentRequestKind::UtilityPrompt,
1086                reasoning_level: ReasoningLevel::default(),
1087            },
1088        )
1089        .await
1090        .expect_err("app-server turn failure should surface");
1091
1092        // Assert
1093        assert!(error.contains("app-server turn failed"));
1094        assert_eq!(
1095            *child_pid.lock().expect("child pid lock should succeed"),
1096            None
1097        );
1098    }
1099
1100    #[tokio::test]
1101    /// Verifies app-server-backed one-shot execution rejects plain-text
1102    /// utility output after both the original parse and the protocol-repair
1103    /// retry fail.
1104    async fn test_submit_one_shot_with_app_server_client_rejects_plain_text_utility_output() {
1105        // Arrange
1106        let temp_directory = tempdir().expect("failed to create temp dir");
1107        let mut app_server_client = MockAppServerClient::new();
1108        app_server_client
1109            .expect_run_turn()
1110            .times(2)
1111            .returning(|request, _| {
1112                assert_eq!(request.model, AgentModel::Gpt55.as_str());
1113
1114                Box::pin(async {
1115                    Ok(AppServerTurnResponse {
1116                        assistant_message: "plain text".to_string(),
1117                        context_reset: false,
1118                        input_tokens: 2,
1119                        output_tokens: 1,
1120                        pid: None,
1121                        provider_conversation_id: None,
1122                    })
1123                })
1124            });
1125        app_server_client
1126            .expect_shutdown_session()
1127            .times(1)
1128            .returning(|_| Box::pin(async {}));
1129
1130        // Act
1131        let error = submit_one_shot_with_app_server_client(
1132            &app_server_client,
1133            OneShotRequest {
1134                agent_kind: AgentKind::Codex,
1135                child_pid: None,
1136                folder: temp_directory.path().to_path_buf(),
1137                model: AgentModel::Gpt55,
1138                prompt: "Generate title".to_string(),
1139                request_kind: AgentRequestKind::UtilityPrompt,
1140                reasoning_level: ReasoningLevel::default(),
1141            },
1142        )
1143        .await
1144        .expect_err("plain-text utility output should fail");
1145
1146        // Assert
1147        assert!(error.contains("did not match the required JSON schema"));
1148        assert!(error.contains("debug_details:"));
1149        assert!(error.contains("response:\nplain text"));
1150    }
1151
1152    #[tokio::test]
1153    /// Verifies app-server-backed non-utility one-shot execution still
1154    /// rejects plain-text output after both the original parse and the
1155    /// protocol-repair retry fail.
1156    async fn test_submit_one_shot_with_app_server_client_rejects_plain_text_non_utility_output() {
1157        // Arrange
1158        let temp_directory = tempdir().expect("failed to create temp dir");
1159        let mut app_server_client = MockAppServerClient::new();
1160        app_server_client
1161            .expect_run_turn()
1162            .times(2)
1163            .returning(|request, _| {
1164                assert!(matches!(
1165                    request.request_kind,
1166                    AgentRequestKind::SessionStart
1167                ));
1168
1169                Box::pin(async {
1170                    Ok(AppServerTurnResponse {
1171                        assistant_message: "plain text".to_string(),
1172                        context_reset: false,
1173                        input_tokens: 2,
1174                        output_tokens: 1,
1175                        pid: None,
1176                        provider_conversation_id: None,
1177                    })
1178                })
1179            });
1180        app_server_client
1181            .expect_shutdown_session()
1182            .times(1)
1183            .returning(|_| Box::pin(async {}));
1184
1185        // Act
1186        let error = submit_one_shot_with_app_server_client(
1187            &app_server_client,
1188            OneShotRequest {
1189                agent_kind: AgentKind::Codex,
1190                child_pid: None,
1191                folder: temp_directory.path().to_path_buf(),
1192                model: AgentModel::Gpt55,
1193                prompt: "Generate title".to_string(),
1194                request_kind: AgentRequestKind::SessionStart,
1195                reasoning_level: ReasoningLevel::default(),
1196            },
1197        )
1198        .await
1199        .expect_err("invalid non-utility output should fail");
1200
1201        // Assert
1202        assert!(error.contains("did not match the required JSON schema"));
1203        assert!(error.contains("debug_details:"));
1204        assert!(error.contains("response:\nplain text"));
1205    }
1206}