zag_agent/
exit_mode.rs

1//! Library-level support for `--exit` interactive mode.
2//!
3//! When a session is launched with `--exit [<hint>]`, the user prompt is
4//! augmented with instructions telling the agent to call
5//! `zag ps kill self <result>` (or `zag ps kill self --file <path>`) to
6//! terminate the session and submit the final result.
7//!
8//! This module owns the prompt template, the typed [`ExitHint`] /
9//! [`ExitConstraints`] state, and the validation logic used by
10//! `zag ps kill` to accept or reject a submitted result.
11
12use crate::json_validation::{validate_json, validate_json_schema};
13use serde::{Deserialize, Deserializer, Serialize, Serializer};
14
15/// State of the `--exit` flag for a single session.
16///
17/// `Bare` means `--exit` was passed with no argument — the agent must
18/// terminate via `zag ps kill self`, but the result is unconstrained.
19/// `Provided(s)` carries a non-empty human-readable description and also
20/// makes `zag ps kill` require a non-empty result.
21///
22/// On-disk JSON shape is intentionally a string: `""` for `Bare`, `"foo"`
23/// for `Provided("foo")`. A missing field (`None` in `Option<ExitHint>`)
24/// means `--exit` was not set at all. This keeps the disk format flat and
25/// human-inspectable while giving Rust callers a typed enum instead of
26/// the previous `Option<Option<String>>`.
27#[derive(Debug, Clone, PartialEq, Eq)]
28pub enum ExitHint {
29    Bare,
30    Provided(String),
31}
32
33impl ExitHint {
34    /// Build from a clap-style `Option<String>`: `None` and `Some("")`
35    /// both yield [`ExitHint::Bare`]; anything else is [`Provided`].
36    pub fn from_optional(s: Option<String>) -> Self {
37        match s {
38            Some(s) if !s.trim().is_empty() => Self::Provided(s),
39            _ => Self::Bare,
40        }
41    }
42
43    /// The hint text iff non-empty. `None` for [`Bare`].
44    pub fn as_str(&self) -> Option<&str> {
45        match self {
46            Self::Provided(s) => Some(s.as_str()),
47            Self::Bare => None,
48        }
49    }
50}
51
52impl Serialize for ExitHint {
53    fn serialize<S: Serializer>(&self, ser: S) -> Result<S::Ok, S::Error> {
54        match self {
55            Self::Bare => ser.serialize_str(""),
56            Self::Provided(s) => ser.serialize_str(s),
57        }
58    }
59}
60
61impl<'de> Deserialize<'de> for ExitHint {
62    fn deserialize<D: Deserializer<'de>>(de: D) -> Result<Self, D::Error> {
63        let s = String::deserialize(de)?;
64        Ok(if s.trim().is_empty() {
65            Self::Bare
66        } else {
67            Self::Provided(s)
68        })
69    }
70}
71
72fn is_false(v: &bool) -> bool {
73    !v
74}
75
76/// All exit-mode constraints captured at session launch.
77///
78/// Stored on [`crate::session::SessionEntry`] as `Option<ExitConstraints>`:
79/// `None` means `--exit` was not set; `Some(_)` means the session is in exit
80/// mode and the agent must terminate via `zag ps kill self <result>`.
81///
82/// The validator [`ExitConstraints::validate`] checks the submitted result
83/// against the hint requirement (non-empty when [`ExitHint::Provided`]),
84/// the JSON-validity requirement (when `json_mode` or `schema` is set), and
85/// the schema (when `schema` is set).
86#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
87pub struct ExitConstraints {
88    #[serde(default, skip_serializing_if = "Option::is_none")]
89    pub hint: Option<ExitHint>,
90    #[serde(default, skip_serializing_if = "is_false")]
91    pub json_mode: bool,
92    #[serde(default, skip_serializing_if = "Option::is_none")]
93    pub schema: Option<serde_json::Value>,
94}
95
96impl ExitConstraints {
97    /// Validate a submitted result against these constraints. See
98    /// [`validate_exit_result`] for the underlying logic — this is a
99    /// thin wrapper.
100    pub fn validate(&self, result: &str) -> Result<(), ExitValidationError> {
101        validate_exit_result(
102            result,
103            self.hint.as_ref().and_then(|h| h.as_str()),
104            self.json_mode,
105            self.schema.as_ref(),
106        )
107    }
108}
109
110/// Raw `prompts/exit/1_0_0.md` source, including YAML front matter.
111const EXIT_TEMPLATE_SOURCE: &str = include_str!("../prompts/exit/1_0_0.md");
112
113/// Exit prompt template (front matter stripped) — `{HINT_SECTION}`,
114/// `{JSON_INSTRUCTION}`, `{SCHEMA_INSTRUCTION}` are replaced at run time.
115pub fn exit_template() -> &'static str {
116    crate::prompts::strip_front_matter(EXIT_TEMPLATE_SOURCE)
117}
118
119/// Build the suffix appended to a user prompt when a session is launched
120/// with `--exit`.
121///
122/// * `hint` — optional human-readable description of the expected result.
123/// * `json_mode` — whether `--json` was set; the agent is told the result
124///   must be valid JSON.
125/// * `json_schema` — optional schema; if present, the schema is rendered
126///   verbatim so the agent knows what shape to produce.
127pub fn build_exit_suffix(
128    hint: Option<&str>,
129    json_mode: bool,
130    json_schema: Option<&serde_json::Value>,
131) -> String {
132    let hint_section = match hint.map(str::trim).filter(|s| !s.is_empty()) {
133        Some(h) => format!("Expected result: {h}\n\n"),
134        None => String::new(),
135    };
136    let json_instruction = if json_mode || json_schema.is_some() {
137        "The result you pass to `zag ps kill self` MUST be valid JSON. \
138         Do not wrap it in markdown fences — pass the raw JSON string.\n\n"
139            .to_string()
140    } else {
141        String::new()
142    };
143    let schema_instruction = match json_schema {
144        Some(schema) => {
145            let pretty = serde_json::to_string_pretty(schema).unwrap_or_default();
146            format!(
147                "The JSON result MUST validate against this schema:\n\n```json\n{pretty}\n```\n\n"
148            )
149        }
150        None => String::new(),
151    };
152    let rendered = exit_template()
153        .replace("{HINT_SECTION}", &hint_section)
154        .replace("{JSON_INSTRUCTION}", &json_instruction)
155        .replace("{SCHEMA_INSTRUCTION}", &schema_instruction);
156    // Belt-and-braces: if the template gains a new placeholder and the
157    // renderer isn't updated, fail loudly in debug/test rather than
158    // leaking `{TOKEN}` into the agent's prompt.
159    debug_assert!(
160        !rendered.contains("{HINT_SECTION}")
161            && !rendered.contains("{JSON_INSTRUCTION}")
162            && !rendered.contains("{SCHEMA_INSTRUCTION}"),
163        "exit prompt template contains unrendered placeholder"
164    );
165    rendered
166}
167
168/// Reason a `zag ps kill` invocation was rejected. The CLI prints the
169/// `Display` impl to stderr; the agent is expected to read the message
170/// and self-correct.
171#[derive(Debug)]
172pub enum ExitValidationError {
173    /// The session was launched with a non-empty `--exit` hint but the
174    /// kill was called with an empty (or missing) result.
175    EmptyResult { hint: String },
176    /// `--json` was set but the result is not valid JSON.
177    InvalidJson { detail: String },
178    /// `--json-schema` was set and the result failed schema validation.
179    SchemaViolations { errors: Vec<String> },
180}
181
182impl std::fmt::Display for ExitValidationError {
183    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
184        match self {
185            Self::EmptyResult { hint } => {
186                write!(
187                    f,
188                    "Cannot terminate: a non-empty result is required (hint: {hint}). \
189                     Re-run with `zag ps kill self \"<your-result>\"` or \
190                     `zag ps kill self --file <path>`."
191                )
192            }
193            Self::InvalidJson { detail } => {
194                write!(
195                    f,
196                    "Result is not valid JSON: {detail}. The session was launched with \
197                     --json, so the result must be a JSON value (object, array, string, \
198                     number, boolean, or null). Do not include markdown fences."
199                )
200            }
201            Self::SchemaViolations { errors } => {
202                writeln!(
203                    f,
204                    "Result failed JSON-schema validation. Fix the result and call kill again:"
205                )?;
206                for e in errors {
207                    writeln!(f, "  - {e}")?;
208                }
209                Ok(())
210            }
211        }
212    }
213}
214
215impl std::error::Error for ExitValidationError {}
216
217/// Validate a result string against the constraints recorded on a session
218/// at launch time. Returns `Ok(())` if the kill should proceed.
219pub fn validate_exit_result(
220    result: &str,
221    exit_hint: Option<&str>,
222    json_mode: bool,
223    json_schema: Option<&serde_json::Value>,
224) -> Result<(), ExitValidationError> {
225    if let Some(hint) = exit_hint
226        && !hint.trim().is_empty()
227        && result.trim().is_empty()
228    {
229        return Err(ExitValidationError::EmptyResult {
230            hint: hint.to_string(),
231        });
232    }
233
234    if let Some(schema) = json_schema {
235        if let Err(errors) = validate_json_schema(result, schema) {
236            return Err(ExitValidationError::SchemaViolations { errors });
237        }
238    } else if json_mode && let Err(detail) = validate_json(result) {
239        return Err(ExitValidationError::InvalidJson { detail });
240    }
241
242    Ok(())
243}
244
245#[cfg(test)]
246#[path = "exit_mode_tests.rs"]
247mod tests;
zag_agent/exit_mode.rs

zag_agent/
exit_mode.rs
