Skip to main content

ag_protocol/
model.rs

1//! Structured response protocol data model and display helpers.
2
3use std::fmt;
4
5use schemars::JsonSchema;
6use serde::{Deserialize, Serialize};
7
8use super::question::QuestionItem;
9use super::subtask::SubtaskItem;
10
11/// Hard cap on the number of clarification questions extracted from one agent
12/// response. Prevents runaway output from flooding the question UI even when
13/// the agent ignores the prompt-level limit.
14///
15/// This constant is also injected into the protocol instruction prompt
16/// templates so the prompt-level guidance and the server-side cap stay in
17/// sync automatically.
18pub(crate) const MAX_QUESTIONS: usize = 5;
19/// Hard cap on the number of subtasks accepted from one orchestrator planning
20/// turn. Bounds how many child sessions, worktrees, and agent CLI processes a
21/// single approved plan can create even when the agent ignores the
22/// prompt-level limit.
23pub(crate) const MAX_SUBTASKS: usize = 8;
24const QUESTIONS_FIELD_DESCRIPTION_TEMPLATE: &str =
25    include_str!("template/questions_field_description.md");
26const SUBTASKS_FIELD_DESCRIPTION_TEMPLATE: &str =
27    include_str!("template/subtasks_field_description.md");
28
29/// Returns the canonical JSON Schema description for the `questions` field.
30///
31/// This is the single source of truth for the runtime-injected schema
32/// description and the matching test expectation. The static `schemars`
33/// metadata on `AgentResponse::questions` is overwritten by
34/// `inject_dynamic_schema_guidance` before any consumer observes the schema,
35/// so all schema-facing call sites must route through this helper to stay in
36/// sync.
37pub(crate) fn questions_field_description() -> String {
38    render_field_description_template(
39        QUESTIONS_FIELD_DESCRIPTION_TEMPLATE,
40        "{{ max_questions }}",
41        MAX_QUESTIONS,
42    )
43}
44
45/// Returns the canonical JSON Schema description for the `subtasks` field.
46///
47/// This mirrors [`questions_field_description`]: the static `schemars`
48/// metadata on `AgentResponse::subtasks` carries only the field title, and
49/// `inject_dynamic_schema_guidance` overwrites the description with this
50/// helper's output before any consumer observes the schema.
51pub(crate) fn subtasks_field_description() -> String {
52    render_field_description_template(
53        SUBTASKS_FIELD_DESCRIPTION_TEMPLATE,
54        "{{ max_subtasks }}",
55        MAX_SUBTASKS,
56    )
57}
58
59/// Substitutes one `{{ name }}` placeholder with a runtime cap value.
60///
61/// The placeholder is matched after collapsing whitespace runs inside every
62/// `{{ ... }}` span, because `mdformat` reflows these templates at a fixed
63/// column width and will break a line in the middle of a placeholder. Matching
64/// the literal text alone silently left the raw `{{ ... }}` in the description
65/// shown to models, so normalization keeps the templates safe to reformat.
66fn render_field_description_template(template: &str, placeholder: &str, value: usize) -> String {
67    let mut rendered = String::with_capacity(template.len());
68    let mut remaining = template.trim_end();
69
70    while let Some(open_index) = remaining.find("{{") {
71        let after_open = &remaining[open_index..];
72        let Some(close_end) = after_open.find("}}").map(|index| index + "}}".len()) else {
73            break;
74        };
75
76        rendered.push_str(&remaining[..open_index]);
77        rendered.push_str(&collapse_whitespace(&after_open[..close_end]));
78        remaining = &after_open[close_end..];
79    }
80    rendered.push_str(remaining);
81
82    rendered.replace(placeholder, &value.to_string())
83}
84
85/// Collapses every whitespace run in `text` to one space.
86fn collapse_whitespace(text: &str) -> String {
87    text.split_whitespace().collect::<Vec<_>>().join(" ")
88}
89
90/// Protocol-owned request family preserved across prompt submission and repair
91/// retries.
92///
93/// Session discussion turns and isolated utility prompts share the same
94/// top-level [`AgentResponse`] schema. Agentty still carries the request
95/// family through transport boundaries so call sites can keep one consistent
96/// protocol contract even when some callers ignore parts of the response, such
97/// as the optional top-level `summary`.
98#[derive(Debug, Clone, Copy, PartialEq, Eq)]
99#[non_exhaustive]
100pub enum ProtocolRequestProfile {
101    /// Interactive session turn.
102    SessionTurn,
103    /// Isolated utility prompt.
104    UtilityPrompt,
105}
106
107/// Structured session summary block emitted alongside protocol messages.
108///
109/// Session-discussion turns use this object instead of embedding the change
110/// summary inside `answer` message text. One-shot prompts set the top-level
111/// `summary` field to `null`.
112#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
113#[schemars(
114    title = "AgentResponseSummary",
115    description = "Structured session summary block emitted alongside protocol messages instead \
116                   of embedding the change summary inside `answer` markdown on session-discussion \
117                   turns."
118)]
119pub struct AgentResponseSummary {
120    /// Cumulative summary of active changes on the current session branch.
121    #[schemars(
122        title = "session",
123        description = "Cumulative summary of active changes on the current session branch."
124    )]
125    pub session: String,
126    /// Concise summary of only the work completed in the current turn.
127    #[schemars(
128        title = "turn",
129        description = "Concise summary of only the work completed in the current turn."
130    )]
131    pub turn: String,
132}
133
134/// Agent-reported disposition for one forge review thread.
135#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
136#[serde(rename_all = "snake_case")]
137#[schemars(
138    title = "ReviewCommentResolution",
139    description = "Disposition reported for one targeted forge review thread."
140)]
141pub enum ReviewCommentResolution {
142    /// The agent completed the requested action and the thread can be
143    /// resolved after its reply is posted.
144    Fixed,
145    /// The agent determined that no code or documentation change was needed;
146    /// its explanatory reply is posted while the thread remains open.
147    NoChangeNeeded,
148}
149
150/// Structured outcome for one forge review thread targeted by the turn.
151#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
152#[schemars(
153    title = "ReviewCommentOutcome",
154    description = "Structured outcome for one forge review thread explicitly included in the turn \
155                   prompt."
156)]
157pub struct ReviewCommentOutcome {
158    /// Concise forge reply explaining what changed or why no change was needed.
159    #[schemars(
160        title = "reply",
161        description = "Concise reply suitable for posting to the forge review thread."
162    )]
163    pub reply: String,
164    /// Whether the thread was fixed or did not require a change.
165    #[schemars(
166        title = "resolution",
167        description = "Whether the targeted thread was fixed or required no change."
168    )]
169    pub resolution: ReviewCommentResolution,
170    /// Opaque forge thread identifier copied exactly from the turn prompt.
171    #[schemars(
172        title = "thread_id",
173        description = "Opaque forge thread identifier copied exactly from the turn prompt."
174    )]
175    pub thread_id: String,
176}
177
178/// Wire-format protocol payload used for schema-driven provider output.
179///
180/// Providers that support output schemas (for example, Codex app-server) are
181/// asked to emit this object as the entire assistant response payload.
182#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
183#[schemars(
184    title = "AgentResponse",
185    description = "Wire-format protocol payload used for schema-driven provider output. Return \
186                   this object as the entire assistant response payload. Providers that support \
187                   output schemas (for example, Codex app-server) are asked to emit this object \
188                   directly."
189)]
190pub struct AgentResponse {
191    /// Markdown answer text emitted for this turn.
192    #[serde(default)]
193    #[schemars(
194        title = "answer",
195        description = "Markdown answer text for delivered work, status updates, or concise \
196                       completion notes. Keep clarification requests out of this field and emit \
197                       them through `questions` instead."
198    )]
199    pub answer: String,
200    /// Ordered clarification questions emitted for this turn.
201    ///
202    /// The canonical JSON Schema description for this field is produced by
203    /// [`questions_field_description`] and injected at schema generation time
204    /// by `inject_dynamic_schema_guidance`. The static `schemars` metadata
205    /// here only sets the field title; the description is intentionally
206    /// omitted so the helper is the single source of truth.
207    #[serde(default)]
208    #[schemars(title = "questions")]
209    pub questions: Vec<QuestionItem>,
210    /// Per-thread outcomes for an agent-driven forge comment-resolution turn.
211    ///
212    /// Ordinary session and utility turns leave this empty. Resolution
213    /// workflows accept only identifiers explicitly allowlisted in the turn
214    /// prompt before applying any forge-side effect.
215    #[serde(default)]
216    #[schemars(
217        title = "review_comment_outcomes",
218        description = "Per-thread outcomes for an agent-driven forge comment-resolution turn. \
219                       Emit an empty array unless the prompt explicitly supplies forge thread \
220                       IDs. Copy each reported `thread_id` exactly from the prompt."
221    )]
222    pub review_comment_outcomes: Vec<ReviewCommentOutcome>,
223    /// Proposed child-session subtasks for an orchestrator planning turn.
224    ///
225    /// The canonical JSON Schema description is produced by
226    /// [`subtasks_field_description`] and injected at schema generation time,
227    /// so the static `schemars` metadata here only sets the field title.
228    /// Ordinary session and utility turns leave this empty, and orchestration
229    /// consumers ignore it unless the turn prompt asked for a plan.
230    #[serde(default)]
231    #[schemars(title = "subtasks")]
232    pub subtasks: Vec<SubtaskItem>,
233    /// Structured summary for session-discussion turns, or `None` for legacy
234    /// payloads and one-shot prompts.
235    #[serde(default)]
236    #[schemars(
237        title = "summary",
238        description = "Structured summary for session-discussion turns, kept outside `answer` \
239                       markdown. Use `null` for one-shot prompts and legacy payloads."
240    )]
241    pub summary: Option<AgentResponseSummary>,
242}
243
244impl AgentResponse {
245    /// Creates a plain response from raw text as one `answer` string.
246    pub fn plain(text: impl Into<String>) -> Self {
247        Self {
248            answer: text.into(),
249            questions: Vec::new(),
250            review_comment_outcomes: Vec::new(),
251            subtasks: Vec::new(),
252            summary: None,
253        }
254    }
255
256    /// Returns display text by joining non-empty answer and question text with
257    /// blank lines.
258    pub fn to_display_text(&self) -> String {
259        let mut display_messages = Vec::new();
260        push_display_message(&mut display_messages, &self.answer);
261        push_question_display_messages(&mut display_messages, &self.questions);
262
263        display_messages.join("\n\n")
264    }
265
266    /// Returns transcript text for session output by joining non-empty
267    /// `answer` content with blank lines.
268    pub fn to_answer_display_text(&self) -> String {
269        let mut display_messages = Vec::new();
270        push_display_message(&mut display_messages, &self.answer);
271
272        display_messages.join("\n\n")
273    }
274
275    /// Returns the answer as one single-item vector when it is non-empty.
276    pub fn answers(&self) -> Vec<String> {
277        let answer = self.to_answer_display_text();
278        if answer.is_empty() {
279            return Vec::new();
280        }
281
282        vec![answer]
283    }
284
285    /// Returns up to [`MAX_QUESTIONS`] clarification questions in response
286    /// order.
287    pub fn question_items(&self) -> Vec<QuestionItem> {
288        self.questions.iter().take(MAX_QUESTIONS).cloned().collect()
289    }
290
291    /// Returns up to [`MAX_SUBTASKS`] proposed subtasks in response order.
292    ///
293    /// Callers must still reject plans whose subtasks overlap or whose keys
294    /// collide; this only bounds how much of a runaway plan is considered.
295    pub fn subtask_items(&self) -> Vec<SubtaskItem> {
296        self.subtasks.iter().take(MAX_SUBTASKS).cloned().collect()
297    }
298}
299
300/// Structured response parsing failure details.
301#[derive(Debug, Clone, PartialEq, Eq)]
302pub enum AgentResponseParseError {
303    /// Response was empty or whitespace-only.
304    Empty,
305    /// Response was JSON, but it did not satisfy the structured protocol
306    /// contract.
307    InvalidFormat {
308        /// Explanation of the protocol contract violation.
309        reason: String,
310    },
311}
312
313impl fmt::Display for AgentResponseParseError {
314    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
315        match self {
316            Self::Empty => write!(formatter, "response is empty"),
317            Self::InvalidFormat { reason } => {
318                write!(formatter, "response is not valid protocol JSON: {reason}")
319            }
320        }
321    }
322}
323
324impl std::error::Error for AgentResponseParseError {}
325
326/// Appends one non-empty display message.
327fn push_display_message(display_messages: &mut Vec<String>, text: &str) {
328    if text.trim().is_empty() {
329        return;
330    }
331
332    display_messages.push(text.to_string());
333}
334
335/// Appends non-empty clarification question text in order.
336fn push_question_display_messages(display_messages: &mut Vec<String>, questions: &[QuestionItem]) {
337    for question in questions {
338        push_display_message(display_messages, &question.text);
339    }
340}
341
342#[cfg(test)]
343mod tests {
344    use super::*;
345
346    #[test]
347    /// Ensures the dynamic `questions` field description renders from the
348    /// checked-in prompt-schema template.
349    fn test_questions_field_description_renders_template_limit() {
350        // Arrange, Act
351        let description = questions_field_description();
352        let normalized_description = description.split_whitespace().collect::<Vec<_>>().join(" ");
353
354        // Assert
355        assert!(normalized_description.contains("Emit at most 5 items"));
356        assert!(normalized_description.contains("Execute the agreed work"));
357        assert!(!description.contains("{{ max_questions }}"));
358    }
359
360    #[test]
361    /// Ensures display text includes the answer and clarification questions in
362    /// order.
363    fn test_agent_response_to_display_text_joins_answer_and_questions() {
364        // Arrange
365        let response = AgentResponse {
366            answer: "Primary answer".to_string(),
367            questions: vec![QuestionItem::new("Need one clarification.")],
368            review_comment_outcomes: Vec::new(),
369            subtasks: Vec::new(),
370            summary: None,
371        };
372
373        // Act
374        let display_text = response.to_display_text();
375
376        // Assert
377        assert_eq!(display_text, "Primary answer\n\nNeed one clarification.");
378    }
379
380    #[test]
381    /// Preserves review-comment outcomes through the wire JSON contract.
382    fn test_agent_response_review_comment_outcomes_round_trip() {
383        // Arrange
384        let response = AgentResponse {
385            answer: "Addressed the comment.".to_string(),
386            questions: Vec::new(),
387            review_comment_outcomes: vec![ReviewCommentOutcome {
388                reply: "Added the missing validation.".to_string(),
389                resolution: ReviewCommentResolution::Fixed,
390                thread_id: "thread-42".to_string(),
391            }],
392            subtasks: Vec::new(),
393            summary: None,
394        };
395
396        // Act
397        let serialized = serde_json::to_string(&response).expect("response should serialize");
398        let deserialized = serde_json::from_str::<AgentResponse>(&serialized)
399            .expect("response should deserialize");
400
401        // Assert
402        assert_eq!(deserialized, response);
403        assert!(serialized.contains(r#""resolution":"fixed""#));
404    }
405
406    #[test]
407    /// Ensures question extraction respects the protocol question cap.
408    fn test_agent_response_question_items_applies_question_cap() {
409        // Arrange
410        let response = AgentResponse {
411            answer: String::new(),
412            questions: (0..=MAX_QUESTIONS)
413                .map(|index| QuestionItem::new(format!("Question {index}")))
414                .collect(),
415            review_comment_outcomes: Vec::new(),
416            subtasks: Vec::new(),
417            summary: None,
418        };
419
420        // Act
421        let questions = response.question_items();
422
423        // Assert
424        assert_eq!(questions.len(), MAX_QUESTIONS);
425    }
426
427    #[test]
428    /// Ensures subtask extraction respects the protocol subtask cap so a
429    /// runaway plan cannot fan out past the bounded child-session budget.
430    fn test_agent_response_subtask_items_applies_subtask_cap() {
431        // Arrange
432        let response = AgentResponse {
433            answer: String::new(),
434            questions: Vec::new(),
435            review_comment_outcomes: Vec::new(),
436            subtasks: (0..=MAX_SUBTASKS).map(test_subtask).collect(),
437            summary: None,
438        };
439
440        // Act
441        let subtasks = response.subtask_items();
442
443        // Assert
444        assert_eq!(subtasks.len(), MAX_SUBTASKS);
445        assert_eq!(subtasks[0].task_key, "task-0");
446    }
447
448    #[test]
449    /// Preserves proposed subtasks through the wire JSON contract and keeps
450    /// them absent from ordinary responses.
451    fn test_agent_response_subtasks_round_trip() {
452        // Arrange
453        let response = AgentResponse {
454            answer: "Proposed a plan.".to_string(),
455            questions: Vec::new(),
456            review_comment_outcomes: Vec::new(),
457            subtasks: vec![test_subtask(1)],
458            summary: None,
459        };
460
461        // Act
462        let serialized = serde_json::to_string(&response).expect("response should serialize");
463        let deserialized = serde_json::from_str::<AgentResponse>(&serialized)
464            .expect("response should deserialize");
465
466        // Assert
467        assert_eq!(deserialized, response);
468        assert!(serialized.contains(r#""task_key":"task-1""#));
469        assert!(AgentResponse::plain("no plan").subtask_items().is_empty());
470    }
471
472    #[test]
473    /// Ensures `touched_areas` defaults to an empty list so a malformed
474    /// subtask still parses and can be rejected by plan validation instead of
475    /// failing the whole turn.
476    fn test_subtask_item_defaults_touched_areas() {
477        // Arrange
478        let raw = r#"{"prompt":"Do the work","task_key":"task-1","title":"Work"}"#;
479
480        // Act
481        let subtask =
482            serde_json::from_str::<SubtaskItem>(raw).expect("subtask should parse without areas");
483
484        // Assert
485        assert!(subtask.touched_areas.is_empty());
486    }
487
488    #[test]
489    /// Keeps the injected `subtasks` schema description in sync with the
490    /// server-side cap the parser enforces.
491    fn test_subtasks_field_description_reports_the_subtask_cap() {
492        // Arrange, Act
493        let description = subtasks_field_description();
494
495        // Assert
496        assert!(description.contains(&format!("at most {MAX_SUBTASKS} items")));
497        assert!(description.contains("without wildcard patterns"));
498        assert!(!description.contains("{{"));
499    }
500
501    #[test]
502    /// Substitutes a cap placeholder that markdown reflowing wrapped across a
503    /// line break, so reformatting a template cannot leak raw `{{ ... }}` text
504    /// into the schema description models read.
505    fn test_field_description_template_survives_a_wrapped_placeholder() {
506        // Arrange
507        let template = "Emit at most {{\nmax_items }} items, and no more.\n";
508
509        // Act
510        let rendered = render_field_description_template(template, "{{ max_items }}", 4);
511
512        // Assert
513        assert_eq!(rendered, "Emit at most 4 items, and no more.");
514    }
515
516    #[test]
517    /// Leaves an unterminated placeholder untouched instead of truncating the
518    /// remaining guidance text.
519    fn test_field_description_template_keeps_unterminated_placeholder_text() {
520        // Arrange
521        let template = "Emit at most {{ max_items items.";
522
523        // Act
524        let rendered = render_field_description_template(template, "{{ max_items }}", 4);
525
526        // Assert
527        assert_eq!(rendered, "Emit at most {{ max_items items.");
528    }
529
530    /// Builds one deterministic subtask with a file-disjoint touched area.
531    fn test_subtask(index: usize) -> SubtaskItem {
532        SubtaskItem {
533            prompt: format!("Complete work item {index}"),
534            task_key: format!("task-{index}"),
535            title: format!("Work item {index}"),
536            touched_areas: vec![format!("crates/area-{index}/")],
537        }
538    }
539}