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

ag_agent/agent/
submission.rs
