use serde::{Deserialize, Serialize};
pub const MAX_FINAL_OUTPUT_BYTES: usize = 256 * 1024;
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct OutputSpec {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub format: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub instructions: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub example: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub schema: Option<serde_json::Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub validator: Option<String>,
}
impl OutputSpec {
pub fn is_empty(&self) -> bool {
self.format.is_none()
&& self.instructions.is_none()
&& self.example.is_none()
&& self.schema.is_none()
&& self.validator.is_none()
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct FinalOutput {
pub content: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub format: Option<String>,
pub stage: String,
pub submitted_at: i64,
#[serde(default)]
pub truncated: bool,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub artifacts: Vec<String>,
}
impl FinalOutput {
pub fn new(content: &str, format: Option<String>, stage: String, submitted_at: i64) -> Self {
let truncated = content.len() > MAX_FINAL_OUTPUT_BYTES;
let kept = crate::text::truncate_at_boundary(content, MAX_FINAL_OUTPUT_BYTES);
Self {
content: kept.to_string(),
format,
stage,
submitted_at,
truncated,
artifacts: Vec::new(),
}
}
pub fn with_artifacts(mut self, artifacts: Vec<String>) -> Self {
self.artifacts = artifacts;
self
}
pub fn descriptor(&self) -> FinalOutputDescriptor {
FinalOutputDescriptor {
format: self.format.clone(),
stage: self.stage.clone(),
submitted_at: self.submitted_at,
bytes: self.content.len(),
truncated: self.truncated,
artifacts: self.artifacts.clone(),
}
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct FinalOutputDescriptor {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub format: Option<String>,
pub stage: String,
pub submitted_at: i64,
#[serde(default)]
pub bytes: usize,
#[serde(default)]
pub truncated: bool,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub artifacts: Vec<String>,
}
pub const FINAL_OUTPUT_FILE: &str = "final_output";
pub fn resolve_output_spec(
agent: Option<&OutputSpec>,
stage: Option<&OutputSpec>,
request: Option<&OutputSpec>,
) -> Option<OutputSpec> {
if agent.is_none() && stage.is_none() && request.is_none() {
return None;
}
fn field<T: Clone>(
agent: Option<&OutputSpec>,
stage: Option<&OutputSpec>,
request: Option<&OutputSpec>,
get: impl Fn(&OutputSpec) -> Option<T>,
) -> Option<T> {
request
.and_then(&get)
.or_else(|| stage.and_then(&get))
.or_else(|| agent.and_then(&get))
}
let declared_format = field(agent, stage, None, |s| s.format.clone());
let requested_format = request.and_then(|r| r.format.clone());
let reshaped = requested_format.is_some() && requested_format != declared_format;
let shape_field = |get: fn(&OutputSpec) -> Option<serde_json::Value>| match reshaped {
true => request.and_then(get),
false => field(agent, stage, request, get),
};
let validator = match reshaped {
true => request.and_then(|r| r.validator.clone()),
false => field(agent, stage, request, |s| s.validator.clone()),
};
Some(OutputSpec {
format: field(agent, stage, request, |s| s.format.clone()),
instructions: field(agent, stage, request, |s| s.instructions.clone()),
example: field(agent, stage, request, |s| s.example.clone()),
schema: shape_field(|s| s.schema.clone()),
validator,
})
}
pub fn describe_spec(spec: &OutputSpec) -> String {
let mut parts = Vec::new();
if let Some(format) = &spec.format {
parts.push(format!("Return it in this format: {format}."));
}
if let Some(instructions) = &spec.instructions {
parts.push(instructions.clone());
}
if let Some(schema) = &spec.schema {
parts.push(format!(
"It must be JSON valid against this schema:\n{schema}"
));
}
if let Some(example) = &spec.example {
parts.push(format!(
"Here is an example of the expected shape:\n{example}"
));
}
if !parts.is_empty() {
parts.push(
"This governs how the answer is presented. Where anything else you were told says \
to present it differently - its length, its structure, what to lead with - follow \
this."
.to_string(),
);
}
parts.join("\n\n")
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
fn spec(format: Option<&str>, schema: Option<serde_json::Value>) -> OutputSpec {
OutputSpec {
format: format.map(str::to_string),
schema,
..OutputSpec::default()
}
}
#[test]
fn artifacts_attach_to_a_submission_and_reach_the_descriptor() {
let output = FinalOutput::new(
"the summary",
Some("markdown".to_string()),
"present".to_string(),
42,
)
.with_artifacts(vec![
"data/dataset.csv".to_string(),
"report.pdf".to_string(),
]);
assert_eq!(output.artifacts, ["data/dataset.csv", "report.pdf"]);
assert_eq!(output.descriptor().artifacts, output.artifacts);
assert_eq!(output.descriptor().bytes, "the summary".len());
}
#[test]
fn a_submission_carries_no_artifacts_unless_given_some() {
assert!(
FinalOutput::new("x", None, "present".to_string(), 0)
.artifacts
.is_empty()
);
}
#[test]
fn empty_spec_constrains_nothing() {
assert!(OutputSpec::default().is_empty());
assert!(!spec(Some("json"), None).is_empty());
assert!(!spec(None, Some(json!({}))).is_empty());
assert!(
!OutputSpec {
instructions: Some("be brief".to_string()),
..OutputSpec::default()
}
.is_empty()
);
assert!(
!OutputSpec {
example: Some("<doc/>".to_string()),
..OutputSpec::default()
}
.is_empty()
);
}
#[test]
fn no_level_asking_for_output_resolves_to_none() {
assert_eq!(resolve_output_spec(None, None, None), None);
}
#[test]
fn later_levels_win_field_by_field() {
let agent = OutputSpec {
format: Some("markdown".to_string()),
instructions: Some("agent guidance".to_string()),
example: Some("agent example".to_string()),
schema: None,
validator: None,
};
let stage = OutputSpec {
instructions: Some("stage guidance".to_string()),
..OutputSpec::default()
};
let resolved = resolve_output_spec(Some(&agent), Some(&stage), None)
.expect("some level asked for an output");
assert_eq!(resolved.instructions.as_deref(), Some("stage guidance"));
assert_eq!(resolved.format.as_deref(), Some("markdown"));
assert_eq!(resolved.example.as_deref(), Some("agent example"));
}
#[test]
fn a_stage_alone_can_ask_for_an_output() {
let stage = spec(Some("a2ui"), None);
let resolved =
resolve_output_spec(None, Some(&stage), None).expect("the stage asked for one");
assert_eq!(resolved.format.as_deref(), Some("a2ui"));
}
#[test]
fn re_stating_the_declared_format_keeps_its_shape_checks() {
let agent = OutputSpec {
format: Some("json".to_string()),
schema: Some(json!({"type": "object"})),
validator: Some("v.rhai".to_string()),
..OutputSpec::default()
};
let request = spec(Some("json"), None);
let resolved = resolve_output_spec(Some(&agent), None, Some(&request))
.expect("the agent asked for one");
assert_eq!(resolved.schema, Some(json!({"type": "object"})));
assert_eq!(resolved.validator.as_deref(), Some("v.rhai"));
}
#[test]
fn reshaping_retires_the_validator_too() {
let agent = OutputSpec {
format: Some("a2ui".to_string()),
validator: Some("a2ui.rhai".to_string()),
..OutputSpec::default()
};
let request = spec(Some("xml"), None);
let resolved = resolve_output_spec(Some(&agent), None, Some(&request))
.expect("the agent asked for one");
assert_eq!(resolved.format.as_deref(), Some("xml"));
assert_eq!(resolved.validator, None);
}
#[test]
fn a_caller_can_supply_shape_checks_with_its_own_format() {
let agent = OutputSpec {
format: Some("a2ui".to_string()),
validator: Some("a2ui.rhai".to_string()),
..OutputSpec::default()
};
let request = OutputSpec {
format: Some("json".to_string()),
schema: Some(json!({"type": "array"})),
..OutputSpec::default()
};
let resolved = resolve_output_spec(Some(&agent), None, Some(&request))
.expect("the agent asked for one");
assert_eq!(resolved.schema, Some(json!({"type": "array"})));
assert_eq!(resolved.validator, None, "the agent's own is still retired");
}
#[test]
fn a_caller_reshaping_the_output_drops_the_declared_schema() {
let agent = spec(Some("json"), Some(json!({"type": "object"})));
let request = spec(Some("a2ui"), None);
let resolved = resolve_output_spec(Some(&agent), None, Some(&request))
.expect("the agent asked for one");
assert_eq!(resolved.format.as_deref(), Some("a2ui"));
assert_eq!(resolved.schema, None);
}
#[test]
fn a_caller_supplying_its_own_schema_keeps_it() {
let agent = spec(Some("json"), Some(json!({"type": "object"})));
let request = spec(Some("json"), Some(json!({"type": "array"})));
let resolved = resolve_output_spec(Some(&agent), None, Some(&request))
.expect("the agent asked for one");
assert_eq!(resolved.schema, Some(json!({"type": "array"})));
}
#[test]
fn a_caller_that_names_no_format_leaves_the_schema_alone() {
let agent = spec(Some("json"), Some(json!({"type": "object"})));
let request = OutputSpec {
instructions: Some("keep it short".to_string()),
..OutputSpec::default()
};
let resolved = resolve_output_spec(Some(&agent), None, Some(&request))
.expect("the agent asked for one");
assert_eq!(resolved.format.as_deref(), Some("json"));
assert_eq!(resolved.schema, Some(json!({"type": "object"})));
}
#[test]
fn short_content_is_stored_verbatim() {
let out = FinalOutput::new(
"done: 3 files",
Some("markdown".to_string()),
"wrap".into(),
7,
);
assert_eq!(out.content, "done: 3 files");
assert_eq!(out.format.as_deref(), Some("markdown"));
assert_eq!(out.stage, "wrap");
assert_eq!(out.submitted_at, 7);
assert!(!out.truncated);
}
#[test]
fn oversized_content_is_cut_at_a_char_boundary_and_flagged() {
let mut content = "a".repeat(MAX_FINAL_OUTPUT_BYTES - 1);
content.push('\u{1f600}');
let out = FinalOutput::new(&content, None, "wrap".into(), 0);
assert!(out.truncated);
assert_eq!(out.content.len(), MAX_FINAL_OUTPUT_BYTES - 1);
assert!(out.format.is_none());
}
#[test]
fn describe_spec_is_empty_when_nothing_is_constrained() {
assert_eq!(describe_spec(&OutputSpec::default()), "");
}
#[test]
fn describe_spec_renders_every_field_it_has() {
let described = describe_spec(&OutputSpec {
format: Some("a2ui".to_string()),
instructions: Some("One card per finding.".to_string()),
example: Some("{\"root\": {}}".to_string()),
schema: Some(json!({"type": "object"})),
validator: None,
});
assert!(described.contains("Return it in this format: a2ui."));
assert!(described.contains("One card per finding."));
assert!(described.contains("valid against this schema"));
assert!(described.contains("{\"root\": {}}"));
}
#[test]
fn a_constrained_spec_says_it_outranks_the_stage_prompt() {
let described = describe_spec(&OutputSpec {
instructions: Some("Reply with only the integer.".to_string()),
..OutputSpec::default()
});
assert!(
described.contains("Where anything else you were told"),
"{described}"
);
assert!(
described.trim_end().ends_with("follow this."),
"{described}"
);
}
#[test]
fn a_format_only_spec_claims_precedence_too() {
let described = describe_spec(&OutputSpec {
format: Some("text".to_string()),
..OutputSpec::default()
});
assert!(
described.contains("Where anything else you were told"),
"{described}"
);
}
#[test]
fn an_unconstrained_spec_claims_nothing() {
assert!(!describe_spec(&OutputSpec::default()).contains("follow this"));
}
#[test]
fn a_spec_round_trips_through_serde() {
let original = spec(Some("a2ui"), Some(json!({"type": "object"})));
let text = serde_json::to_string(&original).expect("a spec serializes");
let back: OutputSpec = serde_json::from_str(&text).expect("and deserializes");
assert_eq!(back, original);
assert!(!text.contains("instructions"));
}
#[test]
fn a_final_output_round_trips_through_serde() {
let original = FinalOutput::new("answer", None, "wrap".into(), 1);
let text = serde_json::to_string(&original).expect("an output serializes");
let back: FinalOutput = serde_json::from_str(&text).expect("and deserializes");
assert_eq!(back, original);
}
}