use crate::llm::structured::{self, PartialObjectCallback, StructuredMode, StructuredRequest};
use crate::llm::LlmClient;
use crate::tools::types::{Tool, ToolContext, ToolErrorKind, ToolOutput, ToolStreamEvent};
use anyhow::Result;
use async_trait::async_trait;
use serde_json::Value;
use std::sync::Arc;
const MAX_SCHEMA_BYTES: usize = 64 * 1024;
const MAX_SCHEMA_DEPTH: usize = 32;
const MAX_PROMPT_BYTES: usize = 128 * 1024;
const MAX_SYSTEM_BYTES: usize = 32 * 1024;
const DEFAULT_TIMEOUT_MS: u64 = 120_000;
const MAX_TIMEOUT_MS: u64 = 600_000;
const PARTIAL_EVENT_INTERVAL: std::time::Duration = std::time::Duration::from_millis(100);
const MAX_PARTIAL_EVENT_BYTES: usize = 64 * 1024;
pub struct GenerateObjectTool {
llm_client: Arc<dyn LlmClient>,
}
impl GenerateObjectTool {
pub fn new(llm_client: Arc<dyn LlmClient>) -> Self {
Self { llm_client }
}
}
#[async_trait]
impl Tool for GenerateObjectTool {
fn name(&self) -> &str {
"generate_object"
}
fn description(&self) -> &str {
"Generate a JSON object that strictly conforms to a provided JSON Schema. \
Use when you need structured output: extracting fields from text, classifying \
data, converting natural language to typed records, or producing machine-readable \
results. Returns the validated object on success."
}
fn parameters(&self) -> Value {
serde_json::json!({
"type": "object",
"required": ["schema", "prompt"],
"additionalProperties": false,
"properties": {
"schema": {
"type": "object",
"description": "JSON Schema that the output object must conform to"
},
"schema_name": {
"type": "string",
"description": "Short name for the schema (used internally for tool naming)",
"default": "result"
},
"schema_description": {
"type": "string",
"description": "Optional description of what the schema represents"
},
"prompt": {
"type": "string",
"description": "The prompt describing what object to generate or extract"
},
"system": {
"type": "string",
"description": "Optional system prompt to guide generation"
},
"mode": {
"type": "string",
"enum": ["auto", "strict", "json", "tool", "prompt"],
"description": "Output mode. 'auto' selects the best mode for the provider. 'tool' uses tool-calling (most reliable cross-provider). 'strict' uses OpenAI native JSON schema. 'json' uses json_object mode. 'prompt' appends schema to prompt.",
"default": "auto"
},
"max_repair_attempts": {
"type": "integer",
"description": "Maximum repair attempts if output fails validation (0-5)",
"default": 2,
"minimum": 0,
"maximum": 5
},
"include_raw_text": {
"type": "boolean",
"description": "Include the raw model text/tool arguments used to extract the final value. Defaults to false to avoid exposing reasoning-channel text.",
"default": false
},
"timeout_ms": {
"type": "integer",
"minimum": 1000,
"maximum": MAX_TIMEOUT_MS,
"description": "Independent generation deadline in milliseconds. Default 120000; maximum 600000."
}
}
})
}
async fn execute(&self, args: &Value, ctx: &ToolContext) -> Result<ToolOutput> {
let schema = match args.get("schema") {
Some(s) if s.is_object() => s.clone(),
Some(_) => {
return Ok(ToolOutput::error(
"'schema' must be a JSON object (a valid JSON Schema)",
));
}
None => {
return Ok(ToolOutput::error("'schema' parameter is required"));
}
};
let schema_bytes = serde_json::to_vec(&schema)?.len();
if schema_bytes > MAX_SCHEMA_BYTES {
return Ok(invalid_argument(format!(
"'schema' exceeds the {MAX_SCHEMA_BYTES} byte limit"
)));
}
if json_depth(&schema) > MAX_SCHEMA_DEPTH {
return Ok(invalid_argument(format!(
"'schema' exceeds the maximum nesting depth of {MAX_SCHEMA_DEPTH}"
)));
}
if let Err(error) = jsonschema::draft202012::options().build(&schema) {
return Ok(invalid_argument(format!(
"'schema' is not a valid JSON Schema: {error}"
)));
}
let prompt = match args.get("prompt").and_then(|v| v.as_str()) {
Some(p) if !p.is_empty() => p.to_string(),
_ => {
return Ok(ToolOutput::error(
"'prompt' parameter is required and must be non-empty",
));
}
};
if prompt.len() > MAX_PROMPT_BYTES {
return Ok(invalid_argument(format!(
"'prompt' exceeds the {MAX_PROMPT_BYTES} byte limit"
)));
}
if schema.get("type").is_none()
&& schema.get("properties").is_none()
&& schema.get("anyOf").is_none()
&& schema.get("oneOf").is_none()
&& schema.get("enum").is_none()
{
return Ok(ToolOutput::error(
"'schema' should contain at least one of: type, properties, anyOf, oneOf, or enum",
));
}
let schema_name: String = args
.get("schema_name")
.and_then(|v| v.as_str())
.unwrap_or("result")
.chars()
.filter(|c| c.is_alphanumeric() || *c == '_' || *c == '-')
.take(64)
.collect();
let schema_name = if schema_name.is_empty() {
"result".to_string()
} else {
schema_name
};
let schema_description = args
.get("schema_description")
.and_then(|v| v.as_str())
.map(|s| s.to_string());
let system = args
.get("system")
.and_then(|v| v.as_str())
.map(|s| s.to_string());
if system
.as_ref()
.is_some_and(|value| value.len() > MAX_SYSTEM_BYTES)
{
return Ok(invalid_argument(format!(
"'system' exceeds the {MAX_SYSTEM_BYTES} byte limit"
)));
}
let requested_mode = args.get("mode").and_then(|v| v.as_str()).unwrap_or("auto");
let mode = match requested_mode {
"strict" => StructuredMode::Strict,
"json" => StructuredMode::Json,
"tool" => StructuredMode::Tool,
"prompt" => StructuredMode::Prompt,
"auto" => StructuredMode::Auto,
other => {
return Ok(ToolOutput::error(format!(
"'mode' must be one of auto, strict, json, tool, or prompt; got '{other}'"
)));
}
};
let max_repair_attempts = args
.get("max_repair_attempts")
.and_then(|v| v.as_u64())
.unwrap_or(2)
.min(5) as u8;
let include_raw_text = args
.get("include_raw_text")
.and_then(|v| v.as_bool())
.unwrap_or(false);
let timeout_ms = args
.get("timeout_ms")
.and_then(|value| value.as_u64())
.unwrap_or(DEFAULT_TIMEOUT_MS)
.clamp(1_000, MAX_TIMEOUT_MS);
let req = StructuredRequest {
prompt,
system,
schema,
schema_name: schema_name.clone(),
schema_description,
mode,
max_repair_attempts,
};
let llm_client = ctx
.llm_client()
.unwrap_or_else(|| Arc::clone(&self.llm_client));
let cancellation = ctx.cancellation_token();
let generation = async {
if let Some(ref tx) = ctx.event_tx {
let tx_clone = tx.clone();
let last_event = Arc::new(std::sync::Mutex::new(None::<std::time::Instant>));
let callback: PartialObjectCallback = Box::new(move |partial: &Value| {
let now = std::time::Instant::now();
let mut last_event = last_event.lock().unwrap();
if last_event
.is_some_and(|last| now.duration_since(last) < PARTIAL_EVENT_INTERVAL)
{
return;
}
*last_event = Some(now);
let encoded = serde_json::to_vec(partial).unwrap_or_default();
let delta = if encoded.len() <= MAX_PARTIAL_EVENT_BYTES {
serde_json::json!({
"object_partial": partial,
"final": false,
})
} else {
serde_json::json!({
"object_partial_omitted": true,
"partial_bytes": encoded.len(),
"final": false,
})
};
let delta_str = serde_json::to_string(&delta).unwrap_or_default();
let _ = tx_clone.try_send(ToolStreamEvent::OutputDelta(delta_str));
});
structured::generate_streaming(&*llm_client, &req, callback).await
} else {
structured::generate_blocking(&*llm_client, &req).await
}
};
let result = tokio::select! {
biased;
_ = cancellation.cancelled() => Err(GenerationStop::Cancelled),
_ = tokio::time::sleep(std::time::Duration::from_millis(timeout_ms)) => Err(GenerationStop::TimedOut),
result = generation => result.map_err(GenerationStop::Failed),
};
match result {
Ok(sr) => {
if let Some(ref tx) = ctx.event_tx {
let object_bytes = serde_json::to_vec(&sr.object).unwrap_or_default().len();
let final_delta = if object_bytes <= MAX_PARTIAL_EVENT_BYTES {
serde_json::json!({
"object_partial": sr.object,
"final": true,
"mode_used": sr.mode_used,
"repair_rounds": sr.repair_rounds,
})
} else {
serde_json::json!({
"object_partial_omitted": true,
"partial_bytes": object_bytes,
"final": true,
"mode_used": sr.mode_used,
"repair_rounds": sr.repair_rounds,
})
};
let _ = tx.try_send(ToolStreamEvent::OutputDelta(
serde_json::to_string(&final_delta).unwrap_or_default(),
));
}
let mut output = serde_json::json!({
"object": sr.object,
"repair_rounds": sr.repair_rounds,
"mode_used": sr.mode_used,
"usage": {
"prompt_tokens": sr.usage.prompt_tokens,
"completion_tokens": sr.usage.completion_tokens,
"total_tokens": sr.usage.total_tokens,
"cache_read_tokens": sr.usage.cache_read_tokens,
"cache_write_tokens": sr.usage.cache_write_tokens,
}
});
if include_raw_text {
output["raw_text"] = sr.raw_text.map(Value::String).unwrap_or(Value::Null);
}
let metadata = serde_json::json!({
"schema_name": schema_name,
"requested_mode": requested_mode,
"mode_used": sr.mode_used,
"repair_rounds": sr.repair_rounds,
"usage": output["usage"].clone(),
"raw_text_included": include_raw_text,
});
Ok(ToolOutput::success(serde_json::to_string(&output)?).with_metadata(metadata))
}
Err(stop) => {
let (message, kind) = match stop {
GenerationStop::Cancelled => (
"generate_object cancelled by caller".to_string(),
Some(ToolErrorKind::Cancelled {
op: "generate_object".to_string(),
}),
),
GenerationStop::TimedOut => (
format!("generate_object timed out after {timeout_ms}ms"),
Some(ToolErrorKind::Timeout {
op: "generate_object".to_string(),
duration_ms: timeout_ms,
}),
),
GenerationStop::Failed(error) => {
let message = error.to_string();
let lower = message.to_ascii_lowercase();
let kind = (lower.contains("rate limit")
|| lower.contains("too many requests"))
.then_some(ToolErrorKind::RateLimited {
retry_after_ms: None,
});
(format!("generate_object failed: {message}"), kind)
}
};
let output = ToolOutput::error(message).with_metadata(serde_json::json!({
"schema_name": schema_name,
"requested_mode": requested_mode,
"mode_requested": mode,
"timeout_ms": timeout_ms,
}));
Ok(match kind {
Some(kind) => output.with_error_kind(kind),
None => output,
})
}
}
}
}
enum GenerationStop {
Cancelled,
TimedOut,
Failed(anyhow::Error),
}
fn invalid_argument(message: String) -> ToolOutput {
ToolOutput::error(&message).with_error_kind(ToolErrorKind::InvalidArgument { message })
}
fn json_depth(value: &Value) -> usize {
match value {
Value::Array(values) => 1 + values.iter().map(json_depth).max().unwrap_or(0),
Value::Object(values) => 1 + values.values().map(json_depth).max().unwrap_or(0),
_ => 1,
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::agent::{AgentConfig, AgentLoop};
use crate::budget::{BudgetDecision, BudgetGuard};
use crate::llm::structured::{NativeStructuredSupport, StructuredDirective, StructuredMode};
use crate::llm::{ContentBlock, LlmResponse, Message, StreamEvent, TokenUsage, ToolDefinition};
use crate::tools::ToolExecutor;
use async_trait::async_trait;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Mutex;
use std::time::Duration;
use tokio::sync::{mpsc, Notify};
use tokio_util::sync::CancellationToken;
struct MockObjectClient {
response: Mutex<Option<LlmResponse>>,
}
impl MockObjectClient {
fn new(response: LlmResponse) -> Self {
Self {
response: Mutex::new(Some(response)),
}
}
fn response() -> LlmResponse {
LlmResponse {
message: Message {
role: "assistant".to_string(),
content: vec![ContentBlock::ToolUse {
id: "call_1".to_string(),
name: "emit_colors".to_string(),
input: serde_json::json!({ "elements": ["red", "blue"] }),
}],
reasoning_content: None,
},
usage: TokenUsage {
prompt_tokens: 11,
completion_tokens: 7,
total_tokens: 18,
cache_read_tokens: None,
cache_write_tokens: None,
},
stop_reason: Some("tool_use".to_string()),
token_logprobs: Vec::new(),
meta: None,
}
}
}
#[async_trait]
impl LlmClient for MockObjectClient {
async fn complete(
&self,
_messages: &[Message],
_system: Option<&str>,
_tools: &[ToolDefinition],
) -> anyhow::Result<LlmResponse> {
self.response
.lock()
.unwrap()
.take()
.ok_or_else(|| anyhow::anyhow!("response already used"))
}
async fn complete_streaming(
&self,
_messages: &[Message],
_system: Option<&str>,
_tools: &[ToolDefinition],
_cancel_token: CancellationToken,
) -> anyhow::Result<mpsc::Receiver<StreamEvent>> {
anyhow::bail!("streaming is not used in this test")
}
fn native_structured_support(&self) -> NativeStructuredSupport {
NativeStructuredSupport::ForcedTool
}
async fn complete_structured(
&self,
messages: &[Message],
system: Option<&str>,
tools: &[ToolDefinition],
directive: &StructuredDirective,
) -> anyhow::Result<LlmResponse> {
assert_eq!(messages.len(), 1);
assert!(system.unwrap_or_default().contains("emit_colors"));
assert_eq!(directive.force_tool.as_deref(), Some("emit_colors"));
assert_eq!(tools[0].parameters["required"][0], "elements");
self.complete(messages, system, tools).await
}
}
struct RepairingObjectClient {
responses: Mutex<Vec<LlmResponse>>,
calls: Arc<AtomicUsize>,
}
#[async_trait]
impl LlmClient for RepairingObjectClient {
async fn complete(
&self,
_messages: &[Message],
_system: Option<&str>,
_tools: &[ToolDefinition],
) -> anyhow::Result<LlmResponse> {
self.calls.fetch_add(1, Ordering::SeqCst);
let mut responses = self.responses.lock().unwrap();
if responses.is_empty() {
anyhow::bail!("no response left")
}
Ok(responses.remove(0))
}
async fn complete_streaming(
&self,
_messages: &[Message],
_system: Option<&str>,
_tools: &[ToolDefinition],
_cancel_token: CancellationToken,
) -> anyhow::Result<mpsc::Receiver<StreamEvent>> {
anyhow::bail!("streaming is not used by repair tests")
}
fn native_structured_support(&self) -> NativeStructuredSupport {
NativeStructuredSupport::ForcedTool
}
}
#[derive(Default)]
struct GenerateObjectBudgetGuard {
checks: AtomicUsize,
records: AtomicUsize,
}
#[async_trait]
impl BudgetGuard for GenerateObjectBudgetGuard {
async fn check_before_llm(
&self,
_session_id: &str,
_estimated_prompt_tokens: usize,
) -> BudgetDecision {
self.checks.fetch_add(1, Ordering::SeqCst);
BudgetDecision::Allow
}
async fn record_after_llm(&self, _session_id: &str, _usage: &TokenUsage) {
self.records.fetch_add(1, Ordering::SeqCst);
}
}
struct BlockingObjectClient {
started: Arc<Notify>,
calls: Arc<AtomicUsize>,
}
#[async_trait]
impl LlmClient for BlockingObjectClient {
async fn complete(
&self,
_messages: &[Message],
_system: Option<&str>,
_tools: &[ToolDefinition],
) -> anyhow::Result<LlmResponse> {
self.calls.fetch_add(1, Ordering::SeqCst);
self.started.notify_one();
std::future::pending::<anyhow::Result<LlmResponse>>().await
}
async fn complete_streaming(
&self,
_messages: &[Message],
_system: Option<&str>,
_tools: &[ToolDefinition],
_cancel_token: CancellationToken,
) -> anyhow::Result<mpsc::Receiver<StreamEvent>> {
anyhow::bail!("streaming is not used by cancellation tests")
}
fn native_structured_support(&self) -> NativeStructuredSupport {
NativeStructuredSupport::ForcedTool
}
}
fn object_tool_response(input: Value) -> LlmResponse {
LlmResponse {
message: Message {
role: "assistant".to_string(),
content: vec![ContentBlock::ToolUse {
id: "call".to_string(),
name: "emit_result".to_string(),
input,
}],
reasoning_content: None,
},
usage: TokenUsage {
prompt_tokens: 3,
completion_tokens: 2,
total_tokens: 5,
cache_read_tokens: None,
cache_write_tokens: None,
},
stop_reason: Some("tool_use".to_string()),
token_logprobs: Vec::new(),
meta: None,
}
}
#[tokio::test]
async fn generate_object_tool_unwraps_array_schema_and_sets_metadata() {
let tool = GenerateObjectTool::new(Arc::new(MockObjectClient::new(
MockObjectClient::response(),
)));
let temp = tempfile::tempdir().unwrap();
let ctx = ToolContext::new(temp.path().to_path_buf());
let output = tool
.execute(
&serde_json::json!({
"schema_name": "colors",
"schema": {
"type": "array",
"items": { "type": "string" },
"minItems": 2
},
"prompt": "Return two colors",
"mode": "tool"
}),
&ctx,
)
.await
.unwrap();
assert!(output.success);
let content: Value = serde_json::from_str(&output.content).unwrap();
assert_eq!(content["object"], serde_json::json!(["red", "blue"]));
assert_eq!(
content["mode_used"],
serde_json::json!(StructuredMode::Tool)
);
assert_eq!(content["usage"]["total_tokens"], 18);
let metadata = output.metadata.unwrap();
assert_eq!(metadata["schema_name"], "colors");
assert_eq!(metadata["requested_mode"], "tool");
assert_eq!(metadata["raw_text_included"], false);
}
#[tokio::test]
async fn generate_object_repairs_use_the_tool_context_llm_budget_scope() {
let temp = tempfile::tempdir().unwrap();
let calls = Arc::new(AtomicUsize::new(0));
let raw_client: Arc<dyn LlmClient> = Arc::new(RepairingObjectClient {
responses: Mutex::new(vec![
object_tool_response(serde_json::json!({})),
object_tool_response(serde_json::json!({"value": "ok"})),
]),
calls: Arc::clone(&calls),
});
let guard = Arc::new(GenerateObjectBudgetGuard::default());
let agent = AgentLoop::new(
Arc::clone(&raw_client),
Arc::new(ToolExecutor::new(temp.path().to_string_lossy().to_string())),
ToolContext::new(temp.path().to_path_buf()),
AgentConfig {
budget_guard: Some(Arc::clone(&guard) as Arc<dyn BudgetGuard>),
..Default::default()
},
);
let cancellation = CancellationToken::new();
let event_tx = None;
let governed =
agent.scoped_llm_client_for_parts(Some("generate-session"), &event_tx, &cancellation);
let ctx = ToolContext::new(temp.path().to_path_buf())
.with_session_id("generate-session")
.with_cancellation(cancellation)
.with_llm_client(governed);
let tool = GenerateObjectTool::new(raw_client);
let output = tool
.execute(
&serde_json::json!({
"schema": {
"type": "object",
"properties": {"value": {"type": "string"}},
"required": ["value"]
},
"prompt": "Return a value",
"mode": "tool",
"max_repair_attempts": 1
}),
&ctx,
)
.await
.unwrap();
assert!(
output.success,
"generate_object should repair: {}",
output.content
);
assert_eq!(calls.load(Ordering::SeqCst), 2);
assert_eq!(guard.checks.load(Ordering::SeqCst), 2);
assert_eq!(guard.records.load(Ordering::SeqCst), 2);
}
#[tokio::test]
async fn generate_object_stops_on_tool_context_cancellation() {
let temp = tempfile::tempdir().unwrap();
let started = Arc::new(Notify::new());
let calls = Arc::new(AtomicUsize::new(0));
let tool = GenerateObjectTool::new(Arc::new(BlockingObjectClient {
started: Arc::clone(&started),
calls: Arc::clone(&calls),
}));
let cancellation = CancellationToken::new();
let ctx =
ToolContext::new(temp.path().to_path_buf()).with_cancellation(cancellation.clone());
let started_wait = started.notified();
let run = tokio::spawn(async move {
tool.execute(
&serde_json::json!({
"schema": {
"type": "object",
"properties": {"value": {"type": "string"}},
"required": ["value"]
},
"prompt": "Wait forever",
"mode": "tool"
}),
&ctx,
)
.await
});
tokio::time::timeout(Duration::from_secs(1), started_wait)
.await
.expect("structured provider call should start");
cancellation.cancel();
let output = tokio::time::timeout(Duration::from_secs(1), run)
.await
.expect("cancellation must stop structured generation")
.expect("generate_object join should succeed")
.expect("generate_object should return a typed failed output");
assert!(!output.success);
assert!(output.content.contains("cancelled"));
assert_eq!(calls.load(Ordering::SeqCst), 1);
}
}