Skip to main content

agent_first_psql/
emit.rs

1use crate::output_fmt;
2use crate::types::Output;
3use agent_first_data::{OutputFormat, OutputTo};
4use serde_json::Value;
5use std::io::Write;
6use std::sync::OnceLock;
7
8static OUTPUT_TO: OnceLock<OutputTo> = OnceLock::new();
9
10/// Resolve the process-wide event destination from raw arguments.
11///
12/// The destination follows the invocation's *consumption mode*. A plain query
13/// emits at most one terminal event and exits, so it is a finite one-shot
14/// command and splits by kind. `--mode pipe` and `--stream-rows` instead emit an
15/// interleaved multi-event stream whose consumer reads it in order, so they are
16/// event streams: every event stays on one stream, because splitting a stream
17/// across stdout and stderr loses the ordering that makes it a stream.
18pub fn install_output_to_from_raw(raw_args: &[String]) -> Result<(), String> {
19    let mut requested: Option<OutputTo> = None;
20    let mut event_stream = false;
21    // Walk arguments the same way `is_psql_mode_requested` does, so this scan
22    // and clap agree on which token is a flag and which is a flag's value.
23    // `--sql` and `--sql-file` take hyphen values, so without the value table
24    // `--sql --stream-rows` (a SQL comment) would read as a streaming flag.
25    let mut top_level = true;
26    let mut i = 1usize;
27    while i < raw_args.len() {
28        let arg = raw_args[i].as_str();
29        if arg == "--" {
30            break;
31        }
32
33        // `--output-to` is global, so it still applies after a subcommand.
34        if let Some(value) = arg.strip_prefix("--output-to=") {
35            requested = Some(OutputTo::parse(value)?);
36            i += 1;
37            continue;
38        }
39        if arg == "--output-to" {
40            let value = raw_args.get(i + 1).ok_or_else(|| {
41                "--output-to requires a value: expected split, stdout, or stderr".to_string()
42            })?;
43            requested = Some(OutputTo::parse(value)?);
44            i += 2;
45            continue;
46        }
47
48        if !top_level {
49            i += 1;
50            continue;
51        }
52
53        // `--stream-rows` and `--mode` are top-level only. Counting them after a
54        // subcommand would pick a destination for an invocation clap is about to
55        // reject, sending the usage error to the wrong stream.
56        if arg == "--stream-rows" || arg == "--mode=pipe" {
57            event_stream = true;
58            i += 1;
59            continue;
60        }
61        if arg == "--mode" {
62            if raw_args.get(i + 1).is_some_and(|value| value == "pipe") {
63                event_stream = true;
64            }
65            i += 2;
66            continue;
67        }
68        if crate::cli::top_level_arg_consumes_two_values(arg) {
69            i += if arg.contains('=') { 2 } else { 3 };
70            continue;
71        }
72        if crate::cli::top_level_arg_consumes_value(arg) {
73            i += if arg.contains('=') { 1 } else { 2 };
74            continue;
75        }
76        if arg.starts_with('-') {
77            i += 1;
78            continue;
79        }
80        top_level = false;
81        i += 1;
82    }
83
84    let output_to = match (requested, event_stream) {
85        (Some(OutputTo::Split), true) => {
86            return Err(
87                "--output-to split cannot be combined with --mode pipe or --stream-rows: an ordered event stream must stay on one stream. Use --output-to stdout (the streaming default) or --output-to stderr"
88                    .to_string(),
89            );
90        }
91        (Some(explicit), _) => explicit,
92        (None, true) => OutputTo::Stdout,
93        (None, false) => OutputTo::Split,
94    };
95    let _ = OUTPUT_TO.set(output_to);
96    Ok(())
97}
98
99pub fn output_to() -> OutputTo {
100    OUTPUT_TO.get().copied().unwrap_or(OutputTo::Split)
101}
102
103pub fn emit_cli_error(
104    msg: &str,
105    hint: Option<&str>,
106    format: OutputFormat,
107) -> Result<(), agent_first_data::CliEmitterError> {
108    let mut emitter =
109        agent_first_data::CliEmitter::from_output_to(output_to(), format).with_strict_protocol();
110    let event = agent_first_data::json_error(crate::protocol::error_code::INVALID_REQUEST, msg)
111        .hint_if_some(hint)
112        .build()
113        .map_err(agent_first_data::CliEmitterError::Build)?;
114    emitter.emit(event)
115}
116
117pub fn emit_value(
118    value: Value,
119    format: OutputFormat,
120) -> Result<(), agent_first_data::CliEmitterError> {
121    let mut emitter =
122        agent_first_data::CliEmitter::from_output_to(output_to(), format).with_strict_protocol();
123    emitter.emit_validated_value(value)
124}
125
126pub fn emit_output(
127    out: &Output,
128    format: OutputFormat,
129) -> Result<(), agent_first_data::CliEmitterError> {
130    output_fmt::emit_process_output(out, format, output_to())
131}
132
133#[allow(clippy::disallowed_methods)]
134pub fn write_result_text(text: &str) -> std::io::Result<()> {
135    let mut writer: Box<dyn Write> = match output_to() {
136        OutputTo::Stderr => Box::new(std::io::stderr()),
137        OutputTo::Split | OutputTo::Stdout => Box::new(std::io::stdout()),
138    };
139    writer.write_all(text.as_bytes())?;
140    writer.flush()
141}