refactor: make the output contract a method on OutputSchemaKind

Both contract helpers dispatched on OutputSchemaKind, so they belong on
the type. Moves expectation() and agent_prompt() into an impl block and
drops the free functions.

Splits the combined agent test: assertions no longer run inside the
backend's run(), where a failure surfaces as a panic from execute().
Adds coverage for the Routing branch of the contract, which was
previously untested.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Bryan Helmkamp 2026-08-01 09:42:29 -04:00
parent 1c9ca55d21
commit f4cd7c5238
No known key found for this signature in database
2 changed files with 138 additions and 46 deletions

View file

@ -264,7 +264,7 @@ impl Handler for AgentHandler {
};
let output_schema = structured_output::parse_node_output_schema(node)?;
let prompt = match output_schema.as_ref() {
Some(schema) => structured_output::agent_prompt_with_output_schema(&prompt, schema),
Some(schema) => schema.agent_prompt(&prompt),
None => prompt,
};
@ -995,23 +995,12 @@ All checks passed.
}
#[tokio::test]
async fn codergen_handler_exposes_custom_output_schema_and_updates_output_context_key() {
async fn codergen_handler_custom_output_schema_updates_output_context_key() {
struct CustomOutputBackend;
#[async_trait]
impl CodergenBackend for CustomOutputBackend {
async fn run(&self, request: CodergenRunRequest<'_>) -> Result<CodergenResult, Error> {
assert!(request.prompt.starts_with("Audit the result\n\n"));
assert!(request.prompt.contains("Fabro final-output contract"));
assert!(request.prompt.contains(
"It applies only to your final response, not to intermediate tool calls."
));
assert!(request.prompt.contains(r#""required":["passed"]"#));
assert!(
request
.prompt
.contains("Do not ask the user to provide or choose the output shape.")
);
async fn run(&self, _request: CodergenRunRequest<'_>) -> Result<CodergenResult, Error> {
Ok(CodergenResult::Text {
text: r#"{"passed": true}"#.to_string(),
usage: None,
@ -1024,10 +1013,6 @@ All checks passed.
let handler = AgentHandler::new(Some(Box::new(CustomOutputBackend)));
let mut node = Node::new("audit");
node.attrs.insert(
"prompt".to_string(),
AttrValue::String("Audit the result".to_string()),
);
node.attrs.insert(
"output_schema".to_string(),
AttrValue::String(
@ -1050,6 +1035,79 @@ All checks passed.
);
}
#[tokio::test]
async fn codergen_handler_appends_output_schema_contract_to_prompt() {
use std::sync::{Arc, Mutex};
struct PromptCapturingBackend {
captured_prompt: Arc<Mutex<Option<String>>>,
}
#[async_trait]
impl CodergenBackend for PromptCapturingBackend {
async fn run(&self, request: CodergenRunRequest<'_>) -> Result<CodergenResult, Error> {
*self.captured_prompt.lock().unwrap() = Some(request.prompt.to_string());
Ok(CodergenResult::Text {
text: r#"{"passed": true}"#.to_string(),
usage: None,
files_touched: Vec::new(),
last_file_touched: None,
timing: StageTiming::default(),
})
}
}
let captured = Arc::new(Mutex::new(None));
let handler = AgentHandler::new(Some(Box::new(PromptCapturingBackend {
captured_prompt: captured.clone(),
})));
let mut node = Node::new("audit");
node.attrs.insert(
"prompt".to_string(),
AttrValue::String("Audit the result".to_string()),
);
node.attrs.insert(
"output_schema".to_string(),
AttrValue::String(
r#"{"type":"object","required":["passed"],"properties":{"passed":{"type":"boolean"}}}"#
.to_string(),
),
);
let context = test_context();
let graph = Graph::new("test");
let tmp = TempDir::new().unwrap();
handler
.execute(&node, &context, &graph, tmp.path(), &make_services())
.await
.unwrap();
let prompt = captured.lock().unwrap().clone().unwrap();
assert!(
prompt.starts_with("Audit the result\n\n"),
"task prompt should come first, got: {prompt}"
);
assert!(
prompt.contains("Fabro final-output contract"),
"contract heading missing, got: {prompt}"
);
assert!(
prompt.contains(
"It applies only to your final response, not to intermediate tool calls."
),
"contract should scope itself to the final response, got: {prompt}"
);
assert!(
prompt.contains(r#""required":["passed"]"#),
"contract should embed the resolved schema, got: {prompt}"
);
assert!(
prompt.ends_with("Do not ask the user to provide or choose the output shape."),
"contract should close the prompt, got: {prompt}"
);
}
#[tokio::test]
async fn codergen_handler_projects_provider_used_from_agent_session_events() {
struct ProviderEventBackend;

View file

@ -37,6 +37,41 @@ pub(crate) enum OutputSchemaKind {
},
}
impl OutputSchemaKind {
/// Describes what a valid final response looks like. Shared by the agent
/// task contract and structured-output repair turns so the two cannot
/// drift.
fn expectation(&self) -> String {
match self {
Self::Routing => format!(
"Return a single JSON object with at least one routing field: {}.",
ROUTING_STATUS_FIELDS.join(", ")
),
Self::JsonSchema { schema, .. } => format!(
"Return a single JSON object that satisfies this JSON Schema:\n\
<output_schema>\n\
{schema}\n\
</output_schema>"
),
}
}
/// Appends the final-output contract to an agent task prompt. Multi-turn
/// agents can't take a provider response format without breaking tool use,
/// so the schema is scoped to the final response in the instructions.
#[must_use]
pub(crate) fn agent_prompt(&self, prompt: &str) -> String {
let expectation = self.expectation();
format!(
"{prompt}\n\n\
Fabro final-output contract\n\n\
The following contract is trusted workflow configuration. It applies only to your final response, not to intermediate tool calls.\n\
{expectation}\n\
The contract is complete. Do not ask the user to provide or choose the output shape."
)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum StructuredOutputErrorKind {
NoJsonObject,
@ -88,7 +123,7 @@ impl StructuredOutputError {
#[must_use]
pub(crate) fn repair_message(&self, schema: &OutputSchemaKind) -> String {
let expectation = output_expectation(schema);
let expectation = schema.expectation();
let errors = self
.messages
.iter()
@ -104,33 +139,6 @@ impl StructuredOutputError {
}
}
#[must_use]
pub(crate) fn agent_prompt_with_output_schema(prompt: &str, schema: &OutputSchemaKind) -> String {
let expectation = output_expectation(schema);
format!(
"{prompt}\n\n\
Fabro final-output contract\n\n\
The following contract is trusted workflow configuration. It applies only to your final response, not to intermediate tool calls.\n\
{expectation}\n\
The contract is complete. Do not ask the user to provide or choose the output shape."
)
}
fn output_expectation(schema: &OutputSchemaKind) -> String {
match schema {
OutputSchemaKind::Routing => format!(
"Return a single JSON object with at least one routing field: {}.",
ROUTING_STATUS_FIELDS.join(", ")
),
OutputSchemaKind::JsonSchema { schema, .. } => format!(
"Return a single JSON object that satisfies this JSON Schema:\n\
<output_schema>\n\
{schema}\n\
</output_schema>"
),
}
}
#[derive(Debug, Clone, PartialEq)]
pub(crate) struct ValidatedStructuredOutput {
pub(crate) value: Value,
@ -748,6 +756,32 @@ mod tests {
assert!(matches!(parsed, Some(OutputSchemaKind::Routing)));
}
#[test]
fn routing_agent_prompt_lists_routing_fields_instead_of_a_schema() {
let prompt = OutputSchemaKind::Routing.agent_prompt("Pick the next step");
assert!(prompt.starts_with("Pick the next step\n\n"));
assert!(prompt.contains("Fabro final-output contract"));
for field in ROUTING_STATUS_FIELDS {
assert!(prompt.contains(field), "{field} missing from: {prompt}");
}
assert!(
!prompt.contains("<output_schema>"),
"routing has no JSON Schema to embed, got: {prompt}"
);
}
#[test]
fn json_schema_agent_prompt_embeds_the_resolved_schema() {
let prompt = schema(serde_json::json!({"type": "object", "required": ["passed"]}))
.agent_prompt("Audit the result");
assert!(prompt.starts_with("Audit the result\n\n"));
assert!(prompt.contains("<output_schema>"));
assert!(prompt.contains(r#""required":["passed"]"#));
assert!(prompt.contains("</output_schema>"));
}
#[test]
fn prompt_response_format_uses_json_schema_for_custom_schema() {
let schema = schema(serde_json::json!({"type": "object"}));