diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 6d612452b..34800c98a 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -2,11 +2,14 @@ # SPDX-License-Identifier: Apache-2.0 # Default ownership for repository files. -* @nvidia/nat-developers +* @nvidia/nemo-relay-developers + +# CODEOWNERS changes require repository administrator review. +/.github/CODEOWNERS @nvidia/nemo-relay-admins # Attribution files require dependency approver review due to change in version(s) -ATTRIBUTIONS-*.md @nvidia/nat-dep-approvers +ATTRIBUTIONS-*.md @nvidia/nemo-relay-dep-approvers -# Documentation changes must go to technical writers and NAT developers. -docs/ @lvojtku @nvidia/NAT-developers -fern/ @lvojtku @nvidia/NAT-developers +# Documentation changes must go to technical writers and NeMo Relay developers. +docs/ @nvidia/nemo-relay-docs-reviewers @nvidia/nemo-relay-developers +fern/ @nvidia/nemo-relay-docs-reviewers @nvidia/nemo-relay-developers diff --git a/README.md b/README.md index 5a9f5facc..54ffef8b8 100644 --- a/README.md +++ b/README.md @@ -21,6 +21,9 @@ requiring changes to the existing agent stack. It gives coding agents, applications, framework integrations, middleware, and observability backends a shared runtime for scopes, policy, plugins, and lifecycle events. +For how Relay complements OpenTelemetry GenAI conventions and observability or +evaluation products, see [the Ecosystem guide](https://docs.nvidia.com/nemo/relay/about-nemo-relay/ecosystem). + ## Where To Start | Goal | Start With | @@ -30,7 +33,7 @@ shared runtime for scopes, policy, plugins, and lifecycle events. | Use LangChain, LangGraph, Deep Agents, or OpenClaw | [Supported Integrations](https://docs.nvidia.com/nemo/relay/supported-integrations/about) | | Build a framework or provider integration | [Integrate into Frameworks](https://docs.nvidia.com/nemo/relay/integrate-into-frameworks/about) | | Export ATOF, ATIF, OpenTelemetry, or OpenInference | [Observability Plugin](https://docs.nvidia.com/nemo/relay/configure-plugins/observability/about) | -| Package reusable middleware or exporters | [Build Plugins](https://docs.nvidia.com/nemo/relay/v0.5.0/build-plugins/about) | +| Package reusable middleware or exporters | [Build Plugins](https://docs.nvidia.com/nemo/relay/build-plugins/about) | | Develop or test this repository from source | [CONTRIBUTING.md](CONTRIBUTING.md) | @@ -211,7 +214,7 @@ uv add nemo-relay # Node.js # Requires Node.js 24 or newer. -npm install nemo-relay-node +npm install nemo-relay-node@0.6.0 # Rust cargo add nemo-relay @@ -335,8 +338,8 @@ End-user documentation lives at Important local entry points: - [Overview](https://docs.nvidia.com/nemo/relay/about-nemo-relay/overview) +- [Getting Started](https://docs.nvidia.com/nemo/relay/getting-started/quick-start) - [Installation](https://docs.nvidia.com/nemo/relay/getting-started/installation) -- [Agent Runtime Primer](https://docs.nvidia.com/nemo/relay/getting-started/agent-runtime-primer) - [Testing and Docs](https://docs.nvidia.com/nemo/relay/contribute/testing-and-docs) For source builds, tests, and contribution workflow, refer to diff --git a/crates/adaptive/src/acg/ir_builder.rs b/crates/adaptive/src/acg/ir_builder.rs index a40702d63..8c08e46f6 100644 --- a/crates/adaptive/src/acg/ir_builder.rs +++ b/crates/adaptive/src/acg/ir_builder.rs @@ -37,6 +37,16 @@ pub fn build_prompt_ir(request: &AnnotatedLlmRequest) -> Result { let mut sequence_index: u32 = 0; let mut inserted_tool_blocks = false; + if let Some(instructions) = &request.instructions { + blocks.push(build_text_block( + &mut sequence_index, + instructions, + PromptRole::System, + ProvenanceLabel::System, + Some("instructions"), + )); + } + for message in &request.messages { if should_insert_tool_blocks_before_message(inserted_tool_blocks, request, message) { append_tool_schema_blocks(&mut blocks, &mut sequence_index, request.tools.as_deref())?; @@ -71,7 +81,9 @@ fn should_insert_tool_blocks_before_message( request: &AnnotatedLlmRequest, message: &Message, ) -> bool { - !inserted_tool_blocks && !matches!(message, Message::System { .. }) && request.tools.is_some() + !inserted_tool_blocks + && !matches!(message, Message::System { .. } | Message::Developer { .. }) + && request.tools.is_some() } fn append_message_blocks( @@ -87,6 +99,13 @@ fn append_message_blocks( ProvenanceLabel::System, None, )), + Message::Developer { content, .. } => blocks.push(build_text_block( + sequence_index, + content, + PromptRole::System, + ProvenanceLabel::Developer, + Some("developer"), + )), Message::User { content, .. } => blocks.push(build_text_block( sequence_index, content, @@ -112,6 +131,37 @@ fn append_message_blocks( content, tool_call_id, )), + Message::Function { content, name } => { + let content = MessageContent::Text(content.clone().unwrap_or_default()); + blocks.push(build_text_block( + sequence_index, + &content, + PromptRole::Tool, + ProvenanceLabel::Tool, + Some(name), + )); + } + Message::ToolCallItem { name, .. } => blocks.push(build_serialized_message_block( + sequence_index, + message, + PromptRole::Assistant, + ProvenanceLabel::Developer, + Some(name), + )?), + Message::ToolResultItem { call_id, .. } => blocks.push(build_serialized_message_block( + sequence_index, + message, + PromptRole::Tool, + ProvenanceLabel::Tool, + Some(call_id), + )?), + Message::ProviderNative { kind, .. } => blocks.push(build_serialized_message_block( + sequence_index, + message, + PromptRole::User, + ProvenanceLabel::Developer, + Some(kind), + )?), } Ok(()) @@ -147,15 +197,36 @@ fn extract_text(content: &MessageContent) -> String { MessageContent::Text(text) => text.clone(), MessageContent::Parts(parts) => parts .iter() - .filter_map(|part| match part { - ContentPart::Text { text } => Some(text.as_str()), - ContentPart::ImageUrl { .. } => None, + .map(|part| match part { + ContentPart::Text { text, .. } => text.clone(), + ContentPart::Refusal { refusal, .. } => refusal.clone(), + ContentPart::ImageUrl { .. } | ContentPart::Image { .. } => String::new(), + ContentPart::Audio { .. } + | ContentPart::File { .. } + | ContentPart::ToolUse { .. } + | ContentPart::ToolResult { .. } + | ContentPart::ProviderNative { .. } => { + serde_json::to_string(part).unwrap_or_default() + } }) + .filter(|text| !text.is_empty()) .collect::>() .join("\n"), } } +fn build_serialized_message_block( + seq: &mut u32, + message: &Message, + role: PromptRole, + provenance: ProvenanceLabel, + suffix: Option<&str>, +) -> Result { + let value = serde_json::to_value(message)?; + let content = MessageContent::Text(canonicalize_value(&value)?); + Ok(build_text_block(seq, &content, role, provenance, suffix)) +} + fn generate_span_id(role: PromptRole, index: u32, suffix: Option<&str>) -> SpanId { let role_str = match role { PromptRole::System => "system", @@ -257,7 +328,7 @@ fn build_tool_schema_block(seq: &mut u32, tool: &ToolDefinition) -> Result Result &str { + match tool { + ToolDefinition::Function { function, .. } => &function.name, + ToolDefinition::ProviderNative { kind, .. } => kind, + } +} + fn compute_request_hash(request: &AnnotatedLlmRequest) -> Result { let value = serde_json::to_value(request)?; let canonical = canonicalize_value(&value)?; diff --git a/crates/adaptive/src/acg_profile.rs b/crates/adaptive/src/acg_profile.rs index 01e5c3b29..fa3f47c0d 100644 --- a/crates/adaptive/src/acg_profile.rs +++ b/crates/adaptive/src/acg_profile.rs @@ -76,26 +76,47 @@ fn derive_key_parts(annotated_request: &AnnotatedLlmRequest) -> AcgKeyParts<'_> fn message_role_tag(message: &Message) -> &'static str { match message { Message::System { .. } => "system", + Message::Developer { .. } => "developer", Message::User { .. } => "user", Message::Assistant { .. } => "assistant", Message::Tool { .. } => "tool", + Message::Function { .. } => "function", + Message::ToolCallItem { .. } => "tool_call", + Message::ToolResultItem { .. } => "tool_result", + Message::ProviderNative { .. } => "provider_native", } } fn system_prompt_fingerprint(annotated_request: &AnnotatedLlmRequest) -> String { - let system_content = annotated_request - .messages - .iter() - .filter_map(|message| match message { - Message::System { content, .. } => Some(extract_text(content)), - _ => None, - }) - .collect::>() - .join("\n"); + let mut system_content = Vec::new(); + if let Some(instructions) = &annotated_request.instructions { + system_content.push(serde_json::json!({ + "source": "instructions", + "content": instructions, + })); + } + system_content.extend( + annotated_request + .messages + .iter() + .filter_map(|message| match message { + Message::System { content, .. } => Some(serde_json::json!({ + "source": "message", + "role": "system", + "content": content, + })), + Message::Developer { content, .. } => Some(serde_json::json!({ + "source": "message", + "role": "developer", + "content": content, + })), + _ => None, + }), + ); if system_content.is_empty() { "no-system".to_string() } else { - sha256_hex(&system_content) + hash_canonical_json(&serde_json::Value::Array(system_content)) } } @@ -148,7 +169,7 @@ fn learning_seed_fingerprint(annotated_request: &AnnotatedLlmRequest) -> String .messages .iter() .find_map(|message| match message { - Message::System { .. } => None, + Message::System { .. } | Message::Developer { .. } => None, Message::User { content, .. } => { Some(format!("user:{}", sha256_hex(&extract_text(content)))) } @@ -160,10 +181,46 @@ fn learning_seed_fingerprint(annotated_request: &AnnotatedLlmRequest) -> String Message::Tool { content, .. } => { Some(format!("tool:{}", sha256_hex(&extract_text(content)))) } + Message::Function { content, name } => Some(format!( + "function:{}", + hash_canonical_json(&serde_json::json!({ + "name": name, + "content": content, + })) + )), + Message::ToolCallItem { + name, arguments, .. + } => Some(format!( + "tool-call:{}", + hash_canonical_json(&serde_json::json!({ + "name": name, + "arguments": arguments, + })) + )), + Message::ToolResultItem { output, .. } => { + Some(format!("tool-result:{}", sha256_hex(&output.to_string()))) + } + Message::ProviderNative { + provider, + kind, + value, + } => Some(format!( + "native:{}", + hash_canonical_json(&serde_json::json!({ + "provider": provider, + "kind": kind, + "value": value, + })) + )), }) .unwrap_or_else(|| "no-seed".to_string()) } +fn hash_canonical_json(value: &serde_json::Value) -> String { + let canonical = canonicalize_value(value).unwrap_or_else(|_| value.to_string()); + sha256_hex(&canonical) +} + fn tool_schema_fingerprint(tools: Option<&[ToolDefinition]>) -> String { let Some(tools) = tools else { return "no-tools".to_string(); @@ -189,12 +246,21 @@ fn extract_text(content: &MessageContent) -> String { MessageContent::Parts(parts) => parts .iter() .map(|part| match part { - ContentPart::Text { text } => text.clone(), - ContentPart::ImageUrl { image_url } => format!( + ContentPart::Text { text, .. } => text.clone(), + ContentPart::Refusal { refusal, .. } => refusal.clone(), + ContentPart::ImageUrl { image_url, .. } => format!( "[image:{}:{}]", image_url.detail.as_deref().unwrap_or("none"), sha256_hex(&image_url.url) ), + ContentPart::Image { .. } + | ContentPart::Audio { .. } + | ContentPart::File { .. } + | ContentPart::ToolUse { .. } + | ContentPart::ToolResult { .. } + | ContentPart::ProviderNative { .. } => { + serde_json::to_string(part).unwrap_or_default() + } }) .collect::>() .join("\n"), diff --git a/crates/adaptive/tests/integration/acg_module_surface_tests.rs b/crates/adaptive/tests/integration/acg_module_surface_tests.rs index 16bf0dd94..374c14189 100644 --- a/crates/adaptive/tests/integration/acg_module_surface_tests.rs +++ b/crates/adaptive/tests/integration/acg_module_surface_tests.rs @@ -129,6 +129,8 @@ fn acg_module_surface_policy_and_ir_builder_symbols_compile_from_canonical_names assert!(envelope.policy.warm_first_enabled); let request = AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![ Message::System { content: MessageContent::Text("You are helpful.".to_string()), @@ -175,6 +177,8 @@ fn acg_module_surface_build_prompt_ir_inserts_tool_schema_before_first_non_syste use nemo_relay_adaptive::acg::prompt_ir::{BlockContentType, PromptRole}; let request = AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![ Message::System { content: MessageContent::Text("You are helpful.".to_string()), @@ -187,8 +191,7 @@ fn acg_module_surface_build_prompt_ir_inserts_tool_schema_before_first_non_syste ], model: Some("claude-sonnet".to_string()), params: None, - tools: Some(vec![ToolDefinition { - tool_type: "function".to_string(), + tools: Some(vec![ToolDefinition::Function { function: FunctionDefinition { name: "get_weather".to_string(), description: Some("Look up the weather".to_string()), @@ -198,7 +201,10 @@ fn acg_module_surface_build_prompt_ir_inserts_tool_schema_before_first_non_syste "location": {"type": "string"} } })), + strict: None, + extra: serde_json::Map::new(), }, + extra: serde_json::Map::new(), }]), tool_choice: None, store: None, diff --git a/crates/adaptive/tests/integration/redis_tests.rs b/crates/adaptive/tests/integration/redis_tests.rs index 9cd730ae6..8e2eca4f4 100644 --- a/crates/adaptive/tests/integration/redis_tests.rs +++ b/crates/adaptive/tests/integration/redis_tests.rs @@ -118,6 +118,8 @@ fn make_test_run(agent_id: &str) -> RunRecord { fn sample_annotated_request(model: &str) -> AnnotatedLlmRequest { AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![ Message::System { content: MessageContent::Text("You are a careful planner".to_string()), diff --git a/crates/adaptive/tests/integration/runtime_integration_tests.rs b/crates/adaptive/tests/integration/runtime_integration_tests.rs index 16da0444f..dd65079ea 100644 --- a/crates/adaptive/tests/integration/runtime_integration_tests.rs +++ b/crates/adaptive/tests/integration/runtime_integration_tests.rs @@ -78,6 +78,8 @@ fn reset_global() { fn sample_annotated_request(model: &str) -> AnnotatedLlmRequest { AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![ Message::System { content: MessageContent::Text("You are a careful planner".to_string()), @@ -112,6 +114,8 @@ fn sample_annotated_request(model: &str) -> AnnotatedLlmRequest { fn sample_growing_chat_requests(model: &str) -> Vec { vec![ AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![ Message::System { content: MessageContent::Text("You are a careful planner".to_string()), @@ -142,6 +146,8 @@ fn sample_growing_chat_requests(model: &str) -> Vec { extra: Map::new(), }, AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![ Message::System { content: MessageContent::Text("You are a careful planner".to_string()), @@ -183,6 +189,8 @@ fn sample_growing_chat_requests(model: &str) -> Vec { extra: Map::new(), }, AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![ Message::System { content: MessageContent::Text("You are a careful planner".to_string()), @@ -517,14 +525,11 @@ async fn runtime_integration_acg_learner_reuses_learning_buckets_across_growing_ let requests = sample_growing_chat_requests("claude-3-5-sonnet"); let learner = AcgLearner::new(agent_id, 8, StabilityThresholds::default()); let learning_key = format!( - "{agent_id}::model=claude-3-5-sonnet::seed={}::system={}::tools=no-tools", + "{agent_id}::model=claude-3-5-sonnet::seed={}::system=sha256:3087d8fd4::tools=no-tools", short_hash(&format!( "user:{}", nemo_relay_adaptive::acg::sha256_hex("Summarize the latest findings") )), - short_hash(&nemo_relay_adaptive::acg::sha256_hex( - "You are a careful planner" - )), ); learner diff --git a/crates/adaptive/tests/unit/acg/ir_builder_tests.rs b/crates/adaptive/tests/unit/acg/ir_builder_tests.rs index 2fee6ba62..9d1ba9934 100644 --- a/crates/adaptive/tests/unit/acg/ir_builder_tests.rs +++ b/crates/adaptive/tests/unit/acg/ir_builder_tests.rs @@ -12,8 +12,7 @@ use super::super::ir_builder::build_prompt_ir; use crate::acg::prompt_ir::{BlockContentType, PromptRole, ProvenanceLabel}; fn sample_tool_definition(name: &str) -> ToolDefinition { - ToolDefinition { - tool_type: "function".to_string(), + ToolDefinition::Function { function: FunctionDefinition { name: name.to_string(), description: Some(format!("describe {name}")), @@ -23,7 +22,10 @@ fn sample_tool_definition(name: &str) -> ToolDefinition { "query": {"type": "string"} } })), + strict: None, + extra: serde_json::Map::new(), }, + extra: serde_json::Map::new(), } } @@ -41,6 +43,8 @@ fn sample_tool_call(name: &str) -> ToolCall { #[test] fn build_prompt_ir_inserts_tools_before_first_non_system_message_and_preserves_all_message_kinds() { let request = AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![ Message::System { content: MessageContent::Text("You are helpful.".to_string()), @@ -50,9 +54,11 @@ fn build_prompt_ir_inserts_tools_before_first_non_system_message_and_preserves_a content: MessageContent::Parts(vec![ ContentPart::Text { text: "Hello".to_string(), + extra: serde_json::Map::new(), }, ContentPart::Text { text: "World".to_string(), + extra: serde_json::Map::new(), }, ]), name: None, @@ -112,6 +118,8 @@ fn build_prompt_ir_inserts_tools_before_first_non_system_message_and_preserves_a #[test] fn build_prompt_ir_appends_tool_blocks_when_request_contains_only_system_messages() { let request = AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![Message::System { content: MessageContent::Text("System only".to_string()), name: None, @@ -157,6 +165,8 @@ fn build_prompt_ir_appends_tool_blocks_when_request_contains_only_system_message #[test] fn build_prompt_ir_omits_tool_schema_hashes_when_no_tools_are_present() { let request = AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![Message::User { content: MessageContent::Text("No tools".to_string()), name: None, @@ -187,3 +197,104 @@ fn build_prompt_ir_omits_tool_schema_hashes_when_no_tools_are_present() { assert!(prompt_ir.tool_schema_hashes.is_none()); assert_eq!(prompt_ir.blocks[0].span_id.0, "user-0"); } + +#[test] +fn build_prompt_ir_covers_extended_request_messages_and_content_parts() { + let request = AnnotatedLlmRequest { + instructions: Some(MessageContent::Text("Top-level instructions".into())), + messages: vec![ + Message::Developer { + content: MessageContent::Text("Developer guide".into()), + name: Some("developer".into()), + }, + Message::User { + content: MessageContent::Parts(vec![ + ContentPart::Refusal { + refusal: "refusal".into(), + extra: serde_json::Map::new(), + }, + ContentPart::Image { + image: serde_json::json!({"file_id": "image_1"}), + extra: serde_json::Map::new(), + }, + ContentPart::Audio { + audio: serde_json::json!({"data": "audio"}), + extra: serde_json::Map::new(), + }, + ContentPart::File { + file: serde_json::json!({"file_id": "file_1"}), + extra: serde_json::Map::new(), + }, + ContentPart::ToolUse { + id: "call_1".into(), + name: "lookup".into(), + input: serde_json::json!({"q": "x"}), + extra: serde_json::Map::new(), + }, + ContentPart::ToolResult { + tool_use_id: "call_1".into(), + content: serde_json::json!("ok"), + is_error: Some(false), + extra: serde_json::Map::new(), + }, + ContentPart::ProviderNative { + provider: "openai_responses".into(), + kind: "future".into(), + value: serde_json::json!({"type": "future"}), + }, + ]), + name: None, + }, + Message::Function { + content: None, + name: "legacy".into(), + }, + Message::ToolCallItem { + id: Some("fc_1".into()), + call_id: "call_1".into(), + name: "lookup".into(), + arguments: serde_json::json!({"q": "x"}), + extra: serde_json::Map::new(), + }, + Message::ToolResultItem { + id: Some("fco_1".into()), + call_id: "call_1".into(), + output: serde_json::json!({"ok": true}), + extra: serde_json::Map::new(), + }, + Message::ProviderNative { + provider: "openai_responses".into(), + kind: "reasoning".into(), + value: serde_json::json!({"type": "reasoning", "summary": []}), + }, + ], + tools: Some(vec![ToolDefinition::ProviderNative { + provider: "openai_responses".into(), + kind: "web_search_preview".into(), + value: serde_json::json!({"type": "web_search_preview"}), + }]), + model: Some("gpt-5".into()), + ..AnnotatedLlmRequest::default() + }; + + let prompt_ir = build_prompt_ir(&request).unwrap(); + assert_eq!(prompt_ir.blocks[0].span_id.0, "system-0-instructions"); + assert_eq!(prompt_ir.blocks[1].span_id.0, "system-1-developer"); + assert_eq!(prompt_ir.blocks[2].span_id.0, "system-2-web_search_preview"); + assert!( + prompt_ir + .blocks + .iter() + .any(|block| block.content_type == BlockContentType::ToolSchema) + ); + assert!( + prompt_ir + .blocks + .iter() + .any(|block| block.span_id.0.ends_with("reasoning")) + ); + assert_eq!( + prompt_ir.tool_schema_hashes.as_ref().unwrap()[0].tool_name, + "web_search_preview" + ); +} diff --git a/crates/adaptive/tests/unit/acg_component_tests.rs b/crates/adaptive/tests/unit/acg_component_tests.rs index fcf09c01a..51477bb3b 100644 --- a/crates/adaptive/tests/unit/acg_component_tests.rs +++ b/crates/adaptive/tests/unit/acg_component_tests.rs @@ -83,6 +83,8 @@ fn sample_openai_chat_request() -> LlmRequest { fn sample_annotated_request(model: &str) -> AnnotatedLlmRequest { AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![ Message::System { content: MessageContent::Text("You are helpful.".to_string()), diff --git a/crates/adaptive/tests/unit/acg_learner_tests.rs b/crates/adaptive/tests/unit/acg_learner_tests.rs index 664ff5126..de23483e2 100644 --- a/crates/adaptive/tests/unit/acg_learner_tests.rs +++ b/crates/adaptive/tests/unit/acg_learner_tests.rs @@ -20,6 +20,8 @@ use crate::types::records::{CallRecord, RunRecord}; fn sample_request(model: &str, system: &str, user: &str) -> AnnotatedLlmRequest { AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![ Message::System { content: MessageContent::Text(system.to_string()), diff --git a/crates/adaptive/tests/unit/acg_profile_tests.rs b/crates/adaptive/tests/unit/acg_profile_tests.rs index 3d5e23f7e..08aebcb27 100644 --- a/crates/adaptive/tests/unit/acg_profile_tests.rs +++ b/crates/adaptive/tests/unit/acg_profile_tests.rs @@ -13,6 +13,8 @@ use super::*; fn request(messages: Vec, tools: Option>) -> AnnotatedLlmRequest { AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages, model: Some("gpt-4o".to_string()), params: None, @@ -36,13 +38,15 @@ fn request(messages: Vec, tools: Option>) -> Annota } fn sample_tool(name: &str) -> ToolDefinition { - ToolDefinition { - tool_type: "function".to_string(), + ToolDefinition::Function { function: FunctionDefinition { name: name.to_string(), description: Some("desc".to_string()), parameters: Some(json!({"type":"object","properties":{"a":{"type":"string"}}})), + strict: None, + extra: serde_json::Map::new(), }, + extra: serde_json::Map::new(), } } @@ -53,18 +57,21 @@ fn acg_profile_derivation_covers_anchor_hash_system_fallback_and_empty_tools() { Message::System { content: MessageContent::Parts(vec![ContentPart::Text { text: "System guide".to_string(), + extra: serde_json::Map::new(), }]), name: None, }, Message::User { content: MessageContent::Parts(vec![ContentPart::Text { text: "Language guide".to_string(), + extra: serde_json::Map::new(), }]), name: None, }, Message::Assistant { content: Some(MessageContent::Parts(vec![ContentPart::Text { text: "Acknowledged".to_string(), + extra: serde_json::Map::new(), }])), tool_calls: None, name: None, @@ -125,6 +132,61 @@ fn acg_profile_helpers_cover_none_paths_and_short_hash() { assert_eq!(message_role_tag(&too_short.messages[0]), "user"); } +#[test] +fn system_fingerprint_preserves_sources_roles_and_content_boundaries() { + let mut instructions_and_system = request( + vec![Message::System { + content: MessageContent::Text("b".into()), + name: None, + }], + None, + ); + instructions_and_system.instructions = Some(MessageContent::Text("a".into())); + let joined_system = request( + vec![Message::System { + content: MessageContent::Text("a\nb".into()), + name: None, + }], + None, + ); + let developer = request( + vec![Message::Developer { + content: MessageContent::Text("a\nb".into()), + name: None, + }], + None, + ); + let split_parts = request( + vec![Message::System { + content: MessageContent::Parts(vec![ + ContentPart::Text { + text: "a".into(), + extra: serde_json::Map::new(), + }, + ContentPart::Text { + text: "b".into(), + extra: serde_json::Map::new(), + }, + ]), + name: None, + }], + None, + ); + + assert_ne!( + system_prompt_fingerprint(&instructions_and_system), + system_prompt_fingerprint(&joined_system) + ); + assert_ne!( + system_prompt_fingerprint(&joined_system), + system_prompt_fingerprint(&developer) + ); + assert_ne!( + system_prompt_fingerprint(&joined_system), + system_prompt_fingerprint(&split_parts) + ); +} + #[test] fn acg_profile_image_parts_contribute_stable_fingerprint_signal() { let with_image_a = request( @@ -134,6 +196,7 @@ fn acg_profile_image_parts_contribute_stable_fingerprint_signal() { url: "https://example.com/a.png".to_string(), detail: Some("high".to_string()), }, + extra: serde_json::Map::new(), }]), name: None, }], @@ -146,6 +209,7 @@ fn acg_profile_image_parts_contribute_stable_fingerprint_signal() { url: "https://example.com/b.png".to_string(), detail: Some("high".to_string()), }, + extra: serde_json::Map::new(), }]), name: None, }], @@ -271,3 +335,164 @@ fn acg_profile_fingerprints_cover_alternate_role_sequences() { ); assert_eq!(learning_seed_fingerprint(&system_only), "no-seed"); } + +#[test] +fn acg_profile_covers_extended_roles_and_native_content() { + let all_roles = request( + vec![ + Message::Developer { + content: MessageContent::Text("developer".into()), + name: None, + }, + Message::Function { + content: Some("legacy result".into()), + name: "legacy".into(), + }, + Message::ToolCallItem { + id: None, + call_id: "call_1".into(), + name: "lookup".into(), + arguments: json!({"q": "x"}), + extra: serde_json::Map::new(), + }, + Message::ToolResultItem { + id: None, + call_id: "call_1".into(), + output: json!({"ok": true}), + extra: serde_json::Map::new(), + }, + Message::ProviderNative { + provider: "openai_responses".into(), + kind: "reasoning".into(), + value: json!({"type": "reasoning"}), + }, + ], + None, + ); + let key = derive_acg_profile_key("agent-extended", &all_roles); + assert!(key.contains("roles=developer.function.tool_call.tool_result.provider_native")); + + for (message, prefix) in [ + ( + Message::Function { + content: None, + name: "legacy".into(), + }, + "function:", + ), + ( + Message::ToolCallItem { + id: None, + call_id: "call_1".into(), + name: "lookup".into(), + arguments: json!({"q": "x"}), + extra: serde_json::Map::new(), + }, + "tool-call:", + ), + ( + Message::ToolResultItem { + id: None, + call_id: "call_1".into(), + output: json!({"ok": true}), + extra: serde_json::Map::new(), + }, + "tool-result:", + ), + ( + Message::ProviderNative { + provider: "openai_responses".into(), + kind: "reasoning".into(), + value: json!({"type": "reasoning"}), + }, + "native:", + ), + ] { + assert!(learning_seed_fingerprint(&request(vec![message], None)).starts_with(prefix)); + } + + let seed_for_part = |part| { + learning_seed_fingerprint(&request( + vec![Message::User { + content: MessageContent::Parts(vec![part]), + name: None, + }], + None, + )) + }; + assert_ne!( + seed_for_part(ContentPart::Refusal { + refusal: "no-a".into(), + extra: serde_json::Map::new(), + }), + seed_for_part(ContentPart::Refusal { + refusal: "no-b".into(), + extra: serde_json::Map::new(), + }) + ); + assert_ne!( + seed_for_part(ContentPart::Audio { + audio: json!({"data": "audio-a"}), + extra: serde_json::Map::new(), + }), + seed_for_part(ContentPart::Audio { + audio: json!({"data": "audio-b"}), + extra: serde_json::Map::new(), + }) + ); + assert_ne!( + seed_for_part(ContentPart::ProviderNative { + provider: "openai_responses".into(), + kind: "future".into(), + value: json!({"type": "future", "value": "a"}), + }), + seed_for_part(ContentPart::ProviderNative { + provider: "openai_responses".into(), + kind: "future".into(), + value: json!({"type": "future", "value": "b"}), + }) + ); +} + +#[test] +fn learning_seed_fingerprint_includes_variant_discriminator_fields() { + let seed_for_message = |message| learning_seed_fingerprint(&request(vec![message], None)); + assert_ne!( + seed_for_message(Message::Function { + content: Some("same".into()), + name: "one".into(), + }), + seed_for_message(Message::Function { + content: Some("same".into()), + name: "two".into(), + }) + ); + assert_ne!( + seed_for_message(Message::ToolCallItem { + id: None, + call_id: "call".into(), + name: "one".into(), + arguments: json!({"same": true}), + extra: serde_json::Map::new(), + }), + seed_for_message(Message::ToolCallItem { + id: None, + call_id: "call".into(), + name: "two".into(), + arguments: json!({"same": true}), + extra: serde_json::Map::new(), + }) + ); + assert_ne!( + seed_for_message(Message::ProviderNative { + provider: "openai_responses".into(), + kind: "reasoning".into(), + value: json!({"same": true}), + }), + seed_for_message(Message::ProviderNative { + provider: "anthropic_messages".into(), + kind: "thinking".into(), + value: json!({"same": true}), + }) + ); +} diff --git a/crates/adaptive/tests/unit/adaptive_hints_intercept_tests.rs b/crates/adaptive/tests/unit/adaptive_hints_intercept_tests.rs index 9d260fac5..e3ef485b5 100644 --- a/crates/adaptive/tests/unit/adaptive_hints_intercept_tests.rs +++ b/crates/adaptive/tests/unit/adaptive_hints_intercept_tests.rs @@ -171,6 +171,8 @@ fn test_adaptive_hints_intercept_injects_prediction_hints_and_manual_override() crate::context_helpers::set_latency_sensitivity(5).unwrap(); let annotated = AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![Message::User { content: MessageContent::Text("hello".into()), name: None, diff --git a/crates/adaptive/tests/unit/cache_diagnostics_tests.rs b/crates/adaptive/tests/unit/cache_diagnostics_tests.rs index b46c7b846..433ab4718 100644 --- a/crates/adaptive/tests/unit/cache_diagnostics_tests.rs +++ b/crates/adaptive/tests/unit/cache_diagnostics_tests.rs @@ -90,6 +90,8 @@ fn make_hot_cache(stable_prefix_length: Option) -> HotCache { fn sample_request(model: Option<&str>) -> AnnotatedLlmRequest { AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![ Message::System { content: MessageContent::Text("You are a careful planner".to_string()), diff --git a/crates/adaptive/tests/unit/runtime_tests.rs b/crates/adaptive/tests/unit/runtime_tests.rs index 77231ccdb..f14d83fb0 100644 --- a/crates/adaptive/tests/unit/runtime_tests.rs +++ b/crates/adaptive/tests/unit/runtime_tests.rs @@ -37,6 +37,8 @@ fn short_hash(value: &str) -> &str { fn sample_annotated_request(model: Option<&str>) -> AnnotatedLlmRequest { AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![ Message::System { content: MessageContent::Text("You are a careful planner".to_string()), @@ -70,6 +72,8 @@ fn sample_annotated_request(model: Option<&str>) -> AnnotatedLlmRequest { fn sample_layered_request(model: Option<&str>, language_guide: &str) -> AnnotatedLlmRequest { AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![ Message::System { content: MessageContent::Text("You are a careful planner".to_string()), @@ -410,7 +414,7 @@ fn adaptive_acg_defaults_and_profile_key_behavior_stay_stable() { ); assert_eq!( profile_key, - "agent-1::model=claude-sonnet-4::roles=system.user::system=sha256:97f793c76::anchor=no-anchor::tools=no-tools" + "agent-1::model=claude-sonnet-4::roles=system.user::system=sha256:3087d8fd4::anchor=no-anchor::tools=no-tools" ); let learning_key = crate::acg_profile::derive_acg_learning_key( "agent-1", @@ -422,11 +426,13 @@ fn adaptive_acg_defaults_and_profile_key_behavior_stay_stable() { "user:{}", crate::acg::sha256_hex("Summarize the latest findings") )), - short_hash(&crate::acg::sha256_hex("You are a careful planner")), + "sha256:3087d8fd4", ); assert_eq!(learning_key, expected_learning_key,); let grown_chat_request = AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![ Message::System { content: MessageContent::Text("You are a careful planner".to_string()), @@ -504,6 +510,8 @@ fn adaptive_acg_defaults_and_profile_key_behavior_stay_stable() { ); let rust_bundle_variant = AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![ Message::System { content: MessageContent::Text("You are a careful planner".to_string()), diff --git a/crates/cli/src/gateway/mod.rs b/crates/cli/src/gateway/mod.rs index 41dcb358f..4af307968 100644 --- a/crates/cli/src/gateway/mod.rs +++ b/crates/cli/src/gateway/mod.rs @@ -28,12 +28,13 @@ use nemo_relay::api::llm::{ use nemo_relay::api::runtime::{ LlmExecutionNextFn, LlmJsonStream, LlmStreamExecutionNextFn, TASK_SCOPE_STACK, }; +use nemo_relay::codec::request::AnnotatedLlmRequest; use nemo_relay::codec::resolve::{ - ProviderSurface, response_codec as build_response_codec, + ProviderSurface, request_codec as build_request_codec, response_codec as build_response_codec, streaming_codec as build_streaming_codec, }; use nemo_relay::codec::streaming::StreamingCodec; -use nemo_relay::codec::traits::LlmResponseCodec; +use nemo_relay::codec::traits::{LlmCodec, LlmResponseCodec}; use nemo_relay::error::{FlowError, UpstreamFailure, UpstreamFailureClass}; use serde_json::{Value, json}; @@ -149,17 +150,48 @@ async fn run_unmanaged_gateway( // Codecs registered for each managed provider route. Routes that emit LLM events but lack a typed // codec (count_tokens) return `None` so the runtime still wraps the call but skips annotation. struct RouteCodecs { + request: Option>, streaming: Option>, response: Option>, } +struct GatewayRequestCodec { + inner: Arc, +} + +impl LlmCodec for GatewayRequestCodec { + fn decode(&self, request: &LlmRequest) -> nemo_relay::error::Result { + self.inner.decode(request) + } + + fn encode( + &self, + annotated: &AnnotatedLlmRequest, + original: &LlmRequest, + ) -> nemo_relay::error::Result { + let original_streaming = self.inner.decode(original)?.stream.unwrap_or(false); + let edited_streaming = annotated.stream.unwrap_or(false); + if edited_streaming != original_streaming { + return Err(FlowError::InvalidArgument( + "gateway request interceptors cannot change stream mode; preserve the original stream value" + .into(), + )); + } + self.inner.encode(annotated, original) + } +} + fn codecs_for_route(route: ProviderRoute) -> RouteCodecs { match route.provider_surface() { Some(surface) => RouteCodecs { + request: Some(Arc::new(GatewayRequestCodec { + inner: build_request_codec(surface), + })), streaming: Some(build_streaming_codec(surface)), response: Some(build_response_codec(surface)), }, None => RouteCodecs { + request: None, streaming: None, response: None, }, @@ -207,6 +239,7 @@ async fn run_managed_buffered( .attributes(attributes) .metadata(metadata) .model_name_opt(model_name) + .codec_opt(codecs.request) .response_codec_opt(codecs.response) .build(); let result = TASK_SCOPE_STACK @@ -407,6 +440,7 @@ async fn run_managed_streaming( .attributes(attributes) .metadata(metadata) .model_name_opt(model_name) + .codec_opt(codecs.request) .response_codec_opt(codecs.response) .build(); let json_stream_result = TASK_SCOPE_STACK diff --git a/crates/cli/tests/coverage/shared/gateway_tests.rs b/crates/cli/tests/coverage/shared/gateway_tests.rs index d6ad7872c..51239ca95 100644 --- a/crates/cli/tests/coverage/shared/gateway_tests.rs +++ b/crates/cli/tests/coverage/shared/gateway_tests.rs @@ -327,6 +327,90 @@ fn selects_provider_routes() { assert_eq!(ProviderRoute::from_path("/unsupported"), None); } +#[test] +fn generation_routes_have_request_codecs_and_passthrough_routes_do_not() { + for route in [ + ProviderRoute::AnthropicMessages, + ProviderRoute::OpenAiChatCompletions, + ProviderRoute::OpenAiResponses, + ] { + let codecs = codecs_for_route(route); + assert!( + codecs.request.is_some(), + "missing request codec for {route:?}" + ); + assert!( + codecs.streaming.is_some(), + "missing stream codec for {route:?}" + ); + assert!( + codecs.response.is_some(), + "missing response codec for {route:?}" + ); + } + + for route in [ + ProviderRoute::AnthropicCountTokens, + ProviderRoute::OpenAiModels, + ] { + let codecs = codecs_for_route(route); + assert!( + codecs.request.is_none(), + "unexpected request codec for {route:?}" + ); + assert!( + codecs.streaming.is_none(), + "unexpected stream codec for {route:?}" + ); + assert!( + codecs.response.is_none(), + "unexpected response codec for {route:?}" + ); + } +} + +#[test] +fn generation_route_codecs_reject_stream_mode_changes() { + for (route, content) in [ + ( + ProviderRoute::AnthropicMessages, + json!({ + "model": "claude-test", + "max_tokens": 32, + "messages": [{"role": "user", "content": "hello"}], + }), + ), + ( + ProviderRoute::OpenAiChatCompletions, + json!({ + "model": "gpt-test", + "messages": [{"role": "user", "content": "hello"}], + }), + ), + ( + ProviderRoute::OpenAiResponses, + json!({"model": "gpt-test", "input": "hello"}), + ), + ] { + let codec = codecs_for_route(route).request.unwrap(); + for original_streaming in [false, true] { + let mut content = content.clone(); + content["stream"] = json!(original_streaming); + let request = LlmRequest { + headers: Map::new(), + content, + }; + let mut annotated = codec.decode(&request).unwrap(); + annotated.stream = Some(!original_streaming); + let error = codec.encode(&annotated, &request).unwrap_err(); + assert!( + error.to_string().contains("cannot change stream mode"), + "unexpected error for {route:?}: {error}" + ); + } + } +} + #[test] fn dispatch_override_routes_cover_models_and_count_tokens() { for alias in ["openai_models", "openai.models", "/models", "/v1/models"] { diff --git a/crates/cli/tests/coverage/shared/plugins_tests.rs b/crates/cli/tests/coverage/shared/plugins_tests.rs index cb4e3b9f2..b40e38d05 100644 --- a/crates/cli/tests/coverage/shared/plugins_tests.rs +++ b/crates/cli/tests/coverage/shared/plugins_tests.rs @@ -310,6 +310,14 @@ fn typed_editor_model_contains_pii_redaction_options() { ); let builtin = schema.field("builtin").unwrap().schema().unwrap(); + assert_eq!(builtin.field("preset").unwrap().kind, EditorFieldKind::Enum); + assert!( + builtin + .field("preset") + .unwrap() + .enum_values + .contains(&"trajectory_context") + ); assert_eq!(builtin.field("action").unwrap().kind, EditorFieldKind::Enum); assert!( builtin @@ -322,6 +330,17 @@ fn typed_editor_model_contains_pii_redaction_options() { builtin.field("target_paths").unwrap().kind, EditorFieldKind::List ); + assert_eq!( + builtin.field("custom_mark_payload_policy").unwrap().kind, + EditorFieldKind::Enum + ); + assert!( + builtin + .field("custom_mark_payload_policy") + .unwrap() + .enum_values + .contains(&"redact_all_leaves") + ); assert_eq!( builtin.field("detector").unwrap().kind, EditorFieldKind::Enum diff --git a/crates/cli/tests/coverage/shared/server_tests.rs b/crates/cli/tests/coverage/shared/server_tests.rs index a18bdbe8a..c624182d8 100644 --- a/crates/cli/tests/coverage/shared/server_tests.rs +++ b/crates/cli/tests/coverage/shared/server_tests.rs @@ -4,18 +4,21 @@ use std::ffi::OsString; use std::future::Future; use std::pin::Pin; -use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; use axum::body::Body; +use axum::extract::State; use axum::http::{Request, StatusCode, header}; -use axum::response::IntoResponse; +use axum::response::{IntoResponse, Response}; use bytes::Bytes; use futures_util::stream; use http_body_util::BodyExt; use nemo_relay::api::event::ScopeCategory; +use nemo_relay::api::llm::LlmRequestInterceptOutcome; use nemo_relay::api::registry::{ - deregister_tool_conditional_execution_guardrail, register_tool_conditional_execution_guardrail, + deregister_llm_request_intercept, deregister_tool_conditional_execution_guardrail, + register_llm_request_intercept, register_tool_conditional_execution_guardrail, }; use nemo_relay::api::subscriber::{deregister_subscriber, flush_subscribers, register_subscriber}; use nemo_relay::plugin::dynamic::DynamicPluginKind; @@ -105,6 +108,14 @@ impl Drop for SubscriberCleanup { } } +struct RequestInterceptCleanup(&'static str); + +impl Drop for RequestInterceptCleanup { + fn drop(&mut self) { + let _ = deregister_llm_request_intercept(self.0); + } +} + fn test_http_client() -> reqwest::Client { reqwest::Client::new() } @@ -2539,6 +2550,313 @@ async fn gateway_accepts_codex_responses_path() { assert_eq!(body["authorization"], json!("Bearer test")); } +#[tokio::test] +async fn gateway_request_codec_exposes_annotations_and_applies_buffered_edits() { + let intercept_name = "server-gateway-request-codec-buffered-edit"; + let _ = deregister_llm_request_intercept(intercept_name); + let _cleanup = RequestInterceptCleanup(intercept_name); + let observed = Arc::new(Mutex::new(None)); + let captured = observed.clone(); + register_llm_request_intercept( + intercept_name, + 1, + false, + Arc::new(move |_name, mut request, annotated| { + if request.headers.get("x-codec-test").and_then(Value::as_str) != Some("buffered") { + return Ok(LlmRequestInterceptOutcome::new(request, annotated)); + } + let mut annotated = annotated.expect("gateway generation route must have a codec"); + *captured.lock().unwrap() = Some(serde_json::to_value(&annotated).unwrap()); + let nemo_relay::codec::request::Message::User { content, .. } = + &mut annotated.messages[0] + else { + panic!("expected portable Responses string input"); + }; + *content = nemo_relay::codec::request::MessageContent::Text("edited".into()); + request + .headers + .insert("x-test-intercept".into(), json!("visible")); + Ok(LlmRequestInterceptOutcome::new(request, Some(annotated))) + }), + ) + .unwrap(); + + let upstream = spawn_upstream(false).await; + let mut config = test_config(); + config.openai_base_url = upstream.url(); + let response = router(config) + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("content-type", "application/json") + .header("x-codec-test", "buffered") + .body(Body::from( + json!({"model": "gpt-test", "input": "original"}).to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + let body: Value = + serde_json::from_slice(&response.into_body().collect().await.unwrap().to_bytes()).unwrap(); + assert_eq!(body["input"], json!("edited")); + assert_eq!(body["x_test_intercept"], json!("visible")); + let annotation = observed.lock().unwrap().clone().unwrap(); + assert_eq!(annotation["messages"][0]["content"], json!("original")); + assert_eq!(annotation["api_specific"]["api"], json!("openai_responses")); +} + +#[tokio::test] +async fn gateway_request_codec_rejects_raw_body_mutation_before_upstream() { + let intercept_name = "server-gateway-request-codec-raw-rejection"; + let _ = deregister_llm_request_intercept(intercept_name); + let _cleanup = RequestInterceptCleanup(intercept_name); + register_llm_request_intercept( + intercept_name, + 1, + false, + Arc::new(|_name, mut request, annotated| { + if request.headers.get("x-codec-test").and_then(Value::as_str) == Some("raw") { + request.content["input"] = json!("forbidden raw edit"); + } + Ok(LlmRequestInterceptOutcome::new(request, annotated)) + }), + ) + .unwrap(); + + let upstream = spawn_upstream(false).await; + let mut config = test_config(); + config.openai_base_url = upstream.url(); + let response = router(config) + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("content-type", "application/json") + .header("x-codec-test", "raw") + .body(Body::from( + json!({"model": "gpt-test", "input": "original"}).to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::BAD_REQUEST); +} + +#[tokio::test] +async fn gateway_request_codec_rejects_malformed_modeled_structure_before_upstream() { + let (upstream, captured_requests) = spawn_request_codec_matrix_upstream().await; + let mut config = test_config(); + config.openai_base_url = upstream.url(); + let response = router(config) + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-test", + "messages": [], + "response_format": "json_object", + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + assert!(captured_requests.lock().unwrap().is_empty()); +} + +#[tokio::test] +async fn gateway_request_codec_rejects_stream_mode_changes_before_upstream() { + let intercept_name = "server-gateway-request-codec-stream-mode-rejection"; + let _ = deregister_llm_request_intercept(intercept_name); + let _cleanup = RequestInterceptCleanup(intercept_name); + register_llm_request_intercept( + intercept_name, + 1, + false, + Arc::new(|_name, request, annotated| { + if request + .headers + .get("x-codec-stream-toggle") + .and_then(Value::as_str) + == Some("true") + { + let mut annotated = annotated.expect("generation route must expose an annotation"); + annotated.stream = Some(!annotated.stream.unwrap_or(false)); + return Ok(LlmRequestInterceptOutcome::new(request, Some(annotated))); + } + Ok(LlmRequestInterceptOutcome::new(request, annotated)) + }), + ) + .unwrap(); + + let (upstream, captured_requests) = spawn_request_codec_matrix_upstream().await; + let mut config = test_config(); + config.openai_base_url = upstream.url(); + let app = router(config); + + for streaming in [false, true] { + let response = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("content-type", "application/json") + .header("x-codec-stream-toggle", "true") + .body(Body::from( + json!({ + "model": "gpt-test", + "input": "original", + "stream": streaming, + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + } + assert!(captured_requests.lock().unwrap().is_empty()); +} + +#[tokio::test] +async fn gateway_request_codecs_apply_buffered_and_streaming_edits_on_all_generation_routes() { + let intercept_name = "server-gateway-request-codec-all-generation-routes"; + let _ = deregister_llm_request_intercept(intercept_name); + let _cleanup = RequestInterceptCleanup(intercept_name); + let observed = Arc::new(Mutex::new(Vec::new())); + let captured_annotations = observed.clone(); + register_llm_request_intercept( + intercept_name, + 1, + false, + Arc::new(move |_name, mut request, annotated| { + let Some(marker) = request + .headers + .get("x-codec-matrix") + .and_then(Value::as_str) + .map(str::to_string) + else { + return Ok(LlmRequestInterceptOutcome::new(request, annotated)); + }; + let mut annotated = annotated.expect("generation route must expose an annotation"); + captured_annotations.lock().unwrap().push(json!({ + "marker": marker, + "annotation": annotated, + })); + let nemo_relay::codec::request::Message::User { content, .. } = + &mut annotated.messages[0] + else { + panic!("expected the first request item to be a portable user message"); + }; + *content = nemo_relay::codec::request::MessageContent::Text(format!("edited-{marker}")); + request + .headers + .insert("x-codec-edited".into(), json!(marker)); + Ok(LlmRequestInterceptOutcome::new(request, Some(annotated))) + }), + ) + .unwrap(); + + let (upstream, captured_requests) = spawn_request_codec_matrix_upstream().await; + let mut config = test_config(); + config.openai_base_url = upstream.url(); + config.anthropic_base_url = upstream.url(); + let app = router(config); + + for (surface, uri, mut payload) in [ + ( + "anthropic", + "/v1/messages", + json!({ + "model": "claude-sonnet-4-20250514", + "max_tokens": 32, + "messages": [{"role": "user", "content": "original"}] + }), + ), + ( + "chat", + "/v1/chat/completions", + json!({ + "model": "gpt-4.1", + "messages": [{"role": "user", "content": "original"}] + }), + ), + ( + "responses", + "/v1/responses", + json!({"model": "gpt-5", "input": "original"}), + ), + ] { + for streaming in [false, true] { + payload["stream"] = json!(streaming); + let marker = format!( + "{surface}-{}", + if streaming { "streaming" } else { "buffered" } + ); + let response = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri(uri) + .header("content-type", "application/json") + .header("x-codec-matrix", &marker) + .body(Body::from(payload.to_string())) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK, "{marker}"); + response.into_body().collect().await.unwrap(); + } + } + + let captured = captured_requests.lock().unwrap(); + assert_eq!(captured.len(), 6); + for request in captured.iter() { + let marker = request["marker"].as_str().unwrap(); + assert_eq!(request["edited_header"], json!(marker)); + let edited = format!("edited-{marker}"); + match request["path"].as_str().unwrap() { + "/v1/messages" | "/v1/chat/completions" => { + assert_eq!(request["body"]["messages"][0]["content"], json!(edited)); + } + "/v1/responses" => assert_eq!(request["body"]["input"], json!(edited)), + path => panic!("unexpected captured path {path}"), + } + } + + let annotations = observed.lock().unwrap(); + assert_eq!(annotations.len(), 6); + for observation in annotations.iter() { + let marker = observation["marker"].as_str().unwrap(); + let expected_api = if marker.starts_with("anthropic-") { + "anthropic_messages" + } else if marker.starts_with("chat-") { + "openai_chat" + } else { + "openai_responses" + }; + assert_eq!( + observation["annotation"]["api_specific"]["api"], + json!(expected_api) + ); + } +} + #[tokio::test] async fn gateway_preserves_streaming_body() { let upstream = spawn_upstream(true).await; @@ -2646,7 +2964,13 @@ async fn gateway_returns_bad_gateway_when_upstream_is_unreachable() { .method("POST") .uri("/v1/chat/completions") .header("content-type", "application/json") - .body(Body::from(json!({ "model": "gpt-test" }).to_string())) + .body(Body::from( + json!({ + "model": "gpt-test", + "messages": [{ "role": "user", "content": "hello" }] + }) + .to_string(), + )) .unwrap(), ) .await @@ -3068,6 +3392,114 @@ async fn spawn_upstream(streaming: bool) -> TestServer { } } +async fn spawn_request_codec_matrix_upstream() -> (TestServer, Arc>>) { + async fn provider( + State(captured): State>>>, + request: Request, + ) -> Response { + let path = request.uri().path().to_string(); + let marker = request + .headers() + .get("x-codec-matrix") + .and_then(|value| value.to_str().ok()) + .unwrap() + .to_string(); + let edited_header = request + .headers() + .get("x-codec-edited") + .and_then(|value| value.to_str().ok()) + .unwrap() + .to_string(); + let body = request.into_body().collect().await.unwrap().to_bytes(); + let payload: Value = serde_json::from_slice(&body).unwrap(); + captured.lock().unwrap().push(json!({ + "path": path, + "marker": marker, + "edited_header": edited_header, + "body": payload, + })); + + if !payload["stream"].as_bool().unwrap_or(false) { + return match path.as_str() { + "/v1/messages" => Json(json!({ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-20250514", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1} + })) + .into_response(), + "/v1/chat/completions" => Json(json!({ + "id": "chatcmpl_1", + "model": "gpt-4.1", + "choices": [{ + "message": {"role": "assistant", "content": "ok"}, + "finish_reason": "stop" + }] + })) + .into_response(), + "/v1/responses" => Json(json!({ + "id": "resp_1", + "model": "gpt-5", + "status": "completed", + "output": [] + })) + .into_response(), + _ => StatusCode::NOT_FOUND.into_response(), + }; + } + + let events: &[&[u8]] = match path.as_str() { + "/v1/messages" => &[ + b"data: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_1\",\"type\":\"message\",\"role\":\"assistant\",\"model\":\"claude-sonnet-4-20250514\",\"content\":[],\"usage\":{\"input_tokens\":1,\"output_tokens\":0}}}\n\n", + b"data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"input_tokens\":1,\"output_tokens\":1}}\n\n", + b"data: {\"type\":\"message_stop\"}\n\n", + ], + "/v1/chat/completions" => &[ + b"data: {\"id\":\"chatcmpl_1\",\"object\":\"chat.completion.chunk\",\"model\":\"gpt-4.1\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\"},\"finish_reason\":null}]}\n\n", + b"data: {\"id\":\"chatcmpl_1\",\"object\":\"chat.completion.chunk\",\"model\":\"gpt-4.1\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"ok\"},\"finish_reason\":\"stop\"}]}\n\n", + b"data: [DONE]\n\n", + ], + "/v1/responses" => &[ + b"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_1\",\"model\":\"gpt-5\"}}\n\n", + b"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_1\",\"model\":\"gpt-5\",\"status\":\"completed\",\"output\":[]}}\n\n", + ], + _ => return StatusCode::NOT_FOUND.into_response(), + }; + let chunks = stream::iter( + events + .iter() + .map(|event| Ok::<_, std::convert::Infallible>(Bytes::copy_from_slice(event))), + ); + ( + [(header::CONTENT_TYPE, "text/event-stream")], + Body::from_stream(chunks), + ) + .into_response() + } + + let captured = Arc::new(Mutex::new(Vec::new())); + let app = Router::new() + .route("/v1/messages", post(provider)) + .route("/v1/chat/completions", post(provider)) + .route("/v1/responses", post(provider)) + .with_state(captured.clone()); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let handle = tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + ( + TestServer { + url: format!("http://{address}"), + handle, + }, + captured, + ) +} + async fn spawn_failing_stream_upstream() -> TestServer { async fn stream_response() -> impl IntoResponse { // First chunk is a valid JSON SSE event so the managed pipeline opens cleanly; the diff --git a/crates/core/src/api/llm.rs b/crates/core/src/api/llm.rs index 1a6e66dac..8484d18e3 100644 --- a/crates/core/src/api/llm.rs +++ b/crates/core/src/api/llm.rs @@ -36,7 +36,9 @@ use crate::error::{FlowError, Result}; use crate::json::Json; use crate::stream::LlmStreamWrapper; -pub use nemo_relay_types::api::llm::{LlmAttributes, LlmRequest, LlmRequestInterceptOutcome}; +pub use nemo_relay_types::api::llm::{ + LLM_REQUEST_INTERCEPT_OUTCOME_SCHEMA, LlmAttributes, LlmRequest, LlmRequestInterceptOutcome, +}; #[derive(Clone)] struct CapturedLlmScopeStack(ScopeStackHandle); diff --git a/crates/core/src/api/skill_load.rs b/crates/core/src/api/skill_load.rs index 88788c752..5ed69c3d2 100644 --- a/crates/core/src/api/skill_load.rs +++ b/crates/core/src/api/skill_load.rs @@ -11,6 +11,30 @@ use strum::{EnumString, IntoStaticStr}; pub(crate) const HANDLED_METADATA_KEY: &str = "nemo_relay.skill_load_handled"; pub(crate) const PRECOMPUTED_METADATA_KEY: &str = "nemo_relay.skill_loads"; +const CAT_REJECTED_OPTIONS: &[&str] = &[ + "-h", + "--help", + "--version", + "--show-all", + "--number-nonblank", + "--show-ends", + "--number", + "--squeeze-blank", + "--show-tabs", + "--show-nonprinting", +]; +const CAT_REJECTED_SHORT_OPTIONS: &[char] = &['A', 'b', 'e', 'E', 'n', 's', 't', 'T', 'v']; +const BAT_ALLOWED_OPTIONS: &[&str] = &["-p", "--plain"]; +const POWERSHELL_ALLOWED_OPTIONS: &[&str] = &[ + "asbytestream", + "encoding", + "force", + "literalpath", + "path", + "raw", + "readcount", +]; + #[derive(Debug, Clone, Copy, PartialEq, Eq, EnumString, IntoStaticStr)] #[strum(serialize_all = "snake_case")] pub(crate) enum SkillLoadSource { @@ -35,7 +59,10 @@ pub(crate) fn detect(tool_name: &str, args: &Value) -> Vec { structured_skill_names(args), ) } else if is_shell_tool(&normalized_tool) { - (SkillLoadSource::ShellRead, shell_skill_names(args)) + ( + SkillLoadSource::ShellRead, + shell_skill_names(&normalized_tool, args), + ) } else { return Vec::new(); }; @@ -92,7 +119,7 @@ fn structured_skill_names(args: &Value) -> Vec { visit_named_values(args, |key, value| { if matches!( key.as_str(), - "path" | "filepath" | "filename" | "file" | "paths" + "path" | "filepath" | "filename" | "file" | "paths" | "uri" | "absolutepath" ) { collect_path_skill_names(value, &mut names); } @@ -100,7 +127,7 @@ fn structured_skill_names(args: &Value) -> Vec { names } -fn shell_skill_names(args: &Value) -> Vec { +fn shell_skill_names(tool_name: &str, args: &Value) -> Vec { let mut commands = Vec::new(); visit_named_values(args, |key, value| { if matches!(key.as_str(), "command" | "cmd") @@ -109,20 +136,24 @@ fn shell_skill_names(args: &Value) -> Vec { commands.push(value.to_string()); } }); + let allow_direct_cat = !matches!(tool_name, "powershell" | "pwsh"); commands .into_iter() - .flat_map(|command| complete_reader_paths(&command)) + .flat_map(|command| complete_reader_paths(&command, allow_direct_cat)) .filter_map(|path| skill_name_from_path(&path)) .collect() } fn is_structured_reader(tool_name: &str) -> bool { - const READERS: [&str; 5] = [ + const READERS: [&str; 8] = [ "read", "readfile", "readtextfile", "readmultiplefiles", "fileread", + "readresource", + "getfilecontents", + "readfilecontent", ]; let segments = tool_name .split(|character: char| !character.is_ascii_alphanumeric()) @@ -136,7 +167,9 @@ fn is_structured_reader(tool_name: &str) -> bool { fn is_shell_tool(tool_name: &str) -> bool { matches!( tool_name, - "bash" + "sh" | "bash" + | "zsh" + | "fish" | "shell" | "shellcommand" | "exec" @@ -147,6 +180,7 @@ fn is_shell_tool(tool_name: &str) -> bool { | "runshellcommand" | "shellexec" | "powershell" + | "pwsh" ) } @@ -156,7 +190,7 @@ fn has_partial_read_controls(value: &Value) -> bool { if !partial { let key = normalize_identifier(&key); partial = match key.as_str() { - "offset" => value.as_i64().is_some_and(|offset| offset != 0), + "offset" => value.as_i64() != Some(0), "limit" | "range" | "head" | "tail" | "startline" | "endline" | "linestart" | "lineend" => !value.is_null(), _ => false, @@ -197,7 +231,12 @@ fn visit_named_values(value: &Value, mut visit: impl FnMut(String, &Value)) { } fn skill_name_from_path(path: &str) -> Option { - let path = path.trim().trim_matches(['\'', '"']); + if path.trim() != path + || path.starts_with(['\'', '"']) + || path.ends_with(['\'', '"', '/', '\\']) + { + return None; + } let components = path .split(['/', '\\']) .filter(|component| !component.is_empty()) @@ -214,45 +253,162 @@ fn skill_name_from_path(path: &str) -> Option { Some((*parent).to_string()) } -fn complete_reader_paths(command: &str) -> Vec { - if command.contains(['\n', '\r']) { +fn complete_reader_paths(command: &str, allow_direct_cat: bool) -> Vec { + let Some(words) = standalone_command_words(command) else { return Vec::new(); + }; + complete_reader_word_paths(&words, allow_direct_cat) +} + +fn standalone_command_words(command: &str) -> Option> { + if command.contains(['\n', '\r']) { + return None; } // Preserve Windows separators: shell-words treats a lone backslash as an escape. let escaped_windows_paths = command.replace('\\', "\\\\"); - let Ok(words) = shell_words::split(&escaped_windows_paths) else { - return Vec::new(); - }; + let words = shell_words::split(&escaped_windows_paths).ok()?; if words.is_empty() || words.iter().any(|word| { - matches!( - word.as_str(), - "|" | "||" | "&" | "&&" | ";" | "<" | ">" | "<<" | ">>" - ) || word.contains("$(") - || word.contains('`') + word.contains(['|', '&', ';', '<', '>']) || word.contains("$(") || word.contains('`') }) { - return Vec::new(); + return None; } + Some(words) +} + +fn complete_reader_word_paths(words: &[String], allow_direct_cat: bool) -> Vec { let Some(executable) = words.first().and_then(|word| executable_name(word)) else { return Vec::new(); }; match executable.as_str() { - "cat" => positional_paths(&words[1..], &[]), - "bat" | "batcat" => positional_paths(&words[1..], &["-r", "--line-range"]), + "cat" if allow_direct_cat => cat_paths(&words[1..]), + "bat" | "batcat" => bat_paths(&words[1..]), "get-content" => powershell_content_paths(&words[1..]), + "sh" | "bash" | "zsh" => posix_shell_wrapper_paths(&executable, &words[1..]), + "fish" => fish_shell_wrapper_paths(&words[1..]), + "powershell" | "pwsh" => powershell_wrapper_paths(&words[1..]), _ => Vec::new(), } } -fn positional_paths(words: &[String], rejected_flags: &[&str]) -> Vec { +fn posix_shell_wrapper_paths(executable: &str, words: &[String]) -> Vec { + let [flag, command] = words else { + return Vec::new(); + }; + if !matches!( + (executable, flag.as_str()), + ("sh", "-c") | ("bash" | "zsh", "-c" | "-lc") + ) { + return Vec::new(); + } + let Some(command_words) = standalone_command_words(command) else { + return Vec::new(); + }; + posix_reader_word_paths(&command_words) +} + +fn fish_shell_wrapper_paths(words: &[String]) -> Vec { + let command = match words { + [flag, command] if flag == "-c" => command.as_str(), + [argument] => argument.strip_prefix("--command=").unwrap_or_default(), + _ => return Vec::new(), + }; + let Some(command_words) = standalone_command_words(command) else { + return Vec::new(); + }; + posix_reader_word_paths(&command_words) +} + +fn powershell_wrapper_paths(words: &[String]) -> Vec { + let [flag, command @ ..] = words else { + return Vec::new(); + }; + if !flag.eq_ignore_ascii_case("-command") || command.is_empty() { + return Vec::new(); + } + if let [command] = command { + let Some(command_words) = standalone_command_words(command) else { + return Vec::new(); + }; + return powershell_reader_word_paths(&command_words); + } + powershell_reader_word_paths(command) +} + +fn posix_reader_word_paths(words: &[String]) -> Vec { + let Some(executable) = words.first().and_then(|word| executable_name(word)) else { + return Vec::new(); + }; + match executable.as_str() { + "cat" => cat_paths(&words[1..]), + "bat" | "batcat" => bat_paths(&words[1..]), + _ => Vec::new(), + } +} + +fn powershell_reader_word_paths(words: &[String]) -> Vec { + let Some(executable) = words.first().and_then(|word| executable_name(word)) else { + return Vec::new(); + }; + if executable != "get-content" { + return Vec::new(); + } + powershell_content_paths(&words[1..]) +} + +fn cat_paths(words: &[String]) -> Vec { if words.iter().any(|word| { - rejected_flags - .iter() - .any(|flag| word.eq_ignore_ascii_case(flag) || word.starts_with(&format!("{flag}="))) + is_rejected_option(word, CAT_REJECTED_OPTIONS, CAT_REJECTED_SHORT_OPTIONS) + || (word.starts_with('-') && word != "--" && !word.eq_ignore_ascii_case("-u")) }) { return Vec::new(); } + positional_paths(words) +} + +fn bat_paths(words: &[String]) -> Vec { + if words + .iter() + .filter(|word| { + BAT_ALLOWED_OPTIONS + .iter() + .any(|option| word.eq_ignore_ascii_case(option)) + }) + .count() + != 1 + || words.iter().any(|word| { + word.starts_with('-') + && word != "--" + && !BAT_ALLOWED_OPTIONS + .iter() + .any(|option| word.eq_ignore_ascii_case(option)) + }) + { + return Vec::new(); + } + positional_paths(words) +} + +fn is_rejected_option( + word: &str, + rejected_options: &[&str], + rejected_short_options: &[char], +) -> bool { + rejected_options + .iter() + .any(|option| word.eq_ignore_ascii_case(option) || word.starts_with(&format!("{option}="))) + || word + .strip_prefix('-') + .filter(|options| !options.starts_with('-')) + .is_some_and(|options| { + options + .chars() + .any(|option| rejected_short_options.contains(&option)) + }) +} + +fn positional_paths(words: &[String]) -> Vec { words .iter() .filter(|word| !word.starts_with('-')) @@ -262,9 +418,12 @@ fn positional_paths(words: &[String], rejected_flags: &[&str]) -> Vec { fn powershell_content_paths(words: &[String]) -> Vec { if words.iter().any(|word| { - ["-totalcount", "-tail", "-head", "-first", "-last"] - .iter() - .any(|flag| word.eq_ignore_ascii_case(flag) || word.starts_with(&format!("{flag}:"))) + word.starts_with('-') + && !powershell_option_name(word).is_some_and(|option| { + POWERSHELL_ALLOWED_OPTIONS + .iter() + .any(|allowed| option.eq_ignore_ascii_case(allowed)) + }) }) { return Vec::new(); } @@ -273,10 +432,21 @@ fn powershell_content_paths(words: &[String]) -> Vec { let mut index = 0; while index < words.len() { let word = &words[index]; - if word.eq_ignore_ascii_case("-path") || word.eq_ignore_ascii_case("-literalpath") { - if let Some(path) = words.get(index + 1) { + let option = powershell_option_name(word); + if option.is_some_and(|option| { + option.eq_ignore_ascii_case("path") || option.eq_ignore_ascii_case("literalpath") + }) { + if let Some(path) = powershell_option_value(word) { + paths.push(path.to_string()); + } else if let Some(path) = words.get(index + 1) { paths.push(path.clone()); + index += 2; + continue; } + } else if option.is_some_and(|option| { + option.eq_ignore_ascii_case("encoding") || option.eq_ignore_ascii_case("readcount") + }) && powershell_option_value(word).is_none() + { index += 2; continue; } @@ -288,6 +458,16 @@ fn powershell_content_paths(words: &[String]) -> Vec { paths } +fn powershell_option_name(word: &str) -> Option<&str> { + word.strip_prefix('-')?.split([':', '=']).next() +} + +fn powershell_option_value(word: &str) -> Option<&str> { + word.split_once([':', '=']) + .map(|(_, value)| value) + .filter(|value| !value.is_empty()) +} + fn executable_name(executable: &str) -> Option { executable .rsplit(['/', '\\']) diff --git a/crates/core/src/api/subscriber.rs b/crates/core/src/api/subscriber.rs index 48223e007..02ce9a322 100644 --- a/crates/core/src/api/subscriber.rs +++ b/crates/core/src/api/subscriber.rs @@ -72,6 +72,10 @@ pub fn deregister_subscriber(name: &str) -> Result { /// Wait for all subscriber callbacks queued before this call to finish. /// +/// Call this helper outside native subscriber callbacks. A re-entrant call returns without +/// waiting to avoid blocking the dispatcher, so callbacks later in the same dispatch snapshot can +/// still run. +/// /// Native targets deliver subscriber callbacks on a background dispatcher so /// event-producing APIs do not wait for observer work. Call this helper from /// tests, shutdown paths, or exporter lifecycle code when callers need a diff --git a/crates/core/src/codec/anthropic.rs b/crates/core/src/codec/anthropic.rs index 22147815f..29e5a78d1 100644 --- a/crates/core/src/codec/anthropic.rs +++ b/crates/core/src/codec/anthropic.rs @@ -23,8 +23,9 @@ use crate::error::{FlowError, Result}; use crate::json::Json; use super::request::{ - AnnotatedLlmRequest, FunctionDefinition, GenerationParams, Message, MessageContent, ToolChoice, - ToolChoiceFunction, ToolChoiceFunctionName, ToolDefinition, + AnnotatedLlmRequest, ApiSpecificRequest, ContentPart, FunctionDefinition, GenerationParams, + Message, MessageContent, ProviderNativeComponent, ToolChoice, ToolChoiceFunction, + ToolChoiceFunctionName, ToolDefinition, }; use super::resolve::{ProviderSurface, ProviderSurfaceDescriptor}; use super::response::{ @@ -128,6 +129,14 @@ const MODELED_REQUEST_KEYS: &[&str] = &[ "tool_choice", "metadata", "service_tier", + "stream", + "cache_control", + "container", + "inference_geo", + "output_config", + "thinking", + "top_k", + "anthropic-user-profile-id", ]; /// Decode the Anthropic `tool_choice` JSON value into a normalized [`ToolChoice`]. @@ -157,153 +166,422 @@ fn decode_anthropic_tool_choice(val: &Json) -> Option { /// Extract Anthropic `disable_parallel_tool_use` from tool_choice and map /// to normalized `parallel_tool_calls` semantics. -fn decode_parallel_tool_calls(val: &Json) -> Option { - let obj = val.as_object()?; - obj.get("disable_parallel_tool_use") - .and_then(|v| v.as_bool()) - .map(|disabled| !disabled) +fn decode_parallel_tool_calls(val: &Json) -> Result> { + let Some(obj) = val.as_object() else { + return Ok(None); + }; + Ok(super::optional_bool( + obj, + "disable_parallel_tool_use", + "Anthropic Messages tool_choice", + )? + .map(|disabled| !disabled)) } /// Encode a normalized [`ToolChoice`] back into Anthropic JSON format. -fn encode_anthropic_tool_choice(tc: &ToolChoice) -> Json { +fn encode_anthropic_tool_choice(tc: &ToolChoice) -> Result { match tc { - ToolChoice::Auto => serde_json::json!({"type": "auto"}), - ToolChoice::Required => serde_json::json!({"type": "any"}), - ToolChoice::None => serde_json::json!({"type": "none"}), + ToolChoice::Auto => Ok(serde_json::json!({"type": "auto"})), + ToolChoice::Required => Ok(serde_json::json!({"type": "any"})), + ToolChoice::None => Ok(serde_json::json!({"type": "none"})), ToolChoice::Specific(func) => { - serde_json::json!({"type": "tool", "name": func.function.name}) + Ok(serde_json::json!({"type": "tool", "name": func.function.name})) + } + ToolChoice::ProviderNative(native) if native.provider == "anthropic_messages" => { + Ok(native.value.clone()) } + ToolChoice::ProviderNative(native) => Err(FlowError::InvalidArgument(format!( + "tool choice for {} cannot be encoded for Anthropic Messages", + native.provider + ))), } } fn encode_tool_choice_with_parallel_hint( tc: &ToolChoice, parallel_tool_calls: Option, -) -> Json { - let mut value = encode_anthropic_tool_choice(tc); +) -> Result { + let mut value = encode_anthropic_tool_choice(tc)?; if let (Some(parallel), Some(obj)) = (parallel_tool_calls, value.as_object_mut()) { obj.insert("disable_parallel_tool_use".into(), Json::Bool(!parallel)); } - value + Ok(value) } -/// Extract the system prompt from an Anthropic top-level `system` field. -/// -/// Handles both string and array-of-content-blocks formats. -fn extract_system_message(system_val: &Json) -> Option { - if let Some(s) = system_val.as_str() { - Some(Message::System { - content: MessageContent::Text(s.to_string()), - name: None, - }) - } else if let Some(arr) = system_val.as_array() { - // Array of content blocks -- extract text from each "text" block. - let texts: Vec<&str> = arr - .iter() - .filter_map(|block| { - let block_type = block.get("type")?.as_str()?; - if block_type == "text" { - block.get("text")?.as_str() - } else { - None - } - }) - .collect(); - if texts.is_empty() { - None - } else { - Some(Message::System { - content: MessageContent::Text(texts.join("\n")), - name: None, - }) - } - } else { - None +fn native_component(provider: &str, value: &Json) -> ProviderNativeComponent { + ProviderNativeComponent { + provider: provider.to_string(), + kind: value + .get("type") + .and_then(Json::as_str) + .unwrap_or("unknown") + .to_string(), + value: value.clone(), } } -/// Extract system text from a [`Message::System`] for encoding back to top-level. -fn extract_system_text(msg: &Message) -> Option { - match msg { - Message::System { - content: MessageContent::Text(s), - .. - } => Some(s.clone()), - Message::System { - content: MessageContent::Parts(parts), - .. - } => { - let texts: Vec<&str> = parts +fn decode_anthropic_content(value: &Json) -> Result { + if let Some(text) = value.as_str() { + return Ok(MessageContent::Text(text.to_string())); + } + let blocks = value.as_array().ok_or_else(|| { + FlowError::InvalidArgument("Anthropic Messages content must be a string or an array".into()) + })?; + let parts = blocks + .iter() + .map(decode_anthropic_content_part) + .collect::>>()?; + Ok(MessageContent::Parts(parts)) +} + +fn decode_anthropic_content_part(value: &Json) -> Result { + let obj = value.as_object().ok_or_else(|| { + FlowError::InvalidArgument("Anthropic Messages content block must be an object".into()) + })?; + let kind = obj.get("type").and_then(Json::as_str).unwrap_or("unknown"); + match kind { + "text" => { + let text = obj.get("text").and_then(Json::as_str).ok_or_else(|| { + FlowError::InvalidArgument("Anthropic text block is missing text".into()) + })?; + Ok(ContentPart::Text { + text: text.to_string(), + extra: obj + .iter() + .filter(|(key, _)| !matches!(key.as_str(), "type" | "text")) + .map(|(key, value)| (key.clone(), value.clone())) + .collect(), + }) + } + "image" => Ok(ContentPart::Image { + image: Json::Object( + obj.iter() + .filter(|(key, _)| key.as_str() != "type") + .map(|(key, value)| (key.clone(), value.clone())) + .collect(), + ), + extra: serde_json::Map::new(), + }), + "document" => Ok(ContentPart::File { + file: Json::Object( + obj.iter() + .filter(|(key, _)| key.as_str() != "type") + .map(|(key, value)| (key.clone(), value.clone())) + .collect(), + ), + extra: serde_json::Map::new(), + }), + "tool_use" => { + let id = obj.get("id").and_then(Json::as_str).ok_or_else(|| { + FlowError::InvalidArgument("Anthropic tool_use block is missing id".into()) + })?; + let name = obj.get("name").and_then(Json::as_str).ok_or_else(|| { + FlowError::InvalidArgument("Anthropic tool_use block is missing name".into()) + })?; + let input = obj.get("input").ok_or_else(|| { + FlowError::InvalidArgument("Anthropic tool_use block is missing input".into()) + })?; + let extra = obj .iter() - .filter_map(|p| match p { - super::request::ContentPart::Text { text } => Some(text.as_str()), - super::request::ContentPart::ImageUrl { .. } => None, + .filter(|(key, _)| !matches!(key.as_str(), "type" | "id" | "name" | "input")) + .map(|(key, value)| (key.clone(), value.clone())) + .collect(); + Ok(ContentPart::ToolUse { + id: id.to_string(), + name: name.to_string(), + input: input.clone(), + extra, + }) + } + "tool_result" => { + let tool_use_id = obj + .get("tool_use_id") + .and_then(Json::as_str) + .ok_or_else(|| { + FlowError::InvalidArgument( + "Anthropic tool_result block is missing tool_use_id".into(), + ) + })?; + let content = obj.get("content").ok_or_else(|| { + FlowError::InvalidArgument("Anthropic tool_result block is missing content".into()) + })?; + let is_error = match obj.get("is_error") { + Some(Json::Null) | None => None, + Some(value) => Some(value.as_bool().ok_or_else(|| { + FlowError::InvalidArgument( + "Anthropic tool_result is_error must be a boolean".into(), + ) + })?), + }; + let extra = obj + .iter() + .filter(|(key, _)| { + !matches!( + key.as_str(), + "type" | "tool_use_id" | "content" | "is_error" + ) }) + .map(|(key, value)| (key.clone(), value.clone())) .collect(); - if texts.is_empty() { - None - } else { - Some(texts.join("\n")) - } + Ok(ContentPart::ToolResult { + tool_use_id: tool_use_id.to_string(), + content: content.clone(), + is_error, + extra, + }) } - _ => None, + _ => { + let native = native_component("anthropic_messages", value); + Ok(ContentPart::ProviderNative { + provider: native.provider, + kind: native.kind, + value: native.value, + }) + } + } +} + +fn decode_anthropic_message(value: &Json) -> Result { + let obj = value.as_object().ok_or_else(|| { + FlowError::InvalidArgument("Anthropic Messages message must be an object".into()) + })?; + let role = obj.get("role").and_then(Json::as_str).ok_or_else(|| { + FlowError::InvalidArgument("Anthropic Messages message is missing role".into()) + })?; + let content = obj.get("content").ok_or_else(|| { + FlowError::InvalidArgument("Anthropic Messages message is missing content".into()) + })?; + if obj + .keys() + .any(|key| !matches!(key.as_str(), "role" | "content")) + { + return Ok(Message::ProviderNative { + provider: "anthropic_messages".into(), + kind: role.to_string(), + value: value.clone(), + }); + } + let content = decode_anthropic_content(content)?; + match role { + "user" => Ok(Message::User { + content, + name: None, + }), + "assistant" => Ok(Message::Assistant { + content: Some(content), + tool_calls: None, + name: None, + }), + "system" => Ok(Message::System { + content, + name: None, + }), + _ => Ok(Message::ProviderNative { + provider: "anthropic_messages".into(), + kind: role.to_string(), + value: value.clone(), + }), } } -fn split_system_and_messages(messages: &[Message]) -> (Option, Vec<&Message>) { - let mut system_text = None; - let mut non_system_messages = Vec::new(); +fn encode_anthropic_content(content: &MessageContent) -> Result { + match content { + MessageContent::Text(text) => Ok(Json::String(text.clone())), + MessageContent::Parts(parts) => Ok(Json::Array( + parts + .iter() + .map(encode_anthropic_content_part) + .collect::>>()?, + )), + } +} - for msg in messages { - if let Some(text) = extract_system_text(msg) { - system_text = Some(text); - } else { - non_system_messages.push(msg); +fn encode_anthropic_content_part(part: &ContentPart) -> Result { + match part { + ContentPart::Text { text, extra } => { + let mut obj = extra.clone(); + obj.insert("type".into(), Json::String("text".into())); + obj.insert("text".into(), Json::String(text.clone())); + Ok(Json::Object(obj)) + } + ContentPart::Image { image, extra } => { + let mut obj = image.as_object().cloned().ok_or_else(|| { + FlowError::InvalidArgument("Anthropic image payload must be an object".into()) + })?; + obj.extend(extra.clone()); + obj.insert("type".into(), Json::String("image".into())); + Ok(Json::Object(obj)) + } + ContentPart::File { file, extra } => { + let mut obj = file.as_object().cloned().ok_or_else(|| { + FlowError::InvalidArgument("Anthropic document payload must be an object".into()) + })?; + obj.extend(extra.clone()); + obj.insert("type".into(), Json::String("document".into())); + Ok(Json::Object(obj)) + } + ContentPart::ToolUse { + id, + name, + input, + extra, + } => { + let mut obj = extra.clone(); + obj.insert("type".into(), Json::String("tool_use".into())); + obj.insert("id".into(), Json::String(id.clone())); + obj.insert("name".into(), Json::String(name.clone())); + obj.insert("input".into(), input.clone()); + Ok(Json::Object(obj)) + } + ContentPart::ToolResult { + tool_use_id, + content, + is_error, + extra, + } => { + let mut obj = extra.clone(); + obj.insert("type".into(), Json::String("tool_result".into())); + obj.insert("tool_use_id".into(), Json::String(tool_use_id.clone())); + obj.insert("content".into(), content.clone()); + if let Some(is_error) = is_error { + obj.insert("is_error".into(), Json::Bool(*is_error)); + } + Ok(Json::Object(obj)) } + ContentPart::ProviderNative { + provider, value, .. + } if provider == "anthropic_messages" => Ok(value.clone()), + other => Err(FlowError::InvalidArgument(format!( + "content part {other:?} cannot be encoded for Anthropic Messages" + ))), } +} - (system_text, non_system_messages) +fn encode_anthropic_message(message: &Message) -> Result { + match message { + Message::User { content, .. } | Message::System { content, .. } => { + let role = if matches!(message, Message::User { .. }) { + "user" + } else { + "system" + }; + let mut obj = serde_json::Map::new(); + obj.insert("role".into(), Json::String(role.into())); + obj.insert("content".into(), encode_anthropic_content(content)?); + Ok(Json::Object(obj)) + } + Message::Assistant { content, .. } => { + let mut obj = serde_json::Map::new(); + obj.insert("role".into(), Json::String("assistant".into())); + obj.insert( + "content".into(), + match content { + Some(content) => encode_anthropic_content(content)?, + None => Json::Array(Vec::new()), + }, + ); + Ok(Json::Object(obj)) + } + Message::ProviderNative { + provider, value, .. + } if provider == "anthropic_messages" => Ok(value.clone()), + other => Err(FlowError::InvalidArgument(format!( + "message {other:?} cannot be encoded for Anthropic Messages" + ))), + } } -fn insert_serialized( - obj: &mut serde_json::Map, - key: &str, - value: &T, - context: &str, -) -> Result<()> { - let json = serde_json::to_value(value) - .map_err(|e| FlowError::Internal(format!("Anthropic Messages {context} encode: {e}")))?; - obj.insert(key.into(), json); - Ok(()) +fn encode_anthropic_tool(tool: &ToolDefinition) -> Result { + match tool { + ToolDefinition::Function { function, extra } => { + let mut obj = extra.clone(); + obj.insert("name".into(), Json::String(function.name.clone())); + if let Some(description) = &function.description { + obj.insert("description".into(), Json::String(description.clone())); + } + if let Some(parameters) = &function.parameters { + obj.insert("input_schema".into(), parameters.clone()); + } + if let Some(strict) = function.strict { + obj.insert("strict".into(), Json::Bool(strict)); + } + obj.extend(function.extra.clone()); + Ok(Json::Object(obj)) + } + ToolDefinition::ProviderNative { + provider, value, .. + } if provider == "anthropic_messages" => Ok(value.clone()), + other => Err(FlowError::InvalidArgument(format!( + "tool {other:?} cannot be encoded for Anthropic Messages" + ))), + } } -fn overlay_generation_params(obj: &mut serde_json::Map, params: &GenerationParams) { - if let Some(temp) = params.temperature { - obj.insert("temperature".into(), json_f64(temp)); +fn decode_anthropic_tool(value: &Json) -> Result { + let obj = value.as_object().ok_or_else(|| { + FlowError::InvalidArgument("Anthropic Messages tool must be an object".into()) + })?; + let is_client_tool = + obj.get("type").is_none() && obj.contains_key("name") && obj.contains_key("input_schema"); + if !is_client_tool { + let native = native_component("anthropic_messages", value); + return Ok(ToolDefinition::ProviderNative { + provider: native.provider, + kind: native.kind, + value: native.value, + }); } - if let Some(top_p) = params.top_p { - obj.insert("top_p".into(), json_f64(top_p)); + + let function_extra = serde_json::Map::new(); + let wrapper_extra = obj + .iter() + .filter(|(key, _)| { + !matches!( + key.as_str(), + "name" | "description" | "input_schema" | "strict" + ) + }) + .map(|(key, value)| (key.clone(), value.clone())) + .collect(); + let name = + super::optional_string(obj, "name", "Anthropic Messages tool")?.ok_or_else(|| { + FlowError::InvalidArgument("Anthropic Messages tool is missing name".into()) + })?; + let description = super::optional_string(obj, "description", "Anthropic Messages tool")?; + let strict = super::optional_bool(obj, "strict", "Anthropic Messages tool")?; + Ok(ToolDefinition::Function { + function: FunctionDefinition { + name, + description, + parameters: obj.get("input_schema").cloned(), + strict, + extra: function_extra, + }, + extra: wrapper_extra, + }) +} + +fn patch_extra_fields( + obj: &mut serde_json::Map, + baseline: &serde_json::Map, + edited: &serde_json::Map, +) { + for key in baseline.keys().filter(|key| !edited.contains_key(*key)) { + obj.remove(key); } - if let Some(max_tokens) = params.max_tokens { - obj.insert("max_tokens".into(), Json::from(max_tokens)); + for (key, value) in edited { + if baseline.get(key) != Some(value) { + obj.insert(key.clone(), value.clone()); + } } } -fn encode_anthropic_tools(tools: &[ToolDefinition]) -> Vec { - tools - .iter() - .map(|td| { - let mut tool = serde_json::Map::new(); - tool.insert("name".into(), Json::String(td.function.name.clone())); - if let Some(ref desc) = td.function.description { - tool.insert("description".into(), Json::String(desc.clone())); - } - if let Some(ref params) = td.function.parameters { - tool.insert("input_schema".into(), params.clone()); - } - Json::Object(tool) - }) - .collect() +fn set_or_remove_json(obj: &mut serde_json::Map, key: &str, value: Option) { + if let Some(value) = value { + obj.insert(key.into(), value); + } else { + obj.remove(key); + } } fn anthropic_text_message(content_blocks: Option<&[Json]>) -> Option { @@ -431,33 +709,36 @@ impl LlmCodec for AnthropicMessagesCodec { .content .as_object() .ok_or_else(|| FlowError::Internal("request content is not an object".into()))?; - - // Extract system from top-level field. - let system_msg = obj.get("system").and_then(extract_system_message); - - // Extract messages (default to empty vec if absent). - let mut messages: Vec = obj - .get("messages") - .map(|v| serde_json::from_value(v.clone()).unwrap_or_default()) - .unwrap_or_default(); - - // Prepend system message if present. - if let Some(sys) = system_msg { - messages.insert(0, sys); - } - - // Extract model. - let model = obj.get("model").and_then(|v| v.as_str()).map(String::from); - - // Extract generation params. - let temperature = obj.get("temperature").and_then(|v| v.as_f64()); - let top_p = obj.get("top_p").and_then(|v| v.as_f64()); - let max_tokens = obj.get("max_tokens").and_then(|v| v.as_u64()); - // Anthropic uses stop_sequences (not stop). + let raw_messages = obj.get("messages").ok_or_else(|| { + FlowError::InvalidArgument("Anthropic Messages request is missing messages".into()) + })?; + let messages = raw_messages + .as_array() + .ok_or_else(|| { + FlowError::InvalidArgument("Anthropic Messages messages must be an array".into()) + })? + .iter() + .map(decode_anthropic_message) + .collect::>>()?; + let instructions = obj + .get("system") + .map(decode_anthropic_content) + .transpose()?; + let model = super::optional_string(obj, "model", "Anthropic Messages")?; + let temperature = super::optional_f64(obj, "temperature", "Anthropic Messages")?; + let top_p = super::optional_f64(obj, "top_p", "Anthropic Messages")?; + let max_tokens = super::optional_u64(obj, "max_tokens", "Anthropic Messages")?; let stop = obj .get("stop_sequences") - .and_then(|v| serde_json::from_value::>(v.clone()).ok()); - + .filter(|value| !value.is_null()) + .map(|value| { + serde_json::from_value::>(value.clone()).map_err(|error| { + FlowError::InvalidArgument(format!( + "invalid Anthropic stop_sequences value: {error}" + )) + }) + }) + .transpose()?; let params = if temperature.is_some() || max_tokens.is_some() || top_p.is_some() || stop.is_some() { Some(GenerationParams { @@ -469,48 +750,50 @@ impl LlmCodec for AnthropicMessagesCodec { } else { None }; - - // Extract tools: Anthropic uses flat structure (name, description, input_schema). - // Normalize to ToolDefinition { type: "function", function: { name, description, parameters } }. - let tools: Option> = obj.get("tools").and_then(|v| { - let arr = v.as_array()?; - let defs: Vec = arr - .iter() - .filter_map(|tool| { - let name = tool.get("name")?.as_str()?.to_string(); - let description = tool - .get("description") - .and_then(|d| d.as_str()) - .map(String::from); - let parameters = tool.get("input_schema").cloned(); - Some(ToolDefinition { - tool_type: "function".into(), - function: FunctionDefinition { - name, - description, - parameters, - }, - }) - }) - .collect(); - if defs.is_empty() { None } else { Some(defs) } + let tools = obj + .get("tools") + .map(|value| { + value + .as_array() + .ok_or_else(|| { + FlowError::InvalidArgument( + "Anthropic Messages tools must be an array".into(), + ) + })? + .iter() + .map(decode_anthropic_tool) + .collect::>>() + }) + .transpose()?; + let tool_choice = obj.get("tool_choice").map(|value| { + decode_anthropic_tool_choice(value).unwrap_or_else(|| { + ToolChoice::ProviderNative(native_component("anthropic_messages", value)) + }) }); - - // Extract tool_choice: Anthropic format. - let tool_choice = obj + let parallel_tool_calls = obj .get("tool_choice") - .and_then(decode_anthropic_tool_choice); - let parallel_tool_calls = obj.get("tool_choice").and_then(decode_parallel_tool_calls); - - // Collect extra fields (keys not in MODELED_REQUEST_KEYS). + .map(decode_parallel_tool_calls) + .transpose()? + .flatten(); + let service_tier = super::optional_string(obj, "service_tier", "Anthropic Messages")?; + let stream = super::optional_bool(obj, "stream", "Anthropic Messages")?; + let container = super::optional_string(obj, "container", "Anthropic Messages")?; + let inference_geo = super::optional_string(obj, "inference_geo", "Anthropic Messages")?; + let top_k = super::optional_u64(obj, "top_k", "Anthropic Messages")?; + let user_profile_id = + super::optional_string(obj, "anthropic-user-profile-id", "Anthropic Messages")?; + let metadata = super::optional_object(obj, "metadata", "Anthropic Messages")?; + let cache_control = super::optional_object(obj, "cache_control", "Anthropic Messages")?; + let output_config = super::optional_object(obj, "output_config", "Anthropic Messages")?; + let thinking = super::optional_object(obj, "thinking", "Anthropic Messages")?; let extra: serde_json::Map = obj .iter() .filter(|(k, _)| !MODELED_REQUEST_KEYS.contains(&k.as_str())) .map(|(k, v)| (k.clone(), v.clone())) .collect(); - Ok(AnnotatedLlmRequest { messages, + instructions, model, params, tools, @@ -521,75 +804,239 @@ impl LlmCodec for AnthropicMessagesCodec { reasoning: None, include: None, user: None, - metadata: obj.get("metadata").cloned(), - service_tier: obj - .get("service_tier") - .and_then(|v| v.as_str()) - .map(String::from), + metadata, + service_tier, parallel_tool_calls, max_output_tokens: None, max_tool_calls: None, top_logprobs: None, - stream: None, + stream, + api_specific: Some(ApiSpecificRequest::AnthropicMessages { + cache_control, + container, + inference_geo, + output_config, + thinking, + top_k, + user_profile_id, + }), extra, }) } fn encode(&self, annotated: &AnnotatedLlmRequest, original: &LlmRequest) -> Result { + let baseline = self.decode(original)?; let mut content = original.content.clone(); let obj = content .as_object_mut() .ok_or_else(|| FlowError::Internal("original content is not an object".into()))?; - - let (system_text, non_system_messages) = split_system_and_messages(&annotated.messages); - - if let Some(text) = system_text { - obj.insert("system".into(), Json::String(text)); + if annotated.messages != baseline.messages { + let original_messages = obj.get("messages").and_then(Json::as_array); + let messages = super::encode_changed_items( + &annotated.messages, + &baseline.messages, + original_messages.map(Vec::as_slice), + encode_anthropic_message, + )?; + obj.insert("messages".into(), Json::Array(messages)); } - - // Overlay messages (non-system only). - insert_serialized(obj, "messages", &non_system_messages, "messages")?; - - // Overlay model if present. - if let Some(ref model) = annotated.model { - obj.insert("model".into(), Json::String(model.clone())); + if annotated.instructions != baseline.instructions { + set_or_remove_json( + obj, + "system", + annotated + .instructions + .as_ref() + .map(encode_anthropic_content) + .transpose()?, + ); } - - // Overlay generation params. - if let Some(ref params) = annotated.params { - overlay_generation_params(obj, params); - // Write stop_sequences (Anthropic key name, not "stop"). - if let Some(ref stop) = params.stop { - insert_serialized(obj, "stop_sequences", stop, "stop_sequences")?; + if annotated.model != baseline.model { + set_or_remove_json(obj, "model", annotated.model.clone().map(Json::String)); + } + if annotated.params != baseline.params { + let edited = annotated.params.as_ref(); + let before = baseline.params.as_ref(); + if edited.and_then(|p| p.temperature) != before.and_then(|p| p.temperature) { + set_or_remove_json( + obj, + "temperature", + edited.and_then(|p| p.temperature).map(json_f64), + ); + } + if edited.and_then(|p| p.top_p) != before.and_then(|p| p.top_p) { + set_or_remove_json(obj, "top_p", edited.and_then(|p| p.top_p).map(json_f64)); + } + if edited.and_then(|p| p.max_tokens) != before.and_then(|p| p.max_tokens) { + set_or_remove_json( + obj, + "max_tokens", + edited.and_then(|p| p.max_tokens).map(Json::from), + ); + } + if edited.and_then(|p| p.stop.as_ref()) != before.and_then(|p| p.stop.as_ref()) { + set_or_remove_json( + obj, + "stop_sequences", + edited + .and_then(|p| p.stop.as_ref()) + .map(|stop| serde_json::json!(stop)), + ); } } - - // Overlay tools in Anthropic format: { name, description, input_schema }. - // Denormalize from ToolDefinition (drop type/function wrapper, rename parameters -> input_schema). - if let Some(ref tools) = annotated.tools { - let anthropic_tools = encode_anthropic_tools(tools); - insert_serialized(obj, "tools", &anthropic_tools, "tools")?; + if annotated.tools != baseline.tools { + let tools = annotated + .tools + .as_deref() + .map(|tools| { + super::encode_changed_items( + tools, + baseline.tools.as_deref().unwrap_or(&[]), + obj.get("tools").and_then(Json::as_array).map(Vec::as_slice), + encode_anthropic_tool, + ) + }) + .transpose()? + .map(Json::Array); + set_or_remove_json(obj, "tools", tools); } - - // Overlay tool_choice in Anthropic format. - if let Some(ref tool_choice) = annotated.tool_choice { - obj.insert( - "tool_choice".into(), - encode_tool_choice_with_parallel_hint(tool_choice, annotated.parallel_tool_calls), + if annotated.tool_choice != baseline.tool_choice + || annotated.parallel_tool_calls != baseline.parallel_tool_calls + { + let tool_choice = match (&annotated.tool_choice, &baseline.tool_choice) { + (Some(edited), Some(before)) => { + let edited = encode_tool_choice_with_parallel_hint( + edited, + annotated.parallel_tool_calls, + )?; + let before = encode_tool_choice_with_parallel_hint( + before, + baseline.parallel_tool_calls, + )?; + Some(match obj.get("tool_choice") { + Some(original) => super::patch_changed_json(original, &before, &edited)?, + None => edited, + }) + } + (Some(edited), None) => Some(encode_tool_choice_with_parallel_hint( + edited, + annotated.parallel_tool_calls, + )?), + (None, _) => annotated + .parallel_tool_calls + .map(|parallel| { + encode_tool_choice_with_parallel_hint(&ToolChoice::Auto, Some(parallel)) + }) + .transpose()?, + }; + set_or_remove_json(obj, "tool_choice", tool_choice); + } + for (key, edited, before) in [("metadata", &annotated.metadata, &baseline.metadata)] { + if edited != before { + set_or_remove_json(obj, key, edited.clone()); + } + } + if annotated.service_tier != baseline.service_tier { + set_or_remove_json( + obj, + "service_tier", + annotated.service_tier.clone().map(Json::String), ); } - - if let Some(ref metadata) = annotated.metadata { - obj.insert("metadata".into(), metadata.clone()); + if annotated.stream != baseline.stream { + set_or_remove_json(obj, "stream", annotated.stream.map(Json::Bool)); } - if let Some(ref service_tier) = annotated.service_tier { - obj.insert("service_tier".into(), Json::String(service_tier.clone())); + + if annotated.store != baseline.store + || annotated.previous_response_id != baseline.previous_response_id + || annotated.truncation != baseline.truncation + || annotated.reasoning != baseline.reasoning + || annotated.include != baseline.include + || annotated.user != baseline.user + || annotated.max_output_tokens != baseline.max_output_tokens + || annotated.max_tool_calls != baseline.max_tool_calls + || annotated.top_logprobs != baseline.top_logprobs + { + return Err(FlowError::InvalidArgument( + "request contains fields that cannot be encoded for Anthropic Messages".into(), + )); } - // Merge extra fields back. - for (k, v) in &annotated.extra { - obj.insert(k.clone(), v.clone()); + match (&annotated.api_specific, &baseline.api_specific) { + ( + Some(ApiSpecificRequest::AnthropicMessages { + cache_control, + container, + inference_geo, + output_config, + thinking, + top_k, + user_profile_id, + }), + Some(ApiSpecificRequest::AnthropicMessages { + cache_control: old_cache_control, + container: old_container, + inference_geo: old_inference_geo, + output_config: old_output_config, + thinking: old_thinking, + top_k: old_top_k, + user_profile_id: old_user_profile_id, + }), + ) => { + if cache_control != old_cache_control { + set_or_remove_json(obj, "cache_control", cache_control.clone()); + } + for (key, edited, before) in [ + ("container", container, old_container), + ("inference_geo", inference_geo, old_inference_geo), + ( + "anthropic-user-profile-id", + user_profile_id, + old_user_profile_id, + ), + ] { + if edited != before { + set_or_remove_json(obj, key, edited.clone().map(Json::String)); + } + } + for (key, edited, before) in [ + ("output_config", output_config, old_output_config), + ("thinking", thinking, old_thinking), + ] { + if edited != before { + set_or_remove_json(obj, key, edited.clone()); + } + } + if top_k != old_top_k { + set_or_remove_json(obj, "top_k", top_k.map(Json::from)); + } + } + (None, Some(ApiSpecificRequest::AnthropicMessages { .. })) => { + for key in [ + "cache_control", + "container", + "inference_geo", + "output_config", + "thinking", + "top_k", + "anthropic-user-profile-id", + ] { + obj.remove(key); + } + } + (Some(_), _) => { + return Err(FlowError::InvalidArgument( + "api_specific provider does not match Anthropic Messages".into(), + )); + } + (None, Some(_)) => { + return Err(FlowError::InvalidArgument( + "api_specific provider does not match Anthropic Messages".into(), + )); + } + (None, None) => {} } + patch_extra_fields(obj, &baseline.extra, &annotated.extra); Ok(LlmRequest { headers: original.headers.clone(), diff --git a/crates/core/src/codec/mod.rs b/crates/core/src/codec/mod.rs index 8c103f97d..a92f00826 100644 --- a/crates/core/src/codec/mod.rs +++ b/crates/core/src/codec/mod.rs @@ -25,6 +25,281 @@ pub mod response; pub mod streaming; pub mod traits; +use nemo_relay_types::Json; + +use crate::error::{FlowError, Result}; + +fn optional_bool( + obj: &serde_json::Map, + key: &str, + surface: &str, +) -> Result> { + match obj.get(key) { + Some(Json::Null) | None => Ok(None), + Some(Json::Bool(value)) => Ok(Some(*value)), + Some(_) => Err(FlowError::InvalidArgument(format!( + "{surface} {key} must be a boolean or null" + ))), + } +} + +fn optional_u64( + obj: &serde_json::Map, + key: &str, + surface: &str, +) -> Result> { + match obj.get(key) { + Some(Json::Null) | None => Ok(None), + Some(value) if value.as_u64().is_some() => Ok(value.as_u64()), + Some(_) => Err(FlowError::InvalidArgument(format!( + "{surface} {key} must be a non-negative integer or null" + ))), + } +} + +fn optional_i64( + obj: &serde_json::Map, + key: &str, + surface: &str, +) -> Result> { + match obj.get(key) { + Some(Json::Null) | None => Ok(None), + Some(value) if value.as_i64().is_some() => Ok(value.as_i64()), + Some(_) => Err(FlowError::InvalidArgument(format!( + "{surface} {key} must be an integer or null" + ))), + } +} + +fn optional_f64( + obj: &serde_json::Map, + key: &str, + surface: &str, +) -> Result> { + match obj.get(key) { + Some(Json::Null) | None => Ok(None), + Some(value) if value.as_f64().is_some() => Ok(value.as_f64()), + Some(_) => Err(FlowError::InvalidArgument(format!( + "{surface} {key} must be a number or null" + ))), + } +} + +fn optional_string( + obj: &serde_json::Map, + key: &str, + surface: &str, +) -> Result> { + match obj.get(key) { + Some(Json::Null) | None => Ok(None), + Some(Json::String(value)) => Ok(Some(value.clone())), + Some(_) => Err(FlowError::InvalidArgument(format!( + "{surface} {key} must be a string or null" + ))), + } +} + +fn optional_object( + obj: &serde_json::Map, + key: &str, + surface: &str, +) -> Result> { + match obj.get(key) { + Some(Json::Null) | None => Ok(None), + Some(value @ Json::Object(_)) => Ok(Some(value.clone())), + Some(_) => Err(FlowError::InvalidArgument(format!( + "{surface} {key} must be an object or null" + ))), + } +} + +fn optional_array( + obj: &serde_json::Map, + key: &str, + surface: &str, +) -> Result> { + match obj.get(key) { + Some(Json::Null) | None => Ok(None), + Some(value @ Json::Array(_)) => Ok(Some(value.clone())), + Some(_) => Err(FlowError::InvalidArgument(format!( + "{surface} {key} must be an array or null" + ))), + } +} + +fn encode_changed_items( + edited: &[T], + baseline: &[T], + original: Option<&[Json]>, + encode: F, +) -> Result> +where + T: PartialEq, + F: FnMut(&T) -> Result, +{ + encode_changed_items_with_patch( + edited, + baseline, + original, + encode, + |original, _, _, baseline_value, edited_value| { + patch_changed_json(original, baseline_value, edited_value) + }, + ) +} + +fn encode_changed_items_with_patch( + edited: &[T], + baseline: &[T], + original: Option<&[Json]>, + mut encode: F, + mut patch: P, +) -> Result> +where + T: PartialEq, + F: FnMut(&T) -> Result, + P: FnMut(&Json, &T, &T, &Json, &Json) -> Result, +{ + let alignment = if original.is_some() { + align_changed_items(edited, baseline)? + } else { + vec![None; edited.len()] + }; + edited + .iter() + .enumerate() + .map(|(index, item)| { + if let Some(baseline_index) = alignment[index] { + let baseline_item = &baseline[baseline_index]; + if let Some(original_item) = original.and_then(|items| items.get(baseline_index)) { + if baseline_item == item { + return Ok(original_item.clone()); + } + let baseline_value = encode(baseline_item)?; + let edited_value = encode(item)?; + return patch( + original_item, + baseline_item, + item, + &baseline_value, + &edited_value, + ); + } + } + encode(item) + }) + .collect() +} + +fn align_changed_items(edited: &[T], baseline: &[T]) -> Result>> { + let mut alignment = vec![None; edited.len()]; + let mut used = vec![false; baseline.len()]; + + for index in 0..edited.len().min(baseline.len()) { + if edited[index] == baseline[index] { + alignment[index] = Some(index); + used[index] = true; + } + } + + let mut structurally_changed = edited.len() != baseline.len(); + for (edited_index, item) in edited.iter().enumerate() { + if alignment[edited_index].is_some() { + continue; + } + if let Some(baseline_index) = baseline + .iter() + .enumerate() + .position(|(baseline_index, candidate)| !used[baseline_index] && candidate == item) + { + alignment[edited_index] = Some(baseline_index); + used[baseline_index] = true; + structurally_changed |= baseline_index != edited_index; + } + } + + let unmatched_edited = alignment + .iter() + .enumerate() + .filter_map(|(index, item)| item.is_none().then_some(index)) + .collect::>(); + let unmatched_baseline = used + .iter() + .enumerate() + .filter_map(|(index, item)| (!item).then_some(index)) + .collect::>(); + + if unmatched_edited.len() > 1 && unmatched_baseline.len() > 1 { + return Err(FlowError::InvalidArgument( + "cannot safely preserve provider fields for multiple edited array items without stable identities" + .into(), + )); + } + + if !structurally_changed && edited.len() == baseline.len() { + for index in unmatched_edited { + alignment[index] = Some(index); + } + } + + Ok(alignment) +} + +fn patch_changed_array(original: &[Json], baseline: &[Json], edited: &[Json]) -> Result> { + let alignment = align_changed_items(edited, baseline)?; + edited + .iter() + .enumerate() + .map(|(edited_index, edited_value)| { + if let Some(baseline_index) = alignment[edited_index] + && let (Some(original_value), Some(baseline_value)) = + (original.get(baseline_index), baseline.get(baseline_index)) + { + return patch_changed_json(original_value, baseline_value, edited_value); + } + Ok(edited_value.clone()) + }) + .collect() +} + +fn patch_changed_json(original: &Json, baseline: &Json, edited: &Json) -> Result { + if baseline == edited { + return Ok(original.clone()); + } + + match (original, baseline, edited) { + (Json::Object(original), Json::Object(baseline), Json::Object(edited)) => { + if ["type", "role"] + .into_iter() + .any(|key| baseline.get(key) != edited.get(key)) + { + return Ok(Json::Object(edited.clone())); + } + let mut patched = original.clone(); + for key in baseline.keys().filter(|key| !edited.contains_key(*key)) { + patched.remove(key); + } + for (key, edited_value) in edited { + if baseline.get(key) == Some(edited_value) { + continue; + } + let value = match (original.get(key), baseline.get(key)) { + (Some(original_value), Some(baseline_value)) => { + patch_changed_json(original_value, baseline_value, edited_value)? + } + _ => edited_value.clone(), + }; + patched.insert(key.clone(), value); + } + Ok(Json::Object(patched)) + } + (Json::Array(original), Json::Array(baseline), Json::Array(edited)) => Ok(Json::Array( + patch_changed_array(original, baseline, edited)?, + )), + _ => Ok(edited.clone()), + } +} + #[cfg(test)] #[path = "../../tests/unit/codec/parity_tests.rs"] mod parity_tests; diff --git a/crates/core/src/codec/openai_chat.rs b/crates/core/src/codec/openai_chat.rs index 50be0b14e..99e94e2bf 100644 --- a/crates/core/src/codec/openai_chat.rs +++ b/crates/core/src/codec/openai_chat.rs @@ -12,7 +12,11 @@ use crate::api::llm::LlmRequest; use crate::error::{FlowError, Result}; use crate::json::Json; -use super::request::{AnnotatedLlmRequest, GenerationParams, Message, ToolChoice, ToolDefinition}; +use super::request::{ + AnnotatedLlmRequest, ApiSpecificRequest, ContentPart, FunctionCall, FunctionDefinition, + GenerationParams, Message, MessageContent, OpenAiImageUrl, ProviderNativeComponent, ToolCall, + ToolChoice, ToolDefinition, +}; use super::resolve::{ProviderSurface, ProviderSurfaceDescriptor}; use super::response::{ AnnotatedLlmResponse, ApiSpecificResponse, FinishReason, RawUsageCost, ResponseToolCall, Usage, @@ -136,8 +140,572 @@ const MODELED_REQUEST_KEYS: &[&str] = &[ "parallel_tool_calls", "top_logprobs", "stream", + "audio", + "frequency_penalty", + "function_call", + "functions", + "logit_bias", + "logprobs", + "modalities", + "moderation", + "n", + "prediction", + "presence_penalty", + "prompt_cache_key", + "prompt_cache_options", + "prompt_cache_retention", + "reasoning_effort", + "response_format", + "safety_identifier", + "seed", + "stream_options", + "verbosity", + "web_search_options", ]; +fn chat_native(kind: &str, value: &Json) -> ProviderNativeComponent { + ProviderNativeComponent { + provider: "openai_chat".into(), + kind: kind.to_string(), + value: value.clone(), + } +} + +fn decode_chat_content(value: &Json) -> Result { + if let Some(text) = value.as_str() { + return Ok(MessageContent::Text(text.to_string())); + } + let parts = value.as_array().ok_or_else(|| { + FlowError::InvalidArgument("OpenAI Chat message content must be a string or array".into()) + })?; + Ok(MessageContent::Parts( + parts + .iter() + .map(decode_chat_content_part) + .collect::>>()?, + )) +} + +fn decode_chat_content_part(value: &Json) -> Result { + let obj = value.as_object().ok_or_else(|| { + FlowError::InvalidArgument("OpenAI Chat content part must be an object".into()) + })?; + let kind = obj.get("type").and_then(Json::as_str).unwrap_or("unknown"); + match kind { + "text" => Ok(ContentPart::Text { + text: obj + .get("text") + .and_then(Json::as_str) + .ok_or_else(|| { + FlowError::InvalidArgument("OpenAI Chat text part is missing text".into()) + })? + .to_string(), + extra: obj + .iter() + .filter(|(key, _)| !matches!(key.as_str(), "type" | "text")) + .map(|(key, value)| (key.clone(), value.clone())) + .collect(), + }), + "image_url" => { + let image_url: OpenAiImageUrl = + serde_json::from_value(obj.get("image_url").cloned().ok_or_else(|| { + FlowError::InvalidArgument("OpenAI Chat image part is missing image_url".into()) + })?) + .map_err(|error| { + FlowError::InvalidArgument(format!("invalid OpenAI Chat image_url: {error}")) + })?; + Ok(ContentPart::ImageUrl { + image_url, + extra: obj + .iter() + .filter(|(key, _)| !matches!(key.as_str(), "type" | "image_url")) + .map(|(key, value)| (key.clone(), value.clone())) + .collect(), + }) + } + "input_audio" => Ok(ContentPart::Audio { + audio: obj.get("input_audio").cloned().ok_or_else(|| { + FlowError::InvalidArgument("OpenAI Chat audio part is missing input_audio".into()) + })?, + extra: obj + .iter() + .filter(|(key, _)| !matches!(key.as_str(), "type" | "input_audio")) + .map(|(key, value)| (key.clone(), value.clone())) + .collect(), + }), + "file" => Ok(ContentPart::File { + file: obj.get("file").cloned().ok_or_else(|| { + FlowError::InvalidArgument("OpenAI Chat file part is missing file".into()) + })?, + extra: obj + .iter() + .filter(|(key, _)| !matches!(key.as_str(), "type" | "file")) + .map(|(key, value)| (key.clone(), value.clone())) + .collect(), + }), + "refusal" => Ok(ContentPart::Refusal { + refusal: obj + .get("refusal") + .and_then(Json::as_str) + .ok_or_else(|| { + FlowError::InvalidArgument("OpenAI Chat refusal part is missing refusal".into()) + })? + .to_string(), + extra: obj + .iter() + .filter(|(key, _)| !matches!(key.as_str(), "type" | "refusal")) + .map(|(key, value)| (key.clone(), value.clone())) + .collect(), + }), + _ => Ok(ContentPart::ProviderNative { + provider: "openai_chat".into(), + kind: kind.to_string(), + value: value.clone(), + }), + } +} + +fn optional_chat_string( + obj: &serde_json::Map, + key: &str, + context: &str, +) -> Result> { + match obj.get(key) { + Some(Json::Null) | None => Ok(None), + Some(Json::String(value)) => Ok(Some(value.clone())), + Some(_) => Err(FlowError::InvalidArgument(format!( + "OpenAI Chat {context} {key} must be a string or null" + ))), + } +} + +fn decode_chat_tool_call(value: &Json) -> Result> { + let obj = value.as_object().ok_or_else(|| { + FlowError::InvalidArgument("OpenAI Chat tool call must be an object".into()) + })?; + if obj.get("type").and_then(Json::as_str) != Some("function") { + return Ok(None); + } + if obj + .keys() + .any(|key| !matches!(key.as_str(), "id" | "type" | "function")) + { + return Ok(None); + } + let function = obj + .get("function") + .and_then(Json::as_object) + .ok_or_else(|| { + FlowError::InvalidArgument("OpenAI Chat function tool call is missing function".into()) + })?; + if function + .keys() + .any(|key| !matches!(key.as_str(), "name" | "arguments")) + { + return Ok(None); + } + let id = obj.get("id").and_then(Json::as_str).ok_or_else(|| { + FlowError::InvalidArgument("OpenAI Chat function tool call is missing id".into()) + })?; + let name = function.get("name").and_then(Json::as_str).ok_or_else(|| { + FlowError::InvalidArgument("OpenAI Chat function tool call is missing name".into()) + })?; + let arguments = function + .get("arguments") + .and_then(Json::as_str) + .ok_or_else(|| { + FlowError::InvalidArgument("OpenAI Chat function tool call is missing arguments".into()) + })?; + Ok(Some(ToolCall { + id: id.to_string(), + call_type: "function".into(), + function: FunctionCall { + name: name.to_string(), + arguments: arguments.to_string(), + }, + })) +} + +fn decode_chat_message(value: &Json) -> Result { + let obj = value.as_object().ok_or_else(|| { + FlowError::InvalidArgument("OpenAI Chat message must be an object".into()) + })?; + let role = obj + .get("role") + .and_then(Json::as_str) + .ok_or_else(|| FlowError::InvalidArgument("OpenAI Chat message is missing role".into()))?; + let native = || Message::ProviderNative { + provider: "openai_chat".into(), + kind: role.to_string(), + value: value.clone(), + }; + match role { + "system" | "developer" | "user" => { + if obj + .keys() + .any(|key| !matches!(key.as_str(), "role" | "content" | "name")) + { + return Ok(native()); + } + let content = decode_chat_content(obj.get("content").ok_or_else(|| { + FlowError::InvalidArgument("OpenAI Chat message is missing content".into()) + })?)?; + let name = optional_chat_string(obj, "name", "message")?; + Ok(match role { + "system" => Message::System { content, name }, + "developer" => Message::Developer { content, name }, + _ => Message::User { content, name }, + }) + } + "assistant" => { + if obj + .keys() + .any(|key| !matches!(key.as_str(), "role" | "content" | "tool_calls" | "name")) + { + return Ok(native()); + } + let content = obj + .get("content") + .filter(|content| !content.is_null()) + .map(decode_chat_content) + .transpose()?; + let tool_calls = match obj.get("tool_calls") { + Some(Json::Null) | None => None, + Some(Json::Array(calls)) => { + let decoded = calls + .iter() + .map(decode_chat_tool_call) + .collect::>>()?; + let Some(decoded) = decoded.into_iter().collect::>>() else { + return Ok(native()); + }; + Some(decoded) + } + Some(_) => { + return Err(FlowError::InvalidArgument( + "OpenAI Chat assistant tool_calls must be an array or null".into(), + )); + } + }; + Ok(Message::Assistant { + content, + tool_calls, + name: optional_chat_string(obj, "name", "assistant message")?, + }) + } + "tool" + if obj + .keys() + .all(|key| matches!(key.as_str(), "role" | "content" | "tool_call_id")) => + { + Ok(Message::Tool { + content: decode_chat_content(obj.get("content").ok_or_else(|| { + FlowError::InvalidArgument("OpenAI Chat tool message is missing content".into()) + })?)?, + tool_call_id: obj + .get("tool_call_id") + .and_then(Json::as_str) + .ok_or_else(|| { + FlowError::InvalidArgument( + "OpenAI Chat tool message is missing tool_call_id".into(), + ) + })? + .to_string(), + }) + } + "function" + if obj + .keys() + .all(|key| matches!(key.as_str(), "role" | "content" | "name")) => + { + Ok(Message::Function { + content: optional_chat_string(obj, "content", "function message")?, + name: obj + .get("name") + .and_then(Json::as_str) + .ok_or_else(|| { + FlowError::InvalidArgument( + "OpenAI Chat function message is missing name".into(), + ) + })? + .to_string(), + }) + } + _ => Ok(native()), + } +} + +fn encode_chat_content(content: &MessageContent) -> Result { + match content { + MessageContent::Text(text) => Ok(Json::String(text.clone())), + MessageContent::Parts(parts) => Ok(Json::Array( + parts + .iter() + .map(|part| match part { + ContentPart::Text { text, extra } => { + let mut obj = extra.clone(); + obj.insert("type".into(), Json::String("text".into())); + obj.insert("text".into(), Json::String(text.clone())); + Ok(Json::Object(obj)) + } + ContentPart::ImageUrl { image_url, extra } => { + let mut obj = extra.clone(); + obj.insert("type".into(), Json::String("image_url".into())); + obj.insert( + "image_url".into(), + serde_json::to_value(image_url).map_err(|error| { + FlowError::Internal(format!( + "OpenAI Chat image URL encode: {error}" + )) + })?, + ); + Ok(Json::Object(obj)) + } + ContentPart::Audio { audio, extra } => { + let mut obj = extra.clone(); + obj.insert("type".into(), Json::String("input_audio".into())); + obj.insert("input_audio".into(), audio.clone()); + Ok(Json::Object(obj)) + } + ContentPart::File { file, extra } => { + let mut obj = extra.clone(); + obj.insert("type".into(), Json::String("file".into())); + obj.insert("file".into(), file.clone()); + Ok(Json::Object(obj)) + } + ContentPart::Refusal { refusal, extra } => { + let mut obj = extra.clone(); + obj.insert("type".into(), Json::String("refusal".into())); + obj.insert("refusal".into(), Json::String(refusal.clone())); + Ok(Json::Object(obj)) + } + ContentPart::ProviderNative { + provider, value, .. + } if provider == "openai_chat" => Ok(value.clone()), + other => Err(FlowError::InvalidArgument(format!( + "content part {other:?} cannot be encoded for OpenAI Chat" + ))), + }) + .collect::>>()?, + )), + } +} + +fn encode_chat_message(message: &Message) -> Result { + let message_with_content = + |role: &str, content: &MessageContent, name: &Option| -> Result { + let mut obj = serde_json::Map::new(); + obj.insert("role".into(), Json::String(role.into())); + obj.insert("content".into(), encode_chat_content(content)?); + if let Some(name) = name { + obj.insert("name".into(), Json::String(name.clone())); + } + Ok(Json::Object(obj)) + }; + match message { + Message::System { content, name } => message_with_content("system", content, name), + Message::Developer { content, name } => message_with_content("developer", content, name), + Message::User { content, name } => message_with_content("user", content, name), + Message::Assistant { + content, + tool_calls, + name, + } => { + let mut obj = serde_json::Map::new(); + obj.insert("role".into(), Json::String("assistant".into())); + if let Some(content) = content { + obj.insert("content".into(), encode_chat_content(content)?); + } + if let Some(tool_calls) = tool_calls { + obj.insert( + "tool_calls".into(), + serde_json::to_value(tool_calls).map_err(|error| { + FlowError::Internal(format!("OpenAI Chat tool calls encode: {error}")) + })?, + ); + } + if let Some(name) = name { + obj.insert("name".into(), Json::String(name.clone())); + } + Ok(Json::Object(obj)) + } + Message::Tool { + content, + tool_call_id, + } => { + let mut obj = serde_json::Map::new(); + obj.insert("role".into(), Json::String("tool".into())); + obj.insert("content".into(), encode_chat_content(content)?); + obj.insert("tool_call_id".into(), Json::String(tool_call_id.clone())); + Ok(Json::Object(obj)) + } + Message::Function { content, name } => { + let mut obj = serde_json::Map::new(); + obj.insert("role".into(), Json::String("function".into())); + obj.insert( + "content".into(), + content.clone().map(Json::String).unwrap_or(Json::Null), + ); + obj.insert("name".into(), Json::String(name.clone())); + Ok(Json::Object(obj)) + } + Message::ProviderNative { + provider, value, .. + } if provider == "openai_chat" => Ok(value.clone()), + other => Err(FlowError::InvalidArgument(format!( + "message {other:?} cannot be encoded for OpenAI Chat" + ))), + } +} + +fn decode_chat_tool(value: &Json) -> Result { + let obj = value + .as_object() + .ok_or_else(|| FlowError::InvalidArgument("OpenAI Chat tool must be an object".into()))?; + if obj.get("type").and_then(Json::as_str) != Some("function") { + let native = chat_native( + obj.get("type").and_then(Json::as_str).unwrap_or("unknown"), + value, + ); + return Ok(ToolDefinition::ProviderNative { + provider: native.provider, + kind: native.kind, + value: native.value, + }); + } + let function = obj + .get("function") + .and_then(Json::as_object) + .ok_or_else(|| { + FlowError::InvalidArgument("OpenAI Chat function tool is missing function".into()) + })?; + let name = function.get("name").and_then(Json::as_str).ok_or_else(|| { + FlowError::InvalidArgument("OpenAI Chat function tool is missing name".into()) + })?; + let description = super::optional_string(function, "description", "OpenAI Chat function tool")?; + let strict = super::optional_bool(function, "strict", "OpenAI Chat function tool")?; + Ok(ToolDefinition::Function { + function: FunctionDefinition { + name: name.to_string(), + description, + parameters: function.get("parameters").cloned(), + strict, + extra: function + .iter() + .filter(|(key, _)| { + !matches!( + key.as_str(), + "name" | "description" | "parameters" | "strict" + ) + }) + .map(|(key, value)| (key.clone(), value.clone())) + .collect(), + }, + extra: obj + .iter() + .filter(|(key, _)| !matches!(key.as_str(), "type" | "function")) + .map(|(key, value)| (key.clone(), value.clone())) + .collect(), + }) +} + +fn encode_chat_tool(tool: &ToolDefinition) -> Result { + match tool { + ToolDefinition::Function { function, extra } => { + let mut function_obj = function.extra.clone(); + function_obj.insert("name".into(), Json::String(function.name.clone())); + if let Some(description) = &function.description { + function_obj.insert("description".into(), Json::String(description.clone())); + } + if let Some(parameters) = &function.parameters { + function_obj.insert("parameters".into(), parameters.clone()); + } + if let Some(strict) = function.strict { + function_obj.insert("strict".into(), Json::Bool(strict)); + } + let mut obj = extra.clone(); + obj.insert("type".into(), Json::String("function".into())); + obj.insert("function".into(), Json::Object(function_obj)); + Ok(Json::Object(obj)) + } + ToolDefinition::ProviderNative { + provider, value, .. + } if provider == "openai_chat" => Ok(value.clone()), + other => Err(FlowError::InvalidArgument(format!( + "tool {other:?} cannot be encoded for OpenAI Chat" + ))), + } +} + +fn decode_chat_tool_choice(value: &Json) -> ToolChoice { + match value.as_str() { + Some("auto") => ToolChoice::Auto, + Some("none") => ToolChoice::None, + Some("required") => ToolChoice::Required, + _ => { + if let Some(function) = value + .as_object() + .filter(|obj| obj.get("type").and_then(Json::as_str) == Some("function")) + .and_then(|obj| obj.get("function")) + .and_then(Json::as_object) + .and_then(|function| function.get("name")) + .and_then(Json::as_str) + { + ToolChoice::Specific(super::request::ToolChoiceFunction { + choice_type: "function".into(), + function: super::request::ToolChoiceFunctionName { + name: function.to_string(), + }, + }) + } else { + ToolChoice::ProviderNative(chat_native("tool_choice", value)) + } + } + } +} + +fn encode_chat_tool_choice(choice: &ToolChoice) -> Result { + match choice { + ToolChoice::Auto => Ok(Json::String("auto".into())), + ToolChoice::None => Ok(Json::String("none".into())), + ToolChoice::Required => Ok(Json::String("required".into())), + ToolChoice::Specific(choice) => Ok(serde_json::json!({ + "type":"function", + "function":{"name":choice.function.name} + })), + ToolChoice::ProviderNative(native) if native.provider == "openai_chat" => { + Ok(native.value.clone()) + } + ToolChoice::ProviderNative(native) => Err(FlowError::InvalidArgument(format!( + "tool choice for {} cannot be encoded for OpenAI Chat", + native.provider + ))), + } +} + +fn patch_extra_fields( + obj: &mut serde_json::Map, + baseline: &serde_json::Map, + edited: &serde_json::Map, +) { + for key in baseline.keys().filter(|key| !edited.contains_key(*key)) { + obj.remove(key); + } + for (key, value) in edited { + if baseline.get(key) != Some(value) { + obj.insert(key.clone(), value.clone()); + } + } +} + +fn set_or_remove_json(obj: &mut serde_json::Map, key: &str, value: Option) { + if let Some(value) = value { + obj.insert(key.into(), value); + } else { + obj.remove(key); + } +} + // --------------------------------------------------------------------------- // LlmResponseCodec implementation // --------------------------------------------------------------------------- @@ -237,29 +805,37 @@ impl LlmCodec for OpenAIChatCodec { .content .as_object() .ok_or_else(|| FlowError::Internal("request content is not an object".into()))?; - - // Extract messages (default to empty vec if absent). - let messages: Vec = obj + let messages = obj .get("messages") - .map(|v| serde_json::from_value(v.clone()).unwrap_or_default()) - .unwrap_or_default(); - - // Extract model. - let model = obj.get("model").and_then(|v| v.as_str()).map(String::from); - - // Extract generation params. - let temperature = obj.get("temperature").and_then(|v| v.as_f64()); - let top_p = obj.get("top_p").and_then(|v| v.as_f64()); - let stop = obj - .get("stop") - .and_then(|v| serde_json::from_value::>(v.clone()).ok()); - - // max_completion_tokens takes priority over max_tokens (newer API key). - let max_tokens = obj - .get("max_completion_tokens") - .and_then(|v| v.as_u64()) - .or_else(|| obj.get("max_tokens").and_then(|v| v.as_u64())); - + .ok_or_else(|| { + FlowError::InvalidArgument("OpenAI Chat request is missing messages".into()) + })? + .as_array() + .ok_or_else(|| { + FlowError::InvalidArgument("OpenAI Chat messages must be an array".into()) + })? + .iter() + .map(decode_chat_message) + .collect::>>()?; + let model = super::optional_string(obj, "model", "OpenAI Chat")?; + let temperature = super::optional_f64(obj, "temperature", "OpenAI Chat")?; + let top_p = super::optional_f64(obj, "top_p", "OpenAI Chat")?; + let stop = match obj.get("stop") { + Some(Json::String(stop)) => Some(vec![stop.clone()]), + Some(Json::Array(_)) => Some( + serde_json::from_value::>(obj["stop"].clone()).map_err(|error| { + FlowError::InvalidArgument(format!("invalid OpenAI Chat stop value: {error}")) + })?, + ), + Some(Json::Null) | None => None, + Some(_) => { + return Err(FlowError::InvalidArgument( + "OpenAI Chat stop must be a string, array, or null".into(), + )); + } + }; + let max_tokens = super::optional_u64(obj, "max_completion_tokens", "OpenAI Chat")? + .or(super::optional_u64(obj, "max_tokens", "OpenAI Chat")?); let params = if temperature.is_some() || max_tokens.is_some() || top_p.is_some() || stop.is_some() { Some(GenerationParams { @@ -271,22 +847,62 @@ impl LlmCodec for OpenAIChatCodec { } else { None }; - - // Extract tools. - let tools: Option> = obj + let tools = obj .get("tools") - .map(|v| serde_json::from_value(v.clone())) - .transpose() - .map_err(|e| FlowError::Internal(format!("OpenAI Chat tools decode: {e}")))?; - - // Extract tool_choice. - let tool_choice: Option = obj - .get("tool_choice") - .map(|v| serde_json::from_value(v.clone())) - .transpose() - .map_err(|e| FlowError::Internal(format!("OpenAI Chat tool_choice decode: {e}")))?; - - // Collect extra fields (keys not in MODELED_REQUEST_KEYS). + .map(|value| { + value + .as_array() + .ok_or_else(|| { + FlowError::InvalidArgument("OpenAI Chat tools must be an array".into()) + })? + .iter() + .map(decode_chat_tool) + .collect::>>() + }) + .transpose()?; + let tool_choice = obj.get("tool_choice").map(decode_chat_tool_choice); + let store = super::optional_bool(obj, "store", "OpenAI Chat")?; + let user = super::optional_string(obj, "user", "OpenAI Chat")?; + let service_tier = super::optional_string(obj, "service_tier", "OpenAI Chat")?; + let parallel_tool_calls = super::optional_bool(obj, "parallel_tool_calls", "OpenAI Chat")?; + let top_logprobs = super::optional_u64(obj, "top_logprobs", "OpenAI Chat")?; + let stream = super::optional_bool(obj, "stream", "OpenAI Chat")?; + let frequency_penalty = super::optional_f64(obj, "frequency_penalty", "OpenAI Chat")?; + let functions = match obj.get("functions") { + Some(Json::Null) | None => None, + Some(Json::Array(functions)) => Some(functions.clone()), + Some(_) => { + return Err(FlowError::InvalidArgument( + "OpenAI Chat functions must be an array or null".into(), + )); + } + }; + let logprobs = super::optional_bool(obj, "logprobs", "OpenAI Chat")?; + let modalities = match obj.get("modalities") { + Some(Json::Null) | None => None, + Some(value) => Some(serde_json::from_value(value.clone()).map_err(|error| { + FlowError::InvalidArgument(format!("invalid OpenAI Chat modalities: {error}")) + })?), + }; + let n = super::optional_u64(obj, "n", "OpenAI Chat")?; + let presence_penalty = super::optional_f64(obj, "presence_penalty", "OpenAI Chat")?; + let prompt_cache_key = super::optional_string(obj, "prompt_cache_key", "OpenAI Chat")?; + let prompt_cache_retention = + super::optional_string(obj, "prompt_cache_retention", "OpenAI Chat")?; + let reasoning_effort = super::optional_string(obj, "reasoning_effort", "OpenAI Chat")?; + let safety_identifier = super::optional_string(obj, "safety_identifier", "OpenAI Chat")?; + let seed = super::optional_i64(obj, "seed", "OpenAI Chat")?; + let verbosity = super::optional_string(obj, "verbosity", "OpenAI Chat")?; + let metadata = super::optional_object(obj, "metadata", "OpenAI Chat")?; + let audio = super::optional_object(obj, "audio", "OpenAI Chat")?; + let logit_bias = super::optional_object(obj, "logit_bias", "OpenAI Chat")?; + let moderation = super::optional_object(obj, "moderation", "OpenAI Chat")?; + let prediction = super::optional_object(obj, "prediction", "OpenAI Chat")?; + let prompt_cache_options = + super::optional_object(obj, "prompt_cache_options", "OpenAI Chat")?; + let response_format = super::optional_object(obj, "response_format", "OpenAI Chat")?; + let stream_options = super::optional_object(obj, "stream_options", "OpenAI Chat")?; + let web_search_options = super::optional_object(obj, "web_search_options", "OpenAI Chat")?; let extra: serde_json::Map = obj .iter() .filter(|(k, _)| !MODELED_REQUEST_KEYS.contains(&k.as_str())) @@ -295,104 +911,351 @@ impl LlmCodec for OpenAIChatCodec { Ok(AnnotatedLlmRequest { messages, + instructions: None, model, params, tools, tool_choice, - store: obj.get("store").and_then(|v| v.as_bool()), + store, previous_response_id: None, truncation: None, reasoning: None, include: None, - user: obj.get("user").and_then(|v| v.as_str()).map(String::from), - metadata: obj.get("metadata").cloned(), - service_tier: obj - .get("service_tier") - .and_then(|v| v.as_str()) - .map(String::from), - parallel_tool_calls: obj.get("parallel_tool_calls").and_then(|v| v.as_bool()), + user, + metadata, + service_tier, + parallel_tool_calls, max_output_tokens: None, max_tool_calls: None, - top_logprobs: obj.get("top_logprobs").and_then(|v| v.as_u64()), - stream: obj.get("stream").and_then(|v| v.as_bool()), + top_logprobs, + stream, + api_specific: Some(ApiSpecificRequest::OpenAIChat { + audio, + frequency_penalty, + function_call: obj.get("function_call").cloned(), + functions, + logit_bias, + logprobs, + modalities, + moderation, + n, + prediction, + presence_penalty, + prompt_cache_key, + prompt_cache_options, + prompt_cache_retention, + reasoning_effort, + response_format, + safety_identifier, + seed, + stream_options, + verbosity, + web_search_options, + }), extra, }) } fn encode(&self, annotated: &AnnotatedLlmRequest, original: &LlmRequest) -> Result { + let baseline = self.decode(original)?; let mut content = original.content.clone(); let obj = content .as_object_mut() .ok_or_else(|| FlowError::Internal("original content is not an object".into()))?; - - insert_serialized(obj, "messages", &annotated.messages, "messages")?; - - if let Some(ref model) = annotated.model { - obj.insert("model".into(), Json::String(model.clone())); + if annotated.messages != baseline.messages { + let original_messages = obj.get("messages").and_then(Json::as_array); + obj.insert( + "messages".into(), + Json::Array(super::encode_changed_items( + &annotated.messages, + &baseline.messages, + original_messages.map(Vec::as_slice), + encode_chat_message, + )?), + ); } - - if let Some(ref params) = annotated.params { - overlay_generation_params(obj, params)?; + if annotated.instructions != baseline.instructions + || annotated.previous_response_id != baseline.previous_response_id + || annotated.truncation != baseline.truncation + || annotated.reasoning != baseline.reasoning + || annotated.include != baseline.include + || annotated.max_output_tokens != baseline.max_output_tokens + || annotated.max_tool_calls != baseline.max_tool_calls + { + return Err(FlowError::InvalidArgument( + "request contains fields that cannot be encoded for OpenAI Chat".into(), + )); } - - if let Some(ref tools) = annotated.tools { - insert_serialized(obj, "tools", tools, "tools")?; + if annotated.model != baseline.model { + set_or_remove_json(obj, "model", annotated.model.clone().map(Json::String)); } - - if let Some(ref tool_choice) = annotated.tool_choice { - insert_serialized(obj, "tool_choice", tool_choice, "tool_choice")?; + if annotated.params != baseline.params { + let edited = annotated.params.as_ref(); + let before = baseline.params.as_ref(); + if edited.and_then(|p| p.temperature) != before.and_then(|p| p.temperature) { + set_or_remove_json( + obj, + "temperature", + edited.and_then(|p| p.temperature).map(json_f64), + ); + } + if edited.and_then(|p| p.top_p) != before.and_then(|p| p.top_p) { + set_or_remove_json(obj, "top_p", edited.and_then(|p| p.top_p).map(json_f64)); + } + if edited.and_then(|p| p.max_tokens) != before.and_then(|p| p.max_tokens) { + let max_tokens = edited.and_then(|p| p.max_tokens).map(Json::from); + if max_tokens.is_none() { + obj.remove("max_completion_tokens"); + obj.remove("max_tokens"); + } else { + let key = if obj.contains_key("max_completion_tokens") + || !obj.contains_key("max_tokens") + { + "max_completion_tokens" + } else { + "max_tokens" + }; + set_or_remove_json(obj, key, max_tokens); + } + } + if edited.and_then(|p| p.stop.as_ref()) != before.and_then(|p| p.stop.as_ref()) { + let stop = edited.and_then(|p| p.stop.as_ref()).map(|values| { + if obj.get("stop").is_some_and(Json::is_string) && values.len() == 1 { + Json::String(values[0].clone()) + } else { + serde_json::json!(values) + } + }); + set_or_remove_json(obj, "stop", stop); + } } - - if let Some(store) = annotated.store { - obj.insert("store".into(), Json::Bool(store)); + if annotated.tools != baseline.tools { + let tools = annotated + .tools + .as_deref() + .map(|tools| { + super::encode_changed_items( + tools, + baseline.tools.as_deref().unwrap_or(&[]), + obj.get("tools").and_then(Json::as_array).map(Vec::as_slice), + encode_chat_tool, + ) + }) + .transpose()? + .map(Json::Array); + set_or_remove_json(obj, "tools", tools); } - if let Some(ref user) = annotated.user { - obj.insert("user".into(), Json::String(user.clone())); + if annotated.tool_choice != baseline.tool_choice { + let tool_choice = match (&annotated.tool_choice, &baseline.tool_choice) { + (Some(edited), Some(before)) => { + let edited = encode_chat_tool_choice(edited)?; + let before = encode_chat_tool_choice(before)?; + Some(match obj.get("tool_choice") { + Some(original) => super::patch_changed_json(original, &before, &edited)?, + None => edited, + }) + } + (Some(edited), None) => Some(encode_chat_tool_choice(edited)?), + (None, _) => None, + }; + set_or_remove_json(obj, "tool_choice", tool_choice); } - if let Some(ref metadata) = annotated.metadata { - obj.insert("metadata".into(), metadata.clone()); + for (key, edited, before) in [("metadata", &annotated.metadata, &baseline.metadata)] { + if edited != before { + set_or_remove_json(obj, key, edited.clone()); + } } - if let Some(ref service_tier) = annotated.service_tier { - obj.insert("service_tier".into(), Json::String(service_tier.clone())); + for (key, edited, before) in [ + ("store", annotated.store, baseline.store), + ( + "parallel_tool_calls", + annotated.parallel_tool_calls, + baseline.parallel_tool_calls, + ), + ("stream", annotated.stream, baseline.stream), + ] { + if edited != before { + set_or_remove_json(obj, key, edited.map(Json::Bool)); + } } - if let Some(parallel_tool_calls) = annotated.parallel_tool_calls { - obj.insert( - "parallel_tool_calls".into(), - Json::Bool(parallel_tool_calls), - ); + for (key, edited, before) in [ + ("user", &annotated.user, &baseline.user), + ( + "service_tier", + &annotated.service_tier, + &baseline.service_tier, + ), + ] { + if edited != before { + set_or_remove_json(obj, key, edited.clone().map(Json::String)); + } } - if let Some(top_logprobs) = annotated.top_logprobs { - obj.insert("top_logprobs".into(), Json::from(top_logprobs)); - } - if let Some(stream) = annotated.stream { - obj.insert("stream".into(), Json::Bool(stream)); - } - - for (k, v) in &annotated.extra { - obj.insert(k.clone(), v.clone()); - } - - // Force `stream_options.include_usage` when the caller did not set it. - // - // Rationale: OpenAI-compatible backends only emit the terminal chunk - // containing `usage` (prompt/completion/total tokens) when this flag - // is true. Without it, Phoenix spans show `token_count=0` for every - // LLM call even though the provider knows the real counts. The - // observability exporter (OpenInference) reads usage off the - // annotated response, so the flag has to be set at the request level - // before bytes go on the wire. - // - // Guarded on `stream == true` per the OpenAI Chat Completions spec, - // which restricts `stream_options` to streaming requests. Caller- - // provided `stream_options` are preserved verbatim (including - // explicit opt-outs such as `include_usage: false`). - let is_streaming = obj.get("stream").and_then(|v| v.as_bool()).unwrap_or(false); - if is_streaming && !obj.contains_key("stream_options") { - obj.insert( - "stream_options".into(), - serde_json::json!({"include_usage": true}), - ); + if annotated.top_logprobs != baseline.top_logprobs { + set_or_remove_json(obj, "top_logprobs", annotated.top_logprobs.map(Json::from)); } + match (&annotated.api_specific, &baseline.api_specific) { + ( + Some(ApiSpecificRequest::OpenAIChat { + audio, + frequency_penalty, + function_call, + functions, + logit_bias, + logprobs, + modalities, + moderation, + n, + prediction, + presence_penalty, + prompt_cache_key, + prompt_cache_options, + prompt_cache_retention, + reasoning_effort, + response_format, + safety_identifier, + seed, + stream_options, + verbosity, + web_search_options, + }), + Some(ApiSpecificRequest::OpenAIChat { + audio: old_audio, + frequency_penalty: old_frequency_penalty, + function_call: old_function_call, + functions: old_functions, + logit_bias: old_logit_bias, + logprobs: old_logprobs, + modalities: old_modalities, + moderation: old_moderation, + n: old_n, + prediction: old_prediction, + presence_penalty: old_presence_penalty, + prompt_cache_key: old_prompt_cache_key, + prompt_cache_options: old_prompt_cache_options, + prompt_cache_retention: old_prompt_cache_retention, + reasoning_effort: old_reasoning_effort, + response_format: old_response_format, + safety_identifier: old_safety_identifier, + seed: old_seed, + stream_options: old_stream_options, + verbosity: old_verbosity, + web_search_options: old_web_search_options, + }), + ) => { + for (key, edited, before) in [ + ("audio", audio, old_audio), + ("function_call", function_call, old_function_call), + ("logit_bias", logit_bias, old_logit_bias), + ("moderation", moderation, old_moderation), + ("prediction", prediction, old_prediction), + ( + "prompt_cache_options", + prompt_cache_options, + old_prompt_cache_options, + ), + ("response_format", response_format, old_response_format), + ("stream_options", stream_options, old_stream_options), + ( + "web_search_options", + web_search_options, + old_web_search_options, + ), + ] { + if edited != before { + set_or_remove_json(obj, key, edited.clone()); + } + } + for (key, edited, before) in [ + ( + "frequency_penalty", + frequency_penalty, + old_frequency_penalty, + ), + ("presence_penalty", presence_penalty, old_presence_penalty), + ] { + if edited != before { + set_or_remove_json(obj, key, edited.map(json_f64)); + } + } + if functions != old_functions { + set_or_remove_json(obj, "functions", functions.clone().map(Json::Array)); + } + if modalities != old_modalities { + set_or_remove_json( + obj, + "modalities", + modalities.as_ref().map(|v| serde_json::json!(v)), + ); + } + if logprobs != old_logprobs { + set_or_remove_json(obj, "logprobs", logprobs.map(Json::Bool)); + } + if n != old_n { + set_or_remove_json(obj, "n", n.map(Json::from)); + } + for (key, edited, before) in [ + ("prompt_cache_key", prompt_cache_key, old_prompt_cache_key), + ( + "prompt_cache_retention", + prompt_cache_retention, + old_prompt_cache_retention, + ), + ("reasoning_effort", reasoning_effort, old_reasoning_effort), + ( + "safety_identifier", + safety_identifier, + old_safety_identifier, + ), + ("verbosity", verbosity, old_verbosity), + ] { + if edited != before { + set_or_remove_json(obj, key, edited.clone().map(Json::String)); + } + } + if seed != old_seed { + set_or_remove_json(obj, "seed", seed.map(Json::from)); + } + } + (None, Some(ApiSpecificRequest::OpenAIChat { .. })) => { + for key in [ + "audio", + "frequency_penalty", + "function_call", + "functions", + "logit_bias", + "logprobs", + "modalities", + "moderation", + "n", + "prediction", + "presence_penalty", + "prompt_cache_key", + "prompt_cache_options", + "prompt_cache_retention", + "reasoning_effort", + "response_format", + "safety_identifier", + "seed", + "stream_options", + "verbosity", + "web_search_options", + ] { + obj.remove(key); + } + } + (Some(_), _) => { + return Err(FlowError::InvalidArgument( + "api_specific provider does not match OpenAI Chat".into(), + )); + } + (None, Some(_)) => { + return Err(FlowError::InvalidArgument( + "api_specific provider does not match OpenAI Chat".into(), + )); + } + (None, None) => {} + } + patch_extra_fields(obj, &baseline.extra, &annotated.extra); Ok(LlmRequest { headers: original.headers.clone(), @@ -408,42 +1271,6 @@ fn json_f64(v: f64) -> Json { .unwrap_or(Json::Null) } -fn insert_serialized( - obj: &mut serde_json::Map, - key: &str, - value: &T, - context: &str, -) -> Result<()> { - let json = serde_json::to_value(value) - .map_err(|e| FlowError::Internal(format!("OpenAI Chat {context} encode: {e}")))?; - obj.insert(key.into(), json); - Ok(()) -} - -fn overlay_generation_params( - obj: &mut serde_json::Map, - params: &GenerationParams, -) -> Result<()> { - if let Some(temp) = params.temperature { - obj.insert("temperature".into(), json_f64(temp)); - } - if let Some(top_p) = params.top_p { - obj.insert("top_p".into(), json_f64(top_p)); - } - if let Some(ref stop) = params.stop { - insert_serialized(obj, "stop", stop, "stop")?; - } - if let Some(max_tokens) = params.max_tokens { - let key = if obj.contains_key("max_completion_tokens") { - "max_completion_tokens" - } else { - "max_tokens" - }; - obj.insert(key.into(), Json::from(max_tokens)); - } - Ok(()) -} - // --------------------------------------------------------------------------- // Streaming codec // --------------------------------------------------------------------------- diff --git a/crates/core/src/codec/openai_responses.rs b/crates/core/src/codec/openai_responses.rs index 0e0e5e346..e14266bbe 100644 --- a/crates/core/src/codec/openai_responses.rs +++ b/crates/core/src/codec/openai_responses.rs @@ -22,7 +22,8 @@ use crate::error::{FlowError, Result}; use crate::json::Json; use super::request::{ - AnnotatedLlmRequest, GenerationParams, Message, MessageContent, ToolChoice, ToolChoiceFunction, + AnnotatedLlmRequest, ApiSpecificRequest, ContentPart, FunctionDefinition, GenerationParams, + Message, MessageContent, ProviderNativeComponent, ToolChoice, ToolChoiceFunction, ToolChoiceFunctionName, ToolDefinition, }; use super::resolve::{ProviderSurface, ProviderSurfaceDescriptor}; @@ -174,8 +175,18 @@ const MODELED_REQUEST_KEYS: &[&str] = &[ "max_tool_calls", "top_logprobs", "stream", + "background", + "context_management", + "conversation", + "moderation", + "prompt", + "prompt_cache_key", + "prompt_cache_options", + "prompt_cache_retention", + "safety_identifier", + "stream_options", + "text", ]; -const UNPARSED_INPUT_ITEMS_KEY: &str = "_openai_responses_unparsed_input_items"; /// Helper to construct a [`Json`] number from an `f64`. fn json_f64(v: f64) -> Json { @@ -277,201 +288,705 @@ fn optional_vec(items: Vec) -> Option> { (!items.is_empty()).then_some(items) } -fn split_system_and_input_messages(messages: &[Message]) -> (Option, Vec<&Message>) { - let mut system_text = None; - let mut input_messages = Vec::new(); - - for msg in messages { - match msg { - Message::System { content, .. } => { - if let MessageContent::Text(text) = content { - system_text = Some(text.clone()); - } - } - other => input_messages.push(other), - } +fn responses_native(kind: &str, value: &Json) -> ProviderNativeComponent { + ProviderNativeComponent { + provider: "openai_responses".into(), + kind: kind.to_string(), + value: value.clone(), } - - (system_text, input_messages) } -fn set_or_remove_string(obj: &mut serde_json::Map, key: &str, value: Option) { - if let Some(value) = value { - obj.insert(key.into(), Json::String(value)); - } else { - obj.remove(key); +fn decode_responses_content(value: &Json) -> Result { + if let Some(text) = value.as_str() { + return Ok(MessageContent::Text(text.to_string())); } + let parts = value.as_array().ok_or_else(|| { + FlowError::InvalidArgument( + "OpenAI Responses message content must be a string or array".into(), + ) + })?; + Ok(MessageContent::Parts( + parts + .iter() + .map(decode_responses_content_part) + .collect::>>()?, + )) } -fn insert_serialized( - obj: &mut serde_json::Map, - key: &str, - value: &T, - context: &str, -) -> Result<()> { - let json = serde_json::to_value(value) - .map_err(|e| FlowError::Internal(format!("OpenAI Responses {context} encode: {e}")))?; - obj.insert(key.into(), json); - Ok(()) +fn decode_responses_content_part(value: &Json) -> Result { + let obj = value.as_object().ok_or_else(|| { + FlowError::InvalidArgument("OpenAI Responses content part must be an object".into()) + })?; + let kind = obj.get("type").and_then(Json::as_str).unwrap_or("unknown"); + match kind { + "input_text" | "output_text" => Ok(ContentPart::Text { + text: obj + .get("text") + .and_then(Json::as_str) + .ok_or_else(|| { + FlowError::InvalidArgument("OpenAI Responses text part is missing text".into()) + })? + .to_string(), + extra: obj + .iter() + .filter(|(key, _)| !matches!(key.as_str(), "type" | "text")) + .map(|(key, value)| (key.clone(), value.clone())) + .collect(), + }), + "input_image" => Ok(ContentPart::Image { + image: Json::Object( + obj.iter() + .filter(|(key, _)| matches!(key.as_str(), "image_url" | "file_id" | "detail")) + .map(|(key, value)| (key.clone(), value.clone())) + .collect(), + ), + extra: obj + .iter() + .filter(|(key, _)| { + !matches!(key.as_str(), "type" | "image_url" | "file_id" | "detail") + }) + .map(|(key, value)| (key.clone(), value.clone())) + .collect(), + }), + "input_file" => Ok(ContentPart::File { + file: Json::Object( + obj.iter() + .filter(|(key, _)| { + matches!( + key.as_str(), + "file_data" | "file_id" | "file_url" | "filename" + ) + }) + .map(|(key, value)| (key.clone(), value.clone())) + .collect(), + ), + extra: obj + .iter() + .filter(|(key, _)| { + !matches!( + key.as_str(), + "type" | "file_data" | "file_id" | "file_url" | "filename" + ) + }) + .map(|(key, value)| (key.clone(), value.clone())) + .collect(), + }), + "refusal" => Ok(ContentPart::Refusal { + refusal: obj + .get("refusal") + .and_then(Json::as_str) + .ok_or_else(|| { + FlowError::InvalidArgument( + "OpenAI Responses refusal part is missing refusal".into(), + ) + })? + .to_string(), + extra: obj + .iter() + .filter(|(key, _)| !matches!(key.as_str(), "type" | "refusal")) + .map(|(key, value)| (key.clone(), value.clone())) + .collect(), + }), + _ => Ok(ContentPart::ProviderNative { + provider: "openai_responses".into(), + kind: kind.to_string(), + value: value.clone(), + }), + } } -fn overlay_generation_params(obj: &mut serde_json::Map, params: &GenerationParams) { - if let Some(temp) = params.temperature { - obj.insert("temperature".into(), json_f64(temp)); - } - if let Some(top_p) = params.top_p { - obj.insert("top_p".into(), json_f64(top_p)); +fn decode_responses_input_item(value: &Json) -> Result { + let obj = value.as_object().ok_or_else(|| { + FlowError::InvalidArgument("OpenAI Responses input item must be an object".into()) + })?; + if let Some(role) = obj.get("role").and_then(Json::as_str) { + if obj + .keys() + .any(|key| !matches!(key.as_str(), "type" | "role" | "content")) + { + return Ok(Message::ProviderNative { + provider: "openai_responses".into(), + kind: "message".into(), + value: value.clone(), + }); + } + let content = decode_responses_content(obj.get("content").ok_or_else(|| { + FlowError::InvalidArgument("OpenAI Responses message is missing content".into()) + })?)?; + return Ok(match role { + "user" => Message::User { + content, + name: None, + }, + "system" => Message::System { + content, + name: None, + }, + "developer" => Message::Developer { + content, + name: None, + }, + "assistant" => Message::Assistant { + content: Some(content), + tool_calls: None, + name: None, + }, + _ => Message::ProviderNative { + provider: "openai_responses".into(), + kind: "message".into(), + value: value.clone(), + }, + }); } - if let Some(max_tokens) = params.max_tokens { - obj.insert("max_output_tokens".into(), Json::from(max_tokens)); - obj.remove("max_tokens"); + + let kind = obj.get("type").and_then(Json::as_str).unwrap_or("unknown"); + match kind { + "function_call" => { + let call_id = obj.get("call_id").and_then(Json::as_str).ok_or_else(|| { + FlowError::InvalidArgument( + "OpenAI Responses function_call is missing call_id".into(), + ) + })?; + let name = obj.get("name").and_then(Json::as_str).ok_or_else(|| { + FlowError::InvalidArgument("OpenAI Responses function_call is missing name".into()) + })?; + let arguments = obj.get("arguments").and_then(Json::as_str).ok_or_else(|| { + FlowError::InvalidArgument( + "OpenAI Responses function_call is missing arguments".into(), + ) + })?; + let id = match obj.get("id") { + Some(Json::String(id)) => Some(id.clone()), + Some(Json::Null) | None => None, + Some(_) => { + return Err(FlowError::InvalidArgument( + "OpenAI Responses function_call id must be a string or null".into(), + )); + } + }; + Ok(Message::ToolCallItem { + id, + call_id: call_id.to_string(), + name: name.to_string(), + arguments: parse_arguments(arguments), + extra: obj + .iter() + .filter(|(key, _)| { + !matches!( + key.as_str(), + "type" | "id" | "call_id" | "name" | "arguments" + ) + }) + .map(|(key, value)| (key.clone(), value.clone())) + .collect(), + }) + } + "function_call_output" => { + let call_id = obj.get("call_id").and_then(Json::as_str).ok_or_else(|| { + FlowError::InvalidArgument( + "OpenAI Responses function_call_output is missing call_id".into(), + ) + })?; + let output = obj.get("output").ok_or_else(|| { + FlowError::InvalidArgument( + "OpenAI Responses function_call_output is missing output".into(), + ) + })?; + let id = match obj.get("id") { + Some(Json::String(id)) => Some(id.clone()), + Some(Json::Null) | None => None, + Some(_) => { + return Err(FlowError::InvalidArgument( + "OpenAI Responses function_call_output id must be a string or null".into(), + )); + } + }; + Ok(Message::ToolResultItem { + id, + call_id: call_id.to_string(), + output: output.clone(), + extra: obj + .iter() + .filter(|(key, _)| { + !matches!(key.as_str(), "type" | "id" | "call_id" | "output") + }) + .map(|(key, value)| (key.clone(), value.clone())) + .collect(), + }) + } + _ => Ok(Message::ProviderNative { + provider: "openai_responses".into(), + kind: kind.into(), + value: value.clone(), + }), } } -fn encode_openai_responses_input( - obj: &mut serde_json::Map, - annotated: &AnnotatedLlmRequest, -) -> Result<()> { - let (system_text, input_messages) = split_system_and_input_messages(&annotated.messages); - set_or_remove_string(obj, "instructions", system_text); - if let Some(raw_input_items) = annotated.extra.get(UNPARSED_INPUT_ITEMS_KEY) { - obj.insert("input".into(), raw_input_items.clone()); - } else { - insert_serialized(obj, "input", &input_messages, "input")?; +fn encode_responses_content(content: &MessageContent, assistant: bool) -> Result { + match content { + MessageContent::Text(text) => Ok(Json::String(text.clone())), + MessageContent::Parts(parts) => Ok(Json::Array( + parts + .iter() + .map(|part| match part { + ContentPart::Text { text, extra } => { + let mut obj = extra.clone(); + obj.insert( + "type".into(), + Json::String( + if assistant { + "output_text" + } else { + "input_text" + } + .into(), + ), + ); + obj.insert("text".into(), Json::String(text.clone())); + Ok(Json::Object(obj)) + } + ContentPart::ImageUrl { image_url, extra } => { + let mut obj = extra.clone(); + obj.insert("type".into(), Json::String("input_image".into())); + obj.insert("image_url".into(), Json::String(image_url.url.clone())); + if let Some(detail) = &image_url.detail { + obj.insert("detail".into(), Json::String(detail.clone())); + } + Ok(Json::Object(obj)) + } + ContentPart::Image { image, extra } => { + let mut obj = image.as_object().cloned().ok_or_else(|| { + FlowError::InvalidArgument( + "OpenAI Responses image content must be an object".into(), + ) + })?; + obj.extend(extra.clone()); + obj.insert("type".into(), Json::String("input_image".into())); + Ok(Json::Object(obj)) + } + ContentPart::File { file, extra } => { + let mut obj = file.as_object().cloned().ok_or_else(|| { + FlowError::InvalidArgument( + "OpenAI Responses file content must be an object".into(), + ) + })?; + obj.extend(extra.clone()); + obj.insert("type".into(), Json::String("input_file".into())); + Ok(Json::Object(obj)) + } + ContentPart::Refusal { refusal, extra } if assistant => { + let mut obj = extra.clone(); + obj.insert("type".into(), Json::String("refusal".into())); + obj.insert("refusal".into(), Json::String(refusal.clone())); + Ok(Json::Object(obj)) + } + ContentPart::ProviderNative { + provider, value, .. + } if provider == "openai_responses" => Ok(value.clone()), + other => Err(FlowError::InvalidArgument(format!( + "content part {other:?} cannot be encoded for OpenAI Responses" + ))), + }) + .collect::>>()?, + )), } - Ok(()) } -fn encode_openai_responses_tools( - obj: &mut serde_json::Map, - annotated: &AnnotatedLlmRequest, -) -> Result<()> { - if let Some(ref tools) = annotated.tools { - insert_serialized(obj, "tools", tools, "tools")?; - } - if let Some(ref tool_choice) = annotated.tool_choice { - insert_serialized(obj, "tool_choice", tool_choice, "tool_choice")?; +fn encode_responses_input_item(message: &Message) -> Result { + match message { + Message::User { content, .. } + | Message::System { content, .. } + | Message::Developer { content, .. } => { + let role = match message { + Message::User { .. } => "user", + Message::System { .. } => "system", + Message::Developer { .. } => "developer", + _ => unreachable!(), + }; + let mut obj = serde_json::Map::new(); + obj.insert("type".into(), Json::String("message".into())); + obj.insert("role".into(), Json::String(role.into())); + obj.insert("content".into(), encode_responses_content(content, false)?); + Ok(Json::Object(obj)) + } + Message::Assistant { + content: Some(content), + .. + } => { + let mut obj = serde_json::Map::new(); + obj.insert("type".into(), Json::String("message".into())); + obj.insert("role".into(), Json::String("assistant".into())); + obj.insert("content".into(), encode_responses_content(content, true)?); + Ok(Json::Object(obj)) + } + Message::ToolCallItem { + id, + call_id, + name, + arguments, + extra, + } => { + let mut obj = extra.clone(); + obj.insert("type".into(), Json::String("function_call".into())); + if let Some(id) = id { + obj.insert("id".into(), Json::String(id.clone())); + } + obj.insert("call_id".into(), Json::String(call_id.clone())); + obj.insert("name".into(), Json::String(name.clone())); + let arguments = match arguments { + Json::String(raw) => raw.clone(), + value => serde_json::to_string(value).map_err(|error| { + FlowError::Internal(format!( + "OpenAI Responses function arguments encode: {error}" + )) + })?, + }; + obj.insert("arguments".into(), Json::String(arguments)); + Ok(Json::Object(obj)) + } + Message::ToolResultItem { + id, + call_id, + output, + extra, + } => { + let mut obj = extra.clone(); + obj.insert("type".into(), Json::String("function_call_output".into())); + if let Some(id) = id { + obj.insert("id".into(), Json::String(id.clone())); + } + obj.insert("call_id".into(), Json::String(call_id.clone())); + obj.insert("output".into(), output.clone()); + Ok(Json::Object(obj)) + } + Message::ProviderNative { + provider, value, .. + } if provider == "openai_responses" => Ok(value.clone()), + other => Err(FlowError::InvalidArgument(format!( + "message {other:?} cannot be encoded for OpenAI Responses" + ))), } - Ok(()) } -fn overlay_openai_responses_fields( - obj: &mut serde_json::Map, - annotated: &AnnotatedLlmRequest, -) { - if let Some(ref model) = annotated.model { - obj.insert("model".into(), Json::String(model.clone())); +fn decode_responses_tool(value: &Json) -> Result { + let obj = value.as_object().ok_or_else(|| { + FlowError::InvalidArgument("OpenAI Responses tool must be an object".into()) + })?; + if obj.get("type").and_then(Json::as_str) != Some("function") { + let kind = obj.get("type").and_then(Json::as_str).unwrap_or("unknown"); + return Ok(ToolDefinition::ProviderNative { + provider: "openai_responses".into(), + kind: kind.into(), + value: value.clone(), + }); } - overlay_openai_responses_json_fields(obj, annotated); - overlay_openai_responses_string_fields(obj, annotated); - overlay_openai_responses_bool_fields(obj, annotated); - overlay_openai_responses_u64_fields(obj, annotated); + let (function, wrapper_extra) = + if let Some(function) = obj.get("function").and_then(Json::as_object) { + ( + function, + obj.iter() + .filter(|(key, _)| !matches!(key.as_str(), "type" | "function")) + .map(|(key, value)| (key.clone(), value.clone())) + .collect(), + ) + } else { + (obj, serde_json::Map::new()) + }; + let name = function.get("name").and_then(Json::as_str).ok_or_else(|| { + FlowError::InvalidArgument("OpenAI Responses function tool is missing name".into()) + })?; + let description = + super::optional_string(function, "description", "OpenAI Responses function tool")?; + let strict = super::optional_bool(function, "strict", "OpenAI Responses function tool")?; + Ok(ToolDefinition::Function { + function: FunctionDefinition { + name: name.into(), + description, + parameters: function.get("parameters").cloned(), + strict, + extra: function + .iter() + .filter(|(key, _)| { + !matches!( + key.as_str(), + "type" | "name" | "description" | "parameters" | "strict" | "function" + ) + }) + .map(|(key, value)| (key.clone(), value.clone())) + .collect(), + }, + extra: wrapper_extra, + }) } -fn overlay_openai_responses_json_fields( - obj: &mut serde_json::Map, - annotated: &AnnotatedLlmRequest, -) { - for (key, value) in [ - ("truncation", &annotated.truncation), - ("reasoning", &annotated.reasoning), - ("include", &annotated.include), - ("metadata", &annotated.metadata), - ] { - if let Some(value) = value { - obj.insert(key.into(), value.clone()); +fn encode_responses_tool(tool: &ToolDefinition) -> Result { + match tool { + ToolDefinition::Function { function, extra } => { + let mut obj = extra.clone(); + let Json::Object(function) = encode_responses_function(function) else { + unreachable!("function definition encodes as an object") + }; + obj.extend(function); + obj.insert("type".into(), Json::String("function".into())); + Ok(Json::Object(obj)) } + ToolDefinition::ProviderNative { + provider, value, .. + } if provider == "openai_responses" => Ok(value.clone()), + other => Err(FlowError::InvalidArgument(format!( + "tool {other:?} cannot be encoded for OpenAI Responses" + ))), } } -fn overlay_openai_responses_string_fields( - obj: &mut serde_json::Map, - annotated: &AnnotatedLlmRequest, -) { - for (key, value) in [ - ("previous_response_id", &annotated.previous_response_id), - ("user", &annotated.user), - ("service_tier", &annotated.service_tier), - ] { - if let Some(value) = value { - obj.insert(key.into(), Json::String(value.clone())); - } +fn encode_responses_function(function: &FunctionDefinition) -> Json { + let mut obj = function.extra.clone(); + obj.insert("name".into(), Json::String(function.name.clone())); + if let Some(description) = &function.description { + obj.insert("description".into(), Json::String(description.clone())); + } + if let Some(parameters) = &function.parameters { + obj.insert("parameters".into(), parameters.clone()); } + if let Some(strict) = function.strict { + obj.insert("strict".into(), Json::Bool(strict)); + } + Json::Object(obj) } -fn overlay_openai_responses_bool_fields( - obj: &mut serde_json::Map, - annotated: &AnnotatedLlmRequest, -) { - for (key, value) in [ - ("store", annotated.store), - ("parallel_tool_calls", annotated.parallel_tool_calls), - ("stream", annotated.stream), - ] { - if let Some(value) = value { - obj.insert(key.into(), Json::Bool(value)); - } +fn patch_responses_tool( + original: &Json, + baseline: &ToolDefinition, + edited: &ToolDefinition, + baseline_value: &Json, + edited_value: &Json, +) -> Result { + if let ( + Some(original), + ToolDefinition::Function { + function: baseline_function, + extra: baseline_extra, + }, + ToolDefinition::Function { + function: edited_function, + extra: edited_extra, + }, + ) = (original.as_object(), baseline, edited) + && let Some(original_function) = original.get("function") + { + let mut patched = original.clone(); + patch_extra_fields(&mut patched, baseline_extra, edited_extra); + patched.insert( + "function".into(), + super::patch_changed_json( + original_function, + &encode_responses_function(baseline_function), + &encode_responses_function(edited_function), + )?, + ); + return Ok(Json::Object(patched)); } + + super::patch_changed_json(original, baseline_value, edited_value) } -fn overlay_openai_responses_u64_fields( - obj: &mut serde_json::Map, - annotated: &AnnotatedLlmRequest, -) { - for (key, value) in [ - ("max_output_tokens", annotated.max_output_tokens), - ("max_tool_calls", annotated.max_tool_calls), - ("top_logprobs", annotated.top_logprobs), - ] { - if let Some(value) = value { - obj.insert(key.into(), Json::from(value)); +fn decode_openai_or_anthropic_tool_choice(value: &Json) -> ToolChoice { + match value.as_str() { + Some("auto") => ToolChoice::Auto, + Some("none") => ToolChoice::None, + Some("required") => ToolChoice::Required, + _ => match value.as_object().and_then(|obj| { + let choice_type = obj.get("type").and_then(Json::as_str)?; + match choice_type { + "auto" => Some(ToolChoice::Auto), + "any" => Some(ToolChoice::Required), + "none" => Some(ToolChoice::None), + "tool" | "function" => obj + .get("name") + .and_then(Json::as_str) + .or_else(|| { + obj.get("function") + .and_then(Json::as_object) + .and_then(|function| function.get("name")) + .and_then(Json::as_str) + }) + .map(|name| { + ToolChoice::Specific(ToolChoiceFunction { + choice_type: "function".into(), + function: ToolChoiceFunctionName { name: name.into() }, + }) + }), + _ => None, + } + }) { + Some(choice) => choice, + None => ToolChoice::ProviderNative(responses_native("tool_choice", value)), + }, + } +} + +fn encode_responses_tool_choice(choice: &ToolChoice) -> Result { + match choice { + ToolChoice::Auto => Ok(Json::String("auto".into())), + ToolChoice::None => Ok(Json::String("none".into())), + ToolChoice::Required => Ok(Json::String("required".into())), + ToolChoice::Specific(choice) => Ok(serde_json::json!({ + "type":"function", + "name":choice.function.name, + })), + ToolChoice::ProviderNative(native) if native.provider == "openai_responses" => { + Ok(native.value.clone()) } + ToolChoice::ProviderNative(native) => Err(FlowError::InvalidArgument(format!( + "tool choice for {} cannot be encoded for OpenAI Responses", + native.provider + ))), } } -fn merge_openai_responses_extra_fields( +fn patch_extra_fields( obj: &mut serde_json::Map, - extra: &serde_json::Map, + baseline: &serde_json::Map, + edited: &serde_json::Map, ) { - for (k, v) in extra { - if k != UNPARSED_INPUT_ITEMS_KEY { - obj.insert(k.clone(), v.clone()); + for key in baseline.keys().filter(|key| !edited.contains_key(*key)) { + obj.remove(key); + } + for (key, value) in edited { + if baseline.get(key) != Some(value) { + obj.insert(key.clone(), value.clone()); } } } -fn decode_openai_or_anthropic_tool_choice(value: &Json) -> Option { - if let Ok(parsed) = serde_json::from_value::(value.clone()) { - return Some(parsed); +fn set_or_remove_json(obj: &mut serde_json::Map, key: &str, value: Option) { + if let Some(value) = value { + obj.insert(key.into(), value); + } else { + obj.remove(key); } +} - let obj = value.as_object()?; - match obj.get("type").and_then(|v| v.as_str()) { - Some("auto") => Some(ToolChoice::Auto), - Some("any") => Some(ToolChoice::Required), - Some("none") => Some(ToolChoice::None), - Some("tool") => { - let name = obj.get("name").and_then(|v| v.as_str())?.to_string(); - Some(ToolChoice::Specific(ToolChoiceFunction { - choice_type: "function".to_string(), - function: ToolChoiceFunctionName { name }, - })) +fn patch_responses_api_specific( + obj: &mut serde_json::Map, + edited: &Option, + baseline: &Option, +) -> Result<()> { + match (edited, baseline) { + ( + Some(ApiSpecificRequest::OpenAIResponses { + background, + context_management, + conversation, + moderation, + prompt, + prompt_cache_key, + prompt_cache_options, + prompt_cache_retention, + safety_identifier, + stream_options, + text, + }), + Some(ApiSpecificRequest::OpenAIResponses { + background: old_background, + context_management: old_context_management, + conversation: old_conversation, + moderation: old_moderation, + prompt: old_prompt, + prompt_cache_key: old_prompt_cache_key, + prompt_cache_options: old_prompt_cache_options, + prompt_cache_retention: old_prompt_cache_retention, + safety_identifier: old_safety_identifier, + stream_options: old_stream_options, + text: old_text, + }), + ) => { + if background != old_background { + set_or_remove_json(obj, "background", background.map(Json::Bool)); + } + for (key, value, old_value) in [ + ( + "context_management", + context_management, + old_context_management, + ), + ("conversation", conversation, old_conversation), + ("moderation", moderation, old_moderation), + ("prompt", prompt, old_prompt), + ( + "prompt_cache_options", + prompt_cache_options, + old_prompt_cache_options, + ), + ("stream_options", stream_options, old_stream_options), + ("text", text, old_text), + ] { + if value != old_value { + set_or_remove_json(obj, key, value.clone()); + } + } + for (key, value, old_value) in [ + ("prompt_cache_key", prompt_cache_key, old_prompt_cache_key), + ( + "prompt_cache_retention", + prompt_cache_retention, + old_prompt_cache_retention, + ), + ( + "safety_identifier", + safety_identifier, + old_safety_identifier, + ), + ] { + if value != old_value { + set_or_remove_json(obj, key, value.clone().map(Json::String)); + } + } + Ok(()) } - _ => None, + (None, Some(ApiSpecificRequest::OpenAIResponses { .. })) => { + for key in [ + "background", + "context_management", + "conversation", + "moderation", + "prompt", + "prompt_cache_key", + "prompt_cache_options", + "prompt_cache_retention", + "safety_identifier", + "stream_options", + "text", + ] { + obj.remove(key); + } + Ok(()) + } + (Some(_), _) => Err(FlowError::InvalidArgument( + "api_specific provider does not match OpenAI Responses".into(), + )), + (None, Some(_)) => Err(FlowError::InvalidArgument( + "api_specific provider does not match OpenAI Responses".into(), + )), + (None, None) => Ok(()), } } fn decode_openai_or_anthropic_parallel_tool_calls( obj: &serde_json::Map, -) -> Option { - if let Some(value) = obj.get("parallel_tool_calls").and_then(|v| v.as_bool()) { - return Some(value); +) -> Result> { + if let Some(value) = super::optional_bool(obj, "parallel_tool_calls", "OpenAI Responses")? { + return Ok(Some(value)); } - let tool_choice = obj.get("tool_choice")?.as_object()?; - tool_choice - .get("disable_parallel_tool_use") - .and_then(|v| v.as_bool()) - .map(|disabled| !disabled) + let Some(tool_choice) = obj.get("tool_choice").and_then(Json::as_object) else { + return Ok(None); + }; + Ok(super::optional_bool( + tool_choice, + "disable_parallel_tool_use", + "OpenAI Responses tool_choice", + )? + .map(|disabled| !disabled)) } // --------------------------------------------------------------------------- @@ -565,47 +1080,39 @@ impl LlmCodec for OpenAIResponsesCodec { .content .as_object() .ok_or_else(|| FlowError::Internal("request content is not an object".into()))?; - - let mut messages: Vec = Vec::new(); - let mut preserved_unparsed_input: Option = None; - - // Extract instructions -> system message (first). - if let Some(instructions) = obj.get("instructions").and_then(|v| v.as_str()) { - messages.push(Message::System { - content: MessageContent::Text(instructions.to_string()), + let input = obj.get("input").ok_or_else(|| { + FlowError::InvalidArgument("OpenAI Responses request is missing input".into()) + })?; + let messages = if let Some(input) = input.as_str() { + vec![Message::User { + content: MessageContent::Text(input.to_string()), name: None, - }); - } - - // Extract input. - if let Some(input) = obj.get("input") { - if let Some(s) = input.as_str() { - // Input is a simple string -> single User message. - messages.push(Message::User { - content: MessageContent::Text(s.to_string()), - name: None, - }); - } else if input.is_array() { - // Strict-first parse to avoid partial normalized state. - match serde_json::from_value::>(input.clone()) { - Ok(input_messages) => messages.extend(input_messages), - Err(_) => { - // Preserve full original array for lossless handling. - preserved_unparsed_input = Some(input.clone()); - } - } + }] + } else { + input + .as_array() + .ok_or_else(|| { + FlowError::InvalidArgument( + "OpenAI Responses input must be a string or an array".into(), + ) + })? + .iter() + .map(decode_responses_input_item) + .collect::>>()? + }; + let instructions = match obj.get("instructions") { + Some(Json::String(instructions)) => Some(MessageContent::Text(instructions.clone())), + Some(Json::Null) | None => None, + Some(_) => { + return Err(FlowError::InvalidArgument( + "OpenAI Responses instructions must be a string or null".into(), + )); } - } - - // Extract model. - let model = obj.get("model").and_then(|v| v.as_str()).map(String::from); - - // Extract generation params. - let temperature = obj.get("temperature").and_then(|v| v.as_f64()); - let top_p = obj.get("top_p").and_then(|v| v.as_f64()); - let max_tokens = obj.get("max_output_tokens").and_then(|v| v.as_u64()); - // Responses API does not support stop sequences. - + }; + let model = super::optional_string(obj, "model", "OpenAI Responses")?; + let temperature = super::optional_f64(obj, "temperature", "OpenAI Responses")?; + let top_p = super::optional_f64(obj, "top_p", "OpenAI Responses")?; + let max_tokens = super::optional_u64(obj, "max_output_tokens", "OpenAI Responses")?; let params = if temperature.is_some() || max_tokens.is_some() || top_p.is_some() { Some(GenerationParams { temperature, @@ -616,71 +1123,269 @@ impl LlmCodec for OpenAIResponsesCodec { } else { None }; - - // Extract tools. - let tools: Option> = obj + let tools = obj .get("tools") - .map(|v| serde_json::from_value(v.clone())) - .transpose() - .map_err(|e| FlowError::Internal(format!("OpenAI Responses tools decode: {e}")))?; - - // Extract tool_choice. - let tool_choice: Option = obj + .map(|value| { + value + .as_array() + .ok_or_else(|| { + FlowError::InvalidArgument("OpenAI Responses tools must be an array".into()) + })? + .iter() + .map(decode_responses_tool) + .collect::>>() + }) + .transpose()?; + let tool_choice = obj .get("tool_choice") - .and_then(decode_openai_or_anthropic_tool_choice); - - // Collect extra fields (keys not in MODELED_REQUEST_KEYS). - let mut extra: serde_json::Map = obj + .map(decode_openai_or_anthropic_tool_choice); + let store = super::optional_bool(obj, "store", "OpenAI Responses")?; + let previous_response_id = + super::optional_string(obj, "previous_response_id", "OpenAI Responses")?; + let user = super::optional_string(obj, "user", "OpenAI Responses")?; + let service_tier = super::optional_string(obj, "service_tier", "OpenAI Responses")?; + let parallel_tool_calls = decode_openai_or_anthropic_parallel_tool_calls(obj)?; + let max_tool_calls = super::optional_u64(obj, "max_tool_calls", "OpenAI Responses")?; + let top_logprobs = super::optional_u64(obj, "top_logprobs", "OpenAI Responses")?; + let stream = super::optional_bool(obj, "stream", "OpenAI Responses")?; + let background = super::optional_bool(obj, "background", "OpenAI Responses")?; + let prompt_cache_key = super::optional_string(obj, "prompt_cache_key", "OpenAI Responses")?; + let prompt_cache_retention = + super::optional_string(obj, "prompt_cache_retention", "OpenAI Responses")?; + let safety_identifier = + super::optional_string(obj, "safety_identifier", "OpenAI Responses")?; + let reasoning = super::optional_object(obj, "reasoning", "OpenAI Responses")?; + let include = super::optional_array(obj, "include", "OpenAI Responses")?; + let metadata = super::optional_object(obj, "metadata", "OpenAI Responses")?; + let context_management = + super::optional_array(obj, "context_management", "OpenAI Responses")?; + let moderation = super::optional_object(obj, "moderation", "OpenAI Responses")?; + let prompt = super::optional_object(obj, "prompt", "OpenAI Responses")?; + let prompt_cache_options = + super::optional_object(obj, "prompt_cache_options", "OpenAI Responses")?; + let stream_options = super::optional_object(obj, "stream_options", "OpenAI Responses")?; + let text = super::optional_object(obj, "text", "OpenAI Responses")?; + let extra: serde_json::Map = obj .iter() .filter(|(k, _)| !MODELED_REQUEST_KEYS.contains(&k.as_str())) .map(|(k, v)| (k.clone(), v.clone())) .collect(); - if let Some(input_items) = preserved_unparsed_input { - extra.insert(UNPARSED_INPUT_ITEMS_KEY.into(), input_items); - } - Ok(AnnotatedLlmRequest { messages, + instructions, model, params, tools, tool_choice, - store: obj.get("store").and_then(|v| v.as_bool()), - previous_response_id: obj - .get("previous_response_id") - .and_then(|v| v.as_str()) - .map(String::from), + store, + previous_response_id, truncation: obj.get("truncation").cloned(), - reasoning: obj.get("reasoning").cloned(), - include: obj.get("include").cloned(), - user: obj.get("user").and_then(|v| v.as_str()).map(String::from), - metadata: obj.get("metadata").cloned(), - service_tier: obj - .get("service_tier") - .and_then(|v| v.as_str()) - .map(String::from), - parallel_tool_calls: decode_openai_or_anthropic_parallel_tool_calls(obj), - max_output_tokens: obj.get("max_output_tokens").and_then(|v| v.as_u64()), - max_tool_calls: obj.get("max_tool_calls").and_then(|v| v.as_u64()), - top_logprobs: obj.get("top_logprobs").and_then(|v| v.as_u64()), - stream: obj.get("stream").and_then(|v| v.as_bool()), + reasoning, + include, + user, + metadata, + service_tier, + parallel_tool_calls, + max_output_tokens: max_tokens, + max_tool_calls, + top_logprobs, + stream, + api_specific: Some(ApiSpecificRequest::OpenAIResponses { + background, + context_management, + conversation: obj.get("conversation").cloned(), + moderation, + prompt, + prompt_cache_key, + prompt_cache_options, + prompt_cache_retention, + safety_identifier, + stream_options, + text, + }), extra, }) } fn encode(&self, annotated: &AnnotatedLlmRequest, original: &LlmRequest) -> Result { + let baseline = self.decode(original)?; let mut content = original.content.clone(); let obj = content .as_object_mut() .ok_or_else(|| FlowError::Internal("original content is not an object".into()))?; - - encode_openai_responses_input(obj, annotated)?; - if let Some(ref params) = annotated.params { - overlay_generation_params(obj, params); + if annotated.messages != baseline.messages { + let input = if original.content.get("input").is_some_and(Json::is_string) + && matches!( + annotated.messages.as_slice(), + [Message::User { + content: MessageContent::Text(_), + name: None + }] + ) { + match &annotated.messages[0] { + Message::User { + content: MessageContent::Text(text), + .. + } => Json::String(text.clone()), + _ => unreachable!(), + } + } else { + Json::Array(super::encode_changed_items( + &annotated.messages, + &baseline.messages, + original + .content + .get("input") + .and_then(Json::as_array) + .map(Vec::as_slice), + encode_responses_input_item, + )?) + }; + obj.insert("input".into(), input); + } + if annotated.instructions != baseline.instructions { + let instructions = match &annotated.instructions { + Some(MessageContent::Text(text)) => Some(Json::String(text.clone())), + Some(MessageContent::Parts(_)) => { + return Err(FlowError::InvalidArgument( + "OpenAI Responses instructions cannot contain content parts".into(), + )); + } + None => None, + }; + set_or_remove_json(obj, "instructions", instructions); + } + if annotated.model != baseline.model { + set_or_remove_json(obj, "model", annotated.model.clone().map(Json::String)); + } + if annotated.params != baseline.params { + let edited = annotated.params.as_ref(); + let before = baseline.params.as_ref(); + if edited.and_then(|params| params.stop.as_ref()) + != before.and_then(|params| params.stop.as_ref()) + && edited.and_then(|params| params.stop.as_ref()).is_some() + { + return Err(FlowError::InvalidArgument( + "OpenAI Responses does not support stop sequences".into(), + )); + } + for (key, value, old_value) in [ + ( + "temperature", + edited.and_then(|params| params.temperature), + before.and_then(|params| params.temperature), + ), + ( + "top_p", + edited.and_then(|params| params.top_p), + before.and_then(|params| params.top_p), + ), + ] { + if value != old_value { + set_or_remove_json(obj, key, value.map(json_f64)); + } + } + let max_tokens = edited.and_then(|params| params.max_tokens); + let old_max_tokens = before.and_then(|params| params.max_tokens); + if max_tokens != old_max_tokens { + set_or_remove_json(obj, "max_output_tokens", max_tokens.map(Json::from)); + } + } + if annotated.tools != baseline.tools { + let tools = annotated + .tools + .as_deref() + .map(|tools| { + super::encode_changed_items_with_patch( + tools, + baseline.tools.as_deref().unwrap_or(&[]), + obj.get("tools").and_then(Json::as_array).map(Vec::as_slice), + encode_responses_tool, + patch_responses_tool, + ) + }) + .transpose()? + .map(Json::Array); + set_or_remove_json(obj, "tools", tools); + } + if annotated.tool_choice != baseline.tool_choice { + let tool_choice = match (&annotated.tool_choice, &baseline.tool_choice) { + (Some(edited), Some(before)) => { + let edited = encode_responses_tool_choice(edited)?; + let before = encode_responses_tool_choice(before)?; + Some(match obj.get("tool_choice") { + Some(original) => super::patch_changed_json(original, &before, &edited)?, + None => edited, + }) + } + (Some(edited), None) => Some(encode_responses_tool_choice(edited)?), + (None, _) => None, + }; + set_or_remove_json(obj, "tool_choice", tool_choice); + } + for (key, value, old_value) in [ + ("truncation", &annotated.truncation, &baseline.truncation), + ("reasoning", &annotated.reasoning, &baseline.reasoning), + ("include", &annotated.include, &baseline.include), + ("metadata", &annotated.metadata, &baseline.metadata), + ] { + if value != old_value { + set_or_remove_json(obj, key, value.clone()); + } + } + for (key, value, old_value) in [ + ( + "previous_response_id", + &annotated.previous_response_id, + &baseline.previous_response_id, + ), + ("user", &annotated.user, &baseline.user), + ( + "service_tier", + &annotated.service_tier, + &baseline.service_tier, + ), + ] { + if value != old_value { + set_or_remove_json(obj, key, value.clone().map(Json::String)); + } + } + for (key, value, old_value) in [ + ("store", annotated.store, baseline.store), + ( + "parallel_tool_calls", + annotated.parallel_tool_calls, + baseline.parallel_tool_calls, + ), + ("stream", annotated.stream, baseline.stream), + ] { + if value != old_value { + set_or_remove_json(obj, key, value.map(Json::Bool)); + } + } + for (key, value, old_value) in [ + ( + "max_output_tokens", + annotated.max_output_tokens, + baseline.max_output_tokens, + ), + ( + "max_tool_calls", + annotated.max_tool_calls, + baseline.max_tool_calls, + ), + ( + "top_logprobs", + annotated.top_logprobs, + baseline.top_logprobs, + ), + ] { + if value != old_value { + set_or_remove_json(obj, key, value.map(Json::from)); + } } - encode_openai_responses_tools(obj, annotated)?; - overlay_openai_responses_fields(obj, annotated); - merge_openai_responses_extra_fields(obj, &annotated.extra); + patch_responses_api_specific(obj, &annotated.api_specific, &baseline.api_specific)?; + patch_extra_fields(obj, &baseline.extra, &annotated.extra); Ok(LlmRequest { headers: original.headers.clone(), diff --git a/crates/core/src/observability/atof.rs b/crates/core/src/observability/atof.rs index f6f312c36..b81bf8df3 100644 --- a/crates/core/src/observability/atof.rs +++ b/crates/core/src/observability/atof.rs @@ -477,7 +477,11 @@ impl AtofExporter { Ok(removed) } - /// Flush the underlying file and drain queued endpoint events. + /// Outside a native subscriber callback, wait for queued subscriber delivery, then flush the + /// configured file sink or ask the configured stream sink to drain for up to its timeout. + /// + /// A re-entrant call does not establish the delivery barrier. A stream timeout is logged and + /// does not by itself return an error. pub fn force_flush(&self) -> Result<()> { flush_subscribers()?; let mut state = self @@ -515,7 +519,11 @@ impl AtofExporter { ) } - /// Shut down the exporter by flushing buffered data and closing endpoints. + /// Outside a native subscriber callback, wait for queued subscriber delivery, then flush the + /// configured file sink or ask the configured stream sink to drain and close up to its timeout. + /// + /// A re-entrant call does not establish the delivery barrier. A stream timeout is logged and + /// does not by itself return an error. pub fn shutdown(&self) -> Result<()> { flush_subscribers()?; let mut state = self diff --git a/crates/core/src/observability/openinference.rs b/crates/core/src/observability/openinference.rs index 8b658d9e7..28f122eef 100644 --- a/crates/core/src/observability/openinference.rs +++ b/crates/core/src/observability/openinference.rs @@ -1095,15 +1095,37 @@ fn push_optimization_attributes( fn push_annotated_input_messages(attributes: &mut Vec, messages: &[Message]) { for (index, message) in messages.iter().enumerate() { - let (role, content) = match message { - Message::System { content, .. } => ("system", Some(content)), - Message::User { content, .. } => ("user", Some(content)), - Message::Assistant { content, .. } => ("assistant", content.as_ref()), - Message::Tool { content, .. } => ("tool", Some(content)), + let role = match message { + Message::System { .. } => "system", + Message::Developer { .. } => "developer", + Message::User { .. } => "user", + Message::Assistant { .. } => "assistant", + Message::Tool { .. } => "tool", + Message::Function { .. } => "function", + Message::ToolCallItem { .. } => "assistant", + Message::ToolResultItem { .. } => "tool", + Message::ProviderNative { value, .. } => value + .get("role") + .and_then(Json::as_str) + .unwrap_or("provider_native"), }; push_message_role(attributes, "llm.input_messages", index, role); + let content = match message { + Message::System { content, .. } + | Message::Developer { content, .. } + | Message::User { content, .. } + | Message::Tool { content, .. } => message_content_text(content), + Message::Assistant { content, .. } => content.as_ref().and_then(message_content_text), + Message::Function { content, .. } => { + content.as_deref().and_then(display_text_from_string) + } + Message::ProviderNative { value, .. } => { + value.get("content").and_then(display_text_from_json) + } + Message::ToolCallItem { .. } | Message::ToolResultItem { .. } => None, + }; if let Some(content) = content { - push_message_text_content(attributes, "llm.input_messages", index, content); + push_message_text_value(attributes, "llm.input_messages", index, content); } } } @@ -1148,18 +1170,16 @@ fn push_message_role( )); } -fn push_message_text_content( +fn push_message_text_value( attributes: &mut Vec, prefix: &'static str, index: usize, - content: &MessageContent, + text: String, ) { - if let Some(text) = message_content_text(content) { - attributes.push(KeyValue::new( - format!("{prefix}.{index}.message.content"), - text, - )); - } + attributes.push(KeyValue::new( + format!("{prefix}.{index}.message.content"), + text, + )); } fn message_content_text(content: &MessageContent) -> Option { @@ -1169,8 +1189,13 @@ fn message_content_text(content: &MessageContent) -> Option { let text = parts .iter() .filter_map(|part| match part { - ContentPart::Text { text } => Some(text.as_str()), - ContentPart::ImageUrl { .. } => None, + ContentPart::Text { text, .. } => Some(text.as_str()), + ContentPart::Refusal { refusal, .. } => Some(refusal.as_str()), + ContentPart::ProviderNative { value, .. } => value + .get("text") + .and_then(Json::as_str) + .or_else(|| value.get("refusal").and_then(Json::as_str)), + _ => None, }) .collect::>() .join("\n") diff --git a/crates/core/src/plugin/dynamic.rs b/crates/core/src/plugin/dynamic.rs index 7dd546c84..fd9ff3653 100644 --- a/crates/core/src/plugin/dynamic.rs +++ b/crates/core/src/plugin/dynamic.rs @@ -9,10 +9,13 @@ //! one file as the feature grows. use chrono::Utc; +use semver::{Version, VersionReq}; use serde::{Deserialize, Serialize}; use strum::{Display, IntoStaticStr}; -use crate::plugin::{PluginDeregistrationOutcome, deregister_plugin_registration_checked}; +use crate::plugin::{ + PluginDeregistrationOutcome, PluginError, deregister_plugin_registration_checked, +}; /// Canonical identifier for one dynamic plugin record. pub type DynamicPluginId = String; @@ -90,6 +93,21 @@ pub(super) fn deregister_tracked_registrations_checked( outcome } +pub(super) fn validate_annotated_request_consumer_compatibility( + relay: &str, + plugin_kind: &str, +) -> crate::plugin::Result<()> { + let requirement = VersionReq::parse(relay).map_err(|error| { + PluginError::InvalidConfig(format!("invalid compat.relay version requirement: {error}")) + })?; + if requirement.matches(&Version::new(0, 5, u64::MAX)) { + return Err(PluginError::InvalidConfig(format!( + "dynamic plugin '{plugin_kind}' registers an LLM request intercept and must declare compat.relay = \">=0.6,<1.0\" or another range that excludes Relay 0.5" + ))); + } + Ok(()) +} + /// Plugin execution lane. #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash, Display)] #[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] diff --git a/crates/core/src/plugin/dynamic/native.rs b/crates/core/src/plugin/dynamic/native.rs index 6b18542d2..be1b658a9 100644 --- a/crates/core/src/plugin/dynamic/native.rs +++ b/crates/core/src/plugin/dynamic/native.rs @@ -58,6 +58,7 @@ use crate::plugin::{ use super::{ DynamicPluginKind, DynamicPluginManifest, DynamicPluginManifestLoad, DynamicPluginTeardownOutcome, deregister_tracked_registrations_checked, + validate_annotated_request_consumer_compatibility, }; /// Native plugin load request derived from host dynamic-plugin state. @@ -258,6 +259,7 @@ fn native_error_diagnostic(plugin_kind: &str, code: &str, message: &str) -> Conf struct NativePluginInstance { plugin_kind: String, + relay_compat: String, allows_multiple_components: bool, plugin: Mutex, _library: Library, @@ -302,6 +304,12 @@ fn load_one_native_plugin( ))); } validate_relay_compatibility(manifest.compat.relay.as_deref())?; + let relay_compat = manifest + .compat + .relay + .as_deref() + .expect("validated native manifest must declare compat.relay") + .to_string(); if manifest.compat.native_api.as_deref().map(str::trim) != Some("1") { return Err(PluginError::InvalidConfig(format!( "dynamic plugin '{}' declares unsupported compat.native_api '{}'; expected 1", @@ -384,6 +392,7 @@ fn load_one_native_plugin( } Ok(Arc::new(NativePluginInstance { plugin_kind, + relay_compat, allows_multiple_components: plugin.allows_multiple_components, plugin: Mutex::new(plugin), _library: library, @@ -1350,6 +1359,12 @@ unsafe extern "C" fn native_plugin_context_register_llm_request_intercept( Err(status) => return status, }; let instance = host_ctx.instance.clone(); + if let Err(error) = validate_annotated_request_consumer_compatibility( + &instance.relay_compat, + &instance.plugin_kind, + ) { + return status_from_plugin_error(error); + } let ctx = unsafe { &mut *host_ctx.ctx }; let name = match read_name(name) { Ok(name) => name, diff --git a/crates/core/src/plugin/dynamic/worker.rs b/crates/core/src/plugin/dynamic/worker.rs index b8e913fc0..ce25e4935 100644 --- a/crates/core/src/plugin/dynamic/worker.rs +++ b/crates/core/src/plugin/dynamic/worker.rs @@ -53,7 +53,7 @@ use tokio_stream::wrappers::UnixListenerStream; use tower::service_fn; use crate::api::event::{Event, EventSanitizeFields}; -use crate::api::llm::LlmRequest; +use crate::api::llm::{LLM_REQUEST_INTERCEPT_OUTCOME_SCHEMA, LlmRequest}; use crate::api::runtime::{ LlmExecutionNextFn, LlmJsonStream, LlmStreamExecutionNextFn, ToolExecutionNextFn, current_scope_stack, with_scope_stack, @@ -63,7 +63,7 @@ use crate::api::scope::{ event as emit_scope_mark, pop_scope, push_scope, }; use crate::api::tool::ToolExecutionInterceptOutcome; -use crate::codec::request::AnnotatedLlmRequest; +use crate::codec::request::{ANNOTATED_LLM_REQUEST_SCHEMA, AnnotatedLlmRequest}; use crate::error::{FlowError, Result as FlowResult}; use crate::plugin::{ ConfigDiagnostic, DiagnosticLevel, Plugin, PluginError, PluginRegistrationContext, @@ -73,12 +73,12 @@ use crate::plugin::{ use super::{ DynamicPluginKind, DynamicPluginManifest, DynamicPluginManifestLoad, DynamicPluginTeardownOutcome, WorkerRuntime, deregister_tracked_registrations_checked, + validate_annotated_request_consumer_compatibility, }; const JSON_SCHEMA: &str = "nemo.relay.Json@1"; const EVENT_SCHEMA: &str = "nemo.relay.Event@1"; const LLM_REQUEST_SCHEMA: &str = "nemo.relay.LlmRequest@1"; -const ANNOTATED_LLM_REQUEST_SCHEMA: &str = "nemo.relay.AnnotatedLlmRequest@1"; const WORKER_STARTUP_TIMEOUT: Duration = Duration::from_secs(10); const WORKER_RPC_TIMEOUT: Duration = Duration::from_secs(30); const WORKER_CONNECT_RETRY: Duration = Duration::from_millis(25); @@ -448,6 +448,12 @@ fn load_one_worker_plugin( ))); } validate_relay_compatibility(manifest.compat.relay.as_deref())?; + let relay_compat = manifest + .compat + .relay + .as_deref() + .expect("validated worker manifest must declare compat.relay") + .to_string(); let DynamicPluginManifestLoad::Worker(load) = &manifest.load else { unreachable!("validated worker manifest must carry worker load contract"); }; @@ -606,6 +612,12 @@ fn load_one_worker_plugin( validate_registration_plan(&spec.plugin_id, ®ister)?; register.registrations }; + if registrations.iter().any(|registration| { + RegistrationSurface::try_from(registration.surface) + .is_ok_and(|surface| surface == RegistrationSurface::LlmRequestIntercept) + }) { + validate_annotated_request_consumer_compatibility(&relay_compat, &spec.plugin_id)?; + } log::info!( target: "nemo_relay.worker", @@ -1626,7 +1638,7 @@ impl WorkerPluginCallback { match response.result { Some(invoke_response_result::Result::LlmRequest(result)) => { let outcome = required_envelope(result.outcome, "llm request intercept outcome")?; - if outcome.schema != "nemo.relay.LlmRequestInterceptOutcome@1" { + if outcome.schema != LLM_REQUEST_INTERCEPT_OUTCOME_SCHEMA { return Err(FlowError::Internal(format!( "worker returned unsupported LLM request intercept outcome schema: {}", outcome.schema diff --git a/crates/core/tests/integration/api_surface_tests.rs b/crates/core/tests/integration/api_surface_tests.rs index a8eb0536d..6b3d83bd9 100644 --- a/crates/core/tests/integration/api_surface_tests.rs +++ b/crates/core/tests/integration/api_surface_tests.rs @@ -388,6 +388,47 @@ fn tool_start_eagerly_emits_deduplicated_skill_load_marks() { deregister_subscriber("skill-load-api-events").unwrap(); } +#[test] +fn mcp_resource_read_emits_minimal_tool_parented_skill_load_mark() { + let _lock = TEST_MUTEX.lock().unwrap(); + reset_global(); + setup_isolated_thread(); + + let events = capture_events("mcp-resource-skill-load-api-events"); + let tool_handle = tool_call( + nemo_relay::api::tool::ToolCallParams::builder() + .name("mcp__resources__read_resource") + .args(json!({"uri": "file:///skills/review/SKILL.md"})) + .build(), + ) + .unwrap(); + tool_call_end( + nemo_relay::api::tool::ToolCallEndParams::builder() + .handle(&tool_handle) + .result(json!({"ok": true})) + .build(), + ) + .unwrap(); + + let captured = captured_events_snapshot(&events); + assert_eq!(captured.len(), 3); + assert_eq!(captured[0].name(), "mcp__resources__read_resource"); + assert_eq!(captured[0].scope_category(), Some(ScopeCategory::Start)); + assert_eq!(captured[1].name(), "skill.load"); + assert_eq!(captured[1].parent_uuid(), Some(tool_handle.uuid)); + assert_eq!(captured[1].data(), Some(&json!({"skill_name": "review"}))); + assert_eq!( + captured[1].metadata(), + Some(&json!({ + "skill_load_source": "structured_read", + "tool_name": "mcp__resources__read_resource" + })) + ); + assert_eq!(captured[2].scope_category(), Some(ScopeCategory::End)); + + deregister_subscriber("mcp-resource-skill-load-api-events").unwrap(); +} + #[test] fn integration_owned_skill_load_metadata_suppresses_core_detection() { let _lock = TEST_MUTEX.lock().unwrap(); diff --git a/crates/core/tests/integration/codec_tests.rs b/crates/core/tests/integration/codec_tests.rs index 12ae1bdda..a2581c72b 100644 --- a/crates/core/tests/integration/codec_tests.rs +++ b/crates/core/tests/integration/codec_tests.rs @@ -26,6 +26,8 @@ struct MockCodec { impl LlmCodec for MockCodec { fn decode(&self, _request: &LlmRequest) -> Result { Ok(AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![], model: Some(self.id.clone()), params: None, @@ -71,6 +73,8 @@ fn dummy_llm_request() -> LlmRequest { #[test] fn test_annotated_llm_request_full_roundtrip() { let req = AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![ Message::System { content: MessageContent::Text("Be helpful".into()), @@ -104,15 +108,17 @@ fn test_annotated_llm_request_full_roundtrip() { top_p: Some(0.9), stop: Some(vec!["END".into()]), }), - tools: Some(vec![ToolDefinition { - tool_type: "function".into(), + tools: Some(vec![ToolDefinition::Function { function: FunctionDefinition { name: "search".into(), description: Some("Search the web".into()), parameters: Some( json!({"type": "object", "properties": {"q": {"type": "string"}}}), ), + strict: None, + extra: serde_json::Map::new(), }, + extra: serde_json::Map::new(), }]), tool_choice: Some(ToolChoice::Auto), store: None, @@ -143,6 +149,8 @@ fn test_annotated_llm_request_full_roundtrip() { #[test] fn test_annotated_llm_request_minimal() { let req = AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![Message::User { content: MessageContent::Text("Hi".into()), name: None, @@ -279,7 +287,10 @@ fn test_message_content_text_serializes_as_string() { #[test] fn test_message_content_parts_serializes_as_array() { - let content = MessageContent::Parts(vec![ContentPart::Text { text: "hi".into() }]); + let content = MessageContent::Parts(vec![ContentPart::Text { + text: "hi".into(), + extra: serde_json::Map::new(), + }]); let json_val = serde_json::to_value(&content).unwrap(); assert_eq!(json_val, json!([{"type": "text", "text": "hi"}])); } @@ -310,13 +321,15 @@ fn test_tool_call_roundtrip() { #[test] fn test_tool_definition_roundtrip() { - let td = ToolDefinition { - tool_type: "function".into(), + let td = ToolDefinition::Function { function: FunctionDefinition { name: "get_weather".into(), description: Some("Get current weather".into()), parameters: Some(json!({"type": "object", "properties": {"city": {"type": "string"}}})), + strict: None, + extra: serde_json::Map::new(), }, + extra: serde_json::Map::new(), }; let json_val = serde_json::to_value(&td).unwrap(); assert_eq!( @@ -399,6 +412,8 @@ fn test_extra_field_flatten() { extra.insert("response_format".into(), json!({"type": "json_object"})); let req = AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![Message::User { content: MessageContent::Text("hi".into()), name: None, @@ -440,6 +455,8 @@ fn test_extra_field_flatten() { #[test] fn test_clone_and_partial_eq() { let a = AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![Message::User { content: MessageContent::Text("hello".into()), name: None, @@ -465,6 +482,8 @@ fn test_clone_and_partial_eq() { }; let b = AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![Message::User { content: MessageContent::Text("hello".into()), name: None, @@ -505,6 +524,8 @@ fn test_clone_and_partial_eq() { #[test] fn test_system_prompt_text() { let req = AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![ Message::System { content: MessageContent::Text("Be helpful".into()), @@ -540,6 +561,8 @@ fn test_system_prompt_text() { #[test] fn test_system_prompt_none() { let req = AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![Message::User { content: MessageContent::Text("Hi".into()), name: None, @@ -569,9 +592,12 @@ fn test_system_prompt_none() { #[test] fn test_system_prompt_parts() { let req = AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![Message::System { content: MessageContent::Parts(vec![ContentPart::Text { text: "Be careful".into(), + extra: serde_json::Map::new(), }]), name: None, }], @@ -600,6 +626,8 @@ fn test_system_prompt_parts() { #[test] fn test_last_user_message_basic() { let req = AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![ Message::User { content: MessageContent::Text("first".into()), @@ -640,6 +668,8 @@ fn test_last_user_message_basic() { #[test] fn test_last_user_message_none() { let req = AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![ Message::System { content: MessageContent::Text("hi".into()), @@ -676,6 +706,8 @@ fn test_last_user_message_none() { #[test] fn test_has_tool_calls_true() { let req = AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![Message::Assistant { content: None, tool_calls: Some(vec![ToolCall { @@ -713,6 +745,8 @@ fn test_has_tool_calls_true() { #[test] fn test_has_tool_calls_false_no_assistant() { let req = AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![Message::User { content: MessageContent::Text("hi".into()), name: None, @@ -742,6 +776,8 @@ fn test_has_tool_calls_false_no_assistant() { #[test] fn test_has_tool_calls_false_empty_calls() { let req = AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![Message::Assistant { content: Some(MessageContent::Text("hello".into())), tool_calls: Some(vec![]), diff --git a/crates/core/tests/integration/native_plugin_tests.rs b/crates/core/tests/integration/native_plugin_tests.rs index b60855717..ef2604f58 100644 --- a/crates/core/tests/integration/native_plugin_tests.rs +++ b/crates/core/tests/integration/native_plugin_tests.rs @@ -844,6 +844,40 @@ fn native_activation_clear_deregisters_plugin_kind() { activation.clear(); } +#[tokio::test] +async fn native_request_intercept_rejects_manifest_that_admits_relay_0_5() { + let _guard = NATIVE_PLUGIN_TEST_LOCK.lock().await; + let fixture = build_fixture_plugin(); + let manifest_ref = write_manifest_text(ManifestOptions { + manifest_dir: fixture.manifest_dir.path(), + plugin_id: "fixture_native", + relay: ">=0.5,<1.0", + library: &fixture.library_path.to_string_lossy(), + symbol: "nemo_relay_fixture_native_plugin", + integrity: None, + }); + let activation = load_native_plugins([load_spec("fixture_native", &manifest_ref)]) + .expect("the host version is in the broad manifest range"); + + let mut plugin_config = PluginConfig::default(); + plugin_config.components.push(PluginComponentSpec { + kind: "fixture_native".into(), + enabled: true, + config: Map::new(), + }); + let error = initialize_plugins_exact(plugin_config) + .await + .expect_err("the request-intercept registration should reject Relay 0.5 compatibility") + .to_string(); + assert!( + error.contains("llm request intercept failed: InvalidArg"), + "{error}" + ); + + clear_plugin_configuration().expect("failed initialization state should clear"); + activation.clear(); +} + #[test] fn native_loader_resolves_manifest_directory_and_relative_library_paths() { let _guard = NATIVE_PLUGIN_TEST_LOCK.blocking_lock(); diff --git a/crates/core/tests/integration/pipeline_tests.rs b/crates/core/tests/integration/pipeline_tests.rs index b44e68056..918407663 100644 --- a/crates/core/tests/integration/pipeline_tests.rs +++ b/crates/core/tests/integration/pipeline_tests.rs @@ -29,12 +29,13 @@ use nemo_relay::api::runtime::{LlmExecutionNextFn, LlmJsonStream, LlmStreamExecu use nemo_relay::api::runtime::{create_scope_stack, set_thread_scope_stack}; use nemo_relay::api::scope::ScopeType; use nemo_relay::api::subscriber::{deregister_subscriber, flush_subscribers, register_subscriber}; +use nemo_relay::codec::anthropic::AnthropicMessagesCodec; use nemo_relay::codec::openai_chat::OpenAIChatCodec; use nemo_relay::codec::optimization::{ LlmOptimizationContribution, LlmOptimizationKind, LlmOptimizationModel, LlmOptimizationModelTransition, }; -use nemo_relay::codec::request::{AnnotatedLlmRequest, Message, MessageContent}; +use nemo_relay::codec::request::{AnnotatedLlmRequest, ContentPart, Message, MessageContent}; use nemo_relay::codec::response::FinishReason; use nemo_relay::codec::response::{ AnnotatedLlmResponse, CostEstimate, CostSource, PricingCatalog, PricingResolver, Usage, @@ -178,6 +179,8 @@ impl LlmCodec for TrackingCodec { } Ok(AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages, model: Some(self.id.clone()), params: None, @@ -455,6 +458,98 @@ async fn test_encode_runs_after_intercepts() { deregister_llm_request_intercept("modify_model").unwrap(); } +#[tokio::test] +async fn anthropic_issue_501_round_trips_and_applies_annotated_edits() { + let _lock = TEST_MUTEX.lock().unwrap(); + reset_global(); + setup_isolated_thread(); + + let original = make_llm_request(json!({ + "model": "claude-sonnet-4-20250514", + "max_tokens": 128, + "system": [{ + "type": "text", + "text": "Keep the cache marker.", + "cache_control": {"type": "ephemeral"} + }], + "messages": [ + {"role": "assistant", "content": [{"type": "tool_use", "id": "toolu_1", "name": "lookup", "input": {"q": "x"}}]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_1", "content": "done"}]} + ], + "future_field": null + })); + + let unchanged_capture = Arc::new(Mutex::new(None::)); + let captured = unchanged_capture.clone(); + let provider: LlmExecutionNextFn = Arc::new(move |request| { + *captured.lock().unwrap() = Some(request); + Box::pin(async { Ok(json!({"content": []})) }) + }); + llm_call_execute( + LlmCallExecuteParams::builder() + .name("anthropic.messages") + .request(original.clone()) + .func(provider) + .codec(Arc::new(AnthropicMessagesCodec)) + .build(), + ) + .await + .unwrap(); + assert_eq!(unchanged_capture.lock().unwrap().as_ref(), Some(&original)); + + register_llm_request_intercept( + "issue_501_annotated_edit", + 1, + false, + Arc::new(|_name, mut request, annotated| { + let mut annotated = annotated.expect("request codec must provide an annotation"); + let Some(MessageContent::Parts(parts)) = &mut annotated.instructions else { + panic!("expected block-form Anthropic system instructions"); + }; + let ContentPart::Text { text, extra } = &mut parts[0] else { + panic!("expected portable text instructions"); + }; + *text = "Edited without dropping cache metadata.".into(); + assert_eq!( + extra.get("cache_control"), + Some(&json!({"type": "ephemeral"})) + ); + request + .headers + .insert("x-annotation-seen".into(), json!("yes")); + Ok(LlmRequestInterceptOutcome::new(request, Some(annotated))) + }), + ) + .unwrap(); + + let edited_capture = Arc::new(Mutex::new(None::)); + let captured = edited_capture.clone(); + let provider: LlmExecutionNextFn = Arc::new(move |request| { + *captured.lock().unwrap() = Some(request); + Box::pin(async { Ok(json!({"content": []})) }) + }); + llm_call_execute( + LlmCallExecuteParams::builder() + .name("anthropic.messages") + .request(original.clone()) + .func(provider) + .codec(Arc::new(AnthropicMessagesCodec)) + .build(), + ) + .await + .unwrap(); + + let edited = edited_capture.lock().unwrap().clone().unwrap(); + let mut expected = original; + expected.content["system"][0]["text"] = json!("Edited without dropping cache metadata."); + expected + .headers + .insert("x-annotation-seen".into(), json!("yes")); + assert_eq!(edited, expected); + + deregister_llm_request_intercept("issue_501_annotated_edit").unwrap(); +} + #[tokio::test] async fn test_codec_rejects_raw_content_mutation_before_lifecycle() { let _lock = TEST_MUTEX.lock().unwrap(); diff --git a/crates/core/tests/integration/worker_plugin_tests.rs b/crates/core/tests/integration/worker_plugin_tests.rs index ea3f447c7..4387b048a 100644 --- a/crates/core/tests/integration/worker_plugin_tests.rs +++ b/crates/core/tests/integration/worker_plugin_tests.rs @@ -935,6 +935,28 @@ fn unsupported_worker_relay_requirement_reports_compatibility_error() { assert!(error.contains("requires relay"), "{error}"); } +#[test] +fn worker_request_intercept_rejects_manifest_that_admits_relay_0_5() { + let _guard = WORKER_PLUGIN_TEST_LOCK.blocking_lock(); + let fixture = build_fixture_worker(); + let (_manifest_dir, manifest_ref) = + write_manifest_with_relay(fixture.binary_path(), ">=0.5,<1.0"); + + let error = match load_worker_plugins([WorkerPluginLoadSpec { + plugin_id: "fixture_worker".into(), + manifest_ref: manifest_ref.to_string_lossy().into_owned(), + environment_ref: None, + config: Map::new(), + }]) { + Ok(activation) => { + activation.clear(); + panic!("the request-intercept registration should reject Relay 0.5 compatibility"); + } + Err(error) => error.to_string(), + }; + assert!(error.contains(">=0.6,<1.0"), "{error}"); +} + #[test] fn invalid_worker_relay_requirement_reports_parse_error() { let _guard = WORKER_PLUGIN_TEST_LOCK.blocking_lock(); @@ -1118,6 +1140,8 @@ struct FixtureCodec; impl LlmCodec for FixtureCodec { fn decode(&self, request: &LlmRequest) -> FlowResult { Ok(AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: Vec::new(), model: Some("fixture-model".into()), params: None, diff --git a/crates/core/tests/unit/codec/anthropic_tests.rs b/crates/core/tests/unit/codec/anthropic_tests.rs index 9fe7d8724..5a47b234c 100644 --- a/crates/core/tests/unit/codec/anthropic_tests.rs +++ b/crates/core/tests/unit/codec/anthropic_tests.rs @@ -6,7 +6,10 @@ use super::*; use serde_json::json; -use super::super::request::{Message, MessageContent, ToolChoice}; +use super::super::request::{ + ContentPart, FunctionDefinition, Message, MessageContent, ProviderNativeComponent, ToolChoice, + ToolChoiceFunction, ToolChoiceFunctionName, +}; use super::super::response::{ApiSpecificResponse, FinishReason}; // ------------------------------------------------------------------- @@ -405,12 +408,12 @@ fn test_decode_request_full() { })); let annotated = codec.decode(&request).unwrap(); - // System should be prepended as Message::System - assert_eq!(annotated.messages.len(), 2); - assert!( - matches!(&annotated.messages[0], Message::System { content: MessageContent::Text(t), .. } if t == "Be helpful") + assert_eq!(annotated.messages.len(), 1); + assert_eq!( + annotated.instructions, + Some(MessageContent::Text("Be helpful".into())) ); - assert!(matches!(&annotated.messages[1], Message::User { .. })); + assert!(matches!(&annotated.messages[0], Message::User { .. })); assert_eq!(annotated.model, Some("claude-sonnet-4-20250514".into())); @@ -420,10 +423,12 @@ fn test_decode_request_full() { // Tools should be normalized: input_schema -> parameters let tools = annotated.tools.unwrap(); assert_eq!(tools.len(), 1); - assert_eq!(tools[0].tool_type, "function"); - assert_eq!(tools[0].function.name, "get_weather"); - assert_eq!(tools[0].function.description, Some("Get weather".into())); - assert!(tools[0].function.parameters.is_some()); + let ToolDefinition::Function { function, .. } = &tools[0] else { + panic!("expected a portable function tool"); + }; + assert_eq!(function.name, "get_weather"); + assert_eq!(function.description, Some("Get weather".into())); + assert!(function.parameters.is_some()); assert_eq!(annotated.tool_choice, Some(ToolChoice::Auto)); } @@ -443,11 +448,10 @@ fn test_decode_request_system_array() { "max_tokens": 100 })); let annotated = codec.decode(&request).unwrap(); - assert_eq!(annotated.messages.len(), 2); + assert_eq!(annotated.messages.len(), 1); assert!(matches!( - &annotated.messages[0], - Message::System { content: MessageContent::Text(t), .. } - if t == "First instruction.\nSecond instruction." + &annotated.instructions, + Some(MessageContent::Parts(parts)) if parts.len() == 2 )); } @@ -521,7 +525,7 @@ fn test_decode_request_extra_fields() { })); let annotated = codec.decode(&request).unwrap(); assert_eq!(annotated.metadata, Some(json!({"user_id": "abc"}))); - assert_eq!(annotated.extra.get("stream"), Some(&json!(true))); + assert_eq!(annotated.stream, Some(true)); } #[test] @@ -601,15 +605,19 @@ fn test_decode_request_litellm_bridge_thinking_output_config_preserved_in_extra( // stable extraction assert_eq!(annotated.tool_choice, Some(ToolChoice::Required)); assert_eq!(annotated.parallel_tool_calls, Some(true)); - // bridge-specific controls preserved losslessly - assert_eq!( - annotated.extra.get("thinking"), - Some(&json!({"type":"enabled","budget_tokens":2048})) - ); + let Some(ApiSpecificRequest::AnthropicMessages { + thinking, + output_config, + .. + }) = annotated.api_specific + else { + panic!("expected Anthropic-specific request controls"); + }; assert_eq!( - annotated.extra.get("output_config"), - Some(&json!({"effort":"low"})) + thinking, + Some(json!({"type":"enabled","budget_tokens":2048})) ); + assert_eq!(output_config, Some(json!({"effort":"low"}))); assert_eq!( annotated.extra.get("reasoning_effort"), Some(&json!("minimal")) @@ -643,6 +651,190 @@ fn test_decode_request_litellm_cache_control_blocks_preserved() { assert_eq!(annotated.system_prompt(), Some("Be terse")); // `system` is a modeled key in Anthropic decode and should not live in extra. assert!(annotated.extra.get("system").is_none()); + let Some(MessageContent::Parts(parts)) = &annotated.instructions else { + panic!("expected block-form system instructions"); + }; + let super::super::request::ContentPart::Text { extra, .. } = &parts[0] else { + panic!("expected a portable text block"); + }; + assert_eq!( + extra.get("cache_control"), + Some(&json!({"type": "ephemeral"})) + ); + assert_eq!(codec.encode(&annotated, &request).unwrap(), request); +} + +#[test] +fn request_schema_fixture_round_trips_and_issue_501_edit_is_surgical() { + let codec = AnthropicMessagesCodec; + let original = make_request(json!({ + "model": "claude-sonnet-4-20250514", + "max_tokens": 512, + "system": [{ + "type": "text", + "text": "You are precise.", + "cache_control": {"type": "ephemeral"} + }], + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Inspect these", "cache_control": {"type": "ephemeral"}}, + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "abcd"}}, + {"type": "document", "source": {"type": "text", "media_type": "text/plain", "data": "notes"}, "title": "Notes"}, + {"type": "tool_result", "tool_use_id": "toolu_1", "content": [{"type": "text", "text": "sunny"}], "is_error": false} + ] + }, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Calling a tool"}, + {"type": "tool_use", "id": "toolu_1", "name": "weather", "input": {"city": "NYC"}}, + {"type": "server_tool_use", "id": "srv_1", "name": "web_search", "input": {"query": "weather"}}, + {"type": "thinking", "thinking": "private", "signature": "sig"} + ] + } + ], + "cache_control": {"type": "ephemeral"}, + "container": "container_1", + "inference_geo": "us", + "metadata": {"user_id": "user_1"}, + "output_config": {"effort": "high"}, + "service_tier": "auto", + "stop_sequences": ["END"], + "stream": true, + "temperature": 0.2, + "thinking": {"type": "enabled", "budget_tokens": 1024}, + "tool_choice": {"type": "auto", "disable_parallel_tool_use": false}, + "tools": [ + {"name": "weather", "description": "Weather", "input_schema": {"type": "object"}, "cache_control": {"type": "ephemeral"}}, + {"type": "web_search_20250305", "name": "web_search", "max_uses": 2} + ], + "top_k": 20, + "top_p": 0.9, + "anthropic-user-profile-id": "profile_1", + "future_field": null + })); + + let mut annotated = codec.decode(&original).unwrap(); + assert_eq!(codec.encode(&annotated, &original).unwrap(), original); + + let Some(MessageContent::Parts(parts)) = &mut annotated.instructions else { + panic!("expected block-form instructions"); + }; + let super::super::request::ContentPart::Text { text, extra } = &mut parts[0] else { + panic!("expected portable instruction text"); + }; + *text = "You are concise.".into(); + assert_eq!( + extra.get("cache_control"), + Some(&json!({"type": "ephemeral"})) + ); + + let encoded = codec.encode(&annotated, &original).unwrap(); + let mut expected = original.clone(); + expected.content["system"][0]["text"] = json!("You are concise."); + assert_eq!(encoded, expected); +} + +#[test] +fn anthropic_tool_edits_preserve_nested_unknown_fields_and_explicit_nulls() { + let codec = AnthropicMessagesCodec; + let original = make_request(json!({ + "model": "claude-sonnet-4-20250514", + "max_tokens": 64, + "messages": [{"role": "user", "content": "hello"}], + "tools": [{ + "name": "lookup_before", + "description": null, + "input_schema": {"type": "object"}, + "strict": null, + "future_tool": {"mode": "keep"} + }], + "tool_choice": { + "type": "auto", + "disable_parallel_tool_use": false, + "future_choice": null + } + })); + let mut annotated = codec.decode(&original).unwrap(); + let ToolDefinition::Function { function, .. } = &mut annotated.tools.as_mut().unwrap()[0] + else { + panic!("expected portable function tool"); + }; + function.name = "lookup_after".into(); + annotated.parallel_tool_calls = Some(false); + + let encoded = codec.encode(&annotated, &original).unwrap(); + let mut expected = original; + expected.content["tools"][0]["name"] = json!("lookup_after"); + expected.content["tool_choice"]["disable_parallel_tool_use"] = json!(true); + assert_eq!(encoded, expected); +} + +#[test] +fn request_schema_rejects_malformed_known_anthropic_components() { + let codec = AnthropicMessagesCodec; + for content in [ + json!([{"type": "tool_use", "id": "toolu_1", "input": {}}]), + json!([{"type": "tool_result", "tool_use_id": "toolu_1"}]), + ] { + let error = codec + .decode(&make_request(json!({ + "model": "claude-sonnet-4-20250514", + "max_tokens": 64, + "messages": [{"role": "user", "content": content}] + }))) + .unwrap_err(); + assert!(error.to_string().contains("missing")); + } + + let error = codec + .decode(&make_request(json!({ + "model": "claude-sonnet-4-20250514", + "max_tokens": 64, + "messages": [], + "stop_sequences": ["ok", 7] + }))) + .unwrap_err(); + assert!(error.to_string().contains("stop_sequences")); + + for (field, malformed) in [ + ("top_k", json!("many")), + ("temperature", json!("high")), + ("stream", json!("yes")), + ("container", json!(7)), + ("metadata", json!("user-123")), + ("cache_control", json!("ephemeral")), + ("output_config", json!("high")), + ("thinking", json!("enabled")), + ] { + let mut content = json!({ + "model": "claude-sonnet-4-20250514", + "max_tokens": 64, + "messages": [] + }); + content[field] = malformed; + let error = codec.decode(&make_request(content)).unwrap_err(); + assert!( + error.to_string().contains(field), + "unexpected error: {error}" + ); + } + + let error = codec + .decode(&make_request(json!({ + "model": "claude-sonnet-4-20250514", + "max_tokens": 64, + "messages": [], + "tools": [{ + "name": "lookup", + "input_schema": {"type": "object"}, + "strict": "yes" + }] + }))) + .unwrap_err(); + assert!(error.to_string().contains("strict")); } // =================================================================== @@ -703,6 +895,24 @@ fn test_encode_writes_anthropic_modeled_controls() { ); } +#[test] +fn test_encode_synthesizes_tool_choice_for_parallel_tool_calls_edit() { + let codec = AnthropicMessagesCodec; + let original = make_request(json!({ + "messages": [{ "role": "user", "content": "Hi" }], + "model": "claude-sonnet-4-20250514", + "max_tokens": 100 + })); + let mut annotated = codec.decode(&original).unwrap(); + annotated.parallel_tool_calls = Some(false); + + let encoded = codec.encode(&annotated, &original).unwrap(); + assert_eq!( + encoded.content["tool_choice"], + json!({"type": "auto", "disable_parallel_tool_use": true}) + ); +} + #[test] fn test_encode_system_as_top_level() { let codec = AnthropicMessagesCodec; @@ -804,33 +1014,22 @@ fn test_helper_and_error_paths_cover_remaining_anthropic_branches() { None ); assert_eq!(decode_anthropic_tool_choice(&json!({"type": "tool"})), None); - assert_eq!( - extract_system_message(&json!([{ "type": "image", "source": "ignored" }])), - None - ); + assert!(matches!( + decode_anthropic_content(&json!([{ "type": "image", "source": "ignored" }])).unwrap(), + MessageContent::Parts(parts) + if matches!(parts.as_slice(), [super::super::request::ContentPart::Image { .. }]) + )); - let system_parts = Message::System { - content: MessageContent::Parts(vec![ - super::super::request::ContentPart::Text { - text: "First".into(), - }, - super::super::request::ContentPart::Text { - text: "Second".into(), - }, - ]), - name: None, - }; - assert_eq!( - extract_system_text(&system_parts), - Some("First\nSecond".to_string()) - ); - assert_eq!( - extract_system_text(&Message::User { - content: MessageContent::Text("hi".into()), - name: None, - }), - None - ); + let system_parts = MessageContent::Parts(vec![ + super::super::request::ContentPart::Text { + text: "First".into(), + extra: serde_json::Map::new(), + }, + super::super::request::ContentPart::Text { + text: "Second".into(), + extra: serde_json::Map::new(), + }, + ]); let codec = AnthropicMessagesCodec; @@ -853,7 +1052,9 @@ fn test_helper_and_error_paths_cover_remaining_anthropic_branches() { assert_eq!(partial_usage.usage.unwrap().total_tokens, None); let annotated = AnnotatedLlmRequest { - messages: vec![system_parts], + instructions: Some(system_parts), + api_specific: None, + messages: vec![], model: Some("claude-sonnet-4-20250514".into()), params: Some(GenerationParams { temperature: Some(0.3), @@ -861,13 +1062,15 @@ fn test_helper_and_error_paths_cover_remaining_anthropic_branches() { top_p: Some(0.8), stop: Some(vec!["END".into()]), }), - tools: Some(vec![ToolDefinition { - tool_type: "function".into(), + tools: Some(vec![ToolDefinition::Function { function: FunctionDefinition { name: "lookup".into(), description: None, parameters: None, + strict: None, + extra: serde_json::Map::new(), }, + extra: serde_json::Map::new(), }]), tool_choice: Some(ToolChoice::None), store: None, @@ -900,7 +1103,13 @@ fn test_helper_and_error_paths_cover_remaining_anthropic_branches() { assert_eq!(obj.get("top_p"), Some(&json!(0.8))); assert_eq!(obj.get("stop_sequences"), Some(&json!(["END"]))); assert_eq!(obj.get("tool_choice"), Some(&json!({"type": "none"}))); - assert_eq!(obj.get("system"), Some(&json!("First\nSecond"))); + assert_eq!( + obj.get("system"), + Some(&json!([ + {"type": "text", "text": "First"}, + {"type": "text", "text": "Second"} + ])) + ); let tools = obj.get("tools").unwrap().as_array().unwrap(); assert_eq!(tools[0].get("name"), Some(&json!("lookup"))); @@ -909,12 +1118,200 @@ fn test_helper_and_error_paths_cover_remaining_anthropic_branches() { match codec.encode(&annotated, &make_request(json!("still-not-an-object"))) { Err(FlowError::Internal(message)) => { - assert!(message.contains("original content is not an object")); + assert!(message.contains("not an object")); } other => panic!("unexpected encode result: {other:?}"), } } +#[test] +fn anthropic_request_component_branch_matrix() { + let specific = ToolChoice::Specific(ToolChoiceFunction { + choice_type: "function".into(), + function: ToolChoiceFunctionName { + name: "lookup".into(), + }, + }); + assert_eq!( + encode_anthropic_tool_choice(&ToolChoice::Required).unwrap(), + json!({"type": "any"}) + ); + assert_eq!( + encode_anthropic_tool_choice(&specific).unwrap(), + json!({"type": "tool", "name": "lookup"}) + ); + assert_eq!( + encode_tool_choice_with_parallel_hint(&specific, Some(false)).unwrap(), + json!({"type": "tool", "name": "lookup", "disable_parallel_tool_use": true}) + ); + let native_choice = ToolChoice::ProviderNative(ProviderNativeComponent { + provider: "anthropic_messages".into(), + kind: "future".into(), + value: json!({"type": "future"}), + }); + assert_eq!( + encode_anthropic_tool_choice(&native_choice).unwrap(), + json!({"type": "future"}) + ); + let mismatched_choice = ToolChoice::ProviderNative(ProviderNativeComponent { + provider: "openai_chat".into(), + kind: "future".into(), + value: json!({"type": "future"}), + }); + assert!(encode_anthropic_tool_choice(&mismatched_choice).is_err()); + + for invalid in [json!(42), json!({"type": "text"})] { + assert!(decode_anthropic_content_part(&invalid).is_err()); + } + for invalid in [ + json!({"type": "tool_use", "name": "lookup", "input": {}}), + json!({"type": "tool_use", "id": "call", "input": {}}), + json!({"type": "tool_use", "id": "call", "name": "lookup"}), + json!({"type": "tool_result", "content": "ok"}), + json!({"type": "tool_result", "tool_use_id": "call"}), + json!({"type": "tool_result", "tool_use_id": "call", "content": "ok", "is_error": "no"}), + ] { + assert!(decode_anthropic_content_part(&invalid).is_err()); + } + assert!(decode_anthropic_content(&json!({})).is_err()); + + let content = MessageContent::Parts(vec![ + ContentPart::Text { + text: "hello".into(), + extra: serde_json::Map::from_iter([( + "cache_control".into(), + json!({"type": "ephemeral"}), + )]), + }, + ContentPart::Image { + image: json!({"source": {"type": "base64", "data": "a"}}), + extra: serde_json::Map::from_iter([("alt".into(), json!("image"))]), + }, + ContentPart::File { + file: json!({"source": {"type": "text", "data": "notes"}}), + extra: serde_json::Map::from_iter([("title".into(), json!("Notes"))]), + }, + ContentPart::ToolUse { + id: "call".into(), + name: "lookup".into(), + input: json!({"q": "x"}), + extra: serde_json::Map::from_iter([("caller".into(), json!("client"))]), + }, + ContentPart::ToolResult { + tool_use_id: "call".into(), + content: json!("ok"), + is_error: Some(false), + extra: serde_json::Map::from_iter([("future".into(), json!(1))]), + }, + ContentPart::ProviderNative { + provider: "anthropic_messages".into(), + kind: "thinking".into(), + value: json!({"type": "thinking", "thinking": "private"}), + }, + ]); + let encoded_content = encode_anthropic_content(&content).unwrap(); + assert_eq!(encoded_content.as_array().unwrap().len(), 6); + assert_eq!(encoded_content[4]["is_error"], json!(false)); + assert!( + encode_anthropic_content_part(&ContentPart::Image { + image: json!("bad"), + extra: serde_json::Map::new(), + }) + .is_err() + ); + assert!( + encode_anthropic_content_part(&ContentPart::File { + file: json!("bad"), + extra: serde_json::Map::new(), + }) + .is_err() + ); + assert!( + encode_anthropic_content_part(&ContentPart::Audio { + audio: json!({}), + extra: serde_json::Map::new(), + }) + .is_err() + ); + + for invalid in [json!(42), json!({"content": "x"}), json!({"role": "user"})] { + assert!(decode_anthropic_message(&invalid).is_err()); + } + assert!(matches!( + decode_anthropic_message(&json!({"role": "user", "content": "x", "name": "Ada"})).unwrap(), + Message::ProviderNative { .. } + )); + assert!(matches!( + decode_anthropic_message(&json!({"role": "system", "content": "x"})).unwrap(), + Message::System { .. } + )); + assert!(matches!( + decode_anthropic_message(&json!({"role": "future", "content": "x"})).unwrap(), + Message::ProviderNative { .. } + )); + + let assistant_without_content = Message::Assistant { + content: None, + tool_calls: None, + name: None, + }; + assert_eq!( + encode_anthropic_message(&assistant_without_content).unwrap()["content"], + json!([]) + ); + let native_message = Message::ProviderNative { + provider: "anthropic_messages".into(), + kind: "future".into(), + value: json!({"role": "future", "content": null}), + }; + assert_eq!( + encode_anthropic_message(&native_message).unwrap()["role"], + json!("future") + ); + assert!( + encode_anthropic_message(&Message::Developer { + content: MessageContent::Text("no".into()), + name: None, + }) + .is_err() + ); + + assert!(decode_anthropic_tool(&json!(42)).is_err()); + assert!(matches!( + decode_anthropic_tool(&json!({"type": "web_search_20250305", "name": "search"})).unwrap(), + ToolDefinition::ProviderNative { .. } + )); + let rich_tool = ToolDefinition::Function { + function: FunctionDefinition { + name: "lookup".into(), + description: Some("Lookup".into()), + parameters: Some(json!({"type": "object"})), + strict: Some(true), + extra: serde_json::Map::from_iter([("future_function".into(), json!(1))]), + }, + extra: serde_json::Map::from_iter([("cache_control".into(), json!({"type": "ephemeral"}))]), + }; + assert_eq!( + encode_anthropic_tool(&rich_tool).unwrap()["strict"], + json!(true) + ); + let native_tool = ToolDefinition::ProviderNative { + provider: "anthropic_messages".into(), + kind: "web_search".into(), + value: json!({"type": "web_search", "name": "search"}), + }; + assert_eq!( + encode_anthropic_tool(&native_tool).unwrap()["type"], + json!("web_search") + ); + let mismatched_tool = ToolDefinition::ProviderNative { + provider: "openai_chat".into(), + kind: "custom".into(), + value: json!({"type": "custom"}), + }; + assert!(encode_anthropic_tool(&mismatched_tool).is_err()); +} + #[test] fn test_encode_tool_choice_any_to_anthropic() { let codec = AnthropicMessagesCodec; diff --git a/crates/core/tests/unit/codec/openai_chat_tests.rs b/crates/core/tests/unit/codec/openai_chat_tests.rs index ce66859ad..872373d8f 100644 --- a/crates/core/tests/unit/codec/openai_chat_tests.rs +++ b/crates/core/tests/unit/codec/openai_chat_tests.rs @@ -6,7 +6,10 @@ use super::*; use serde_json::json; -use super::super::request::{ContentPart, MessageContent, OpenAiImageUrl}; +use super::super::request::{ + ContentPart, FunctionCall, FunctionDefinition, Message, MessageContent, OpenAiImageUrl, + ProviderNativeComponent, ToolCall, ToolChoiceFunction, ToolChoiceFunctionName, +}; use super::super::response::{ApiSpecificResponse, CostSource, FinishReason}; // ------------------------------------------------------------------- @@ -566,7 +569,10 @@ fn test_decode_request_full() { let tools = annotated.tools.unwrap(); assert_eq!(tools.len(), 1); - assert_eq!(tools[0].function.name, "get_weather"); + let ToolDefinition::Function { function, .. } = &tools[0] else { + panic!("expected a portable function tool"); + }; + assert_eq!(function.name, "get_weather"); assert_eq!(annotated.tool_choice, Some(ToolChoice::Auto)); } @@ -596,11 +602,16 @@ fn test_decode_request_extra_fields() { })); let annotated = codec.decode(&request).unwrap(); assert_eq!(annotated.stream, Some(true)); - assert_eq!(annotated.extra.get("seed"), Some(&json!(42))); - assert_eq!( - annotated.extra.get("response_format"), - Some(&json!({"type": "json_object"})) - ); + let Some(ApiSpecificRequest::OpenAIChat { + seed, + response_format, + .. + }) = annotated.api_specific + else { + panic!("expected Chat-specific request controls"); + }; + assert_eq!(seed, Some(42)); + assert_eq!(response_format, Some(json!({"type": "json_object"}))); } #[test] @@ -633,8 +644,8 @@ fn test_decode_request_no_messages_key() { let request = make_request(json!({ "model": "gpt-4o" })); - let annotated = codec.decode(&request).unwrap(); - assert!(annotated.messages.is_empty()); + let error = codec.decode(&request).unwrap_err(); + assert!(error.to_string().contains("missing messages")); } #[test] @@ -658,13 +669,15 @@ fn test_decode_request_multimodal_image_url_parts() { parts, &vec![ ContentPart::Text { - text: "describe this".into() + text: "describe this".into(), + extra: serde_json::Map::new(), }, ContentPart::ImageUrl { image_url: OpenAiImageUrl { url: "https://example.com/cat.png".into(), detail: Some("high".into()) - } + }, + extra: serde_json::Map::new(), } ] ); @@ -779,12 +792,14 @@ fn test_encode_request_multimodal_image_url_parts() { content: MessageContent::Parts(vec![ ContentPart::Text { text: "describe this".into(), + extra: serde_json::Map::new(), }, ContentPart::ImageUrl { image_url: OpenAiImageUrl { url: "https://example.com/cat.png".into(), detail: Some("low".into()), }, + extra: serde_json::Map::new(), }, ]), name: None, @@ -843,24 +858,26 @@ fn test_helper_and_error_paths_cover_remaining_chat_branches() { }))) .unwrap_err() { - FlowError::Internal(message) => assert!(message.contains("OpenAI Chat tools decode")), + FlowError::InvalidArgument(message) => { + assert!(message.contains("tools must be an array")) + } other => panic!("unexpected tools decode error: {other}"), } - match codec + let native_choice = codec .decode(&make_request(json!({ "messages": [], "tool_choice": [] }))) - .unwrap_err() - { - FlowError::Internal(message) => { - assert!(message.contains("OpenAI Chat tool_choice decode")); - } - other => panic!("unexpected tool_choice decode error: {other}"), - } + .unwrap(); + assert!(matches!( + native_choice.tool_choice, + Some(ToolChoice::ProviderNative(_)) + )); let annotated = AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![], model: Some("gpt-4.1-mini".into()), params: Some(GenerationParams { @@ -869,13 +886,15 @@ fn test_helper_and_error_paths_cover_remaining_chat_branches() { top_p: Some(0.9), stop: Some(vec!["END".into()]), }), - tools: Some(vec![ToolDefinition { - tool_type: "function".into(), + tools: Some(vec![ToolDefinition::Function { function: super::super::request::FunctionDefinition { name: "lookup".into(), description: Some("Look up data".into()), parameters: Some(json!({"type": "object"})), + strict: None, + extra: serde_json::Map::new(), }, + extra: serde_json::Map::new(), }]), tool_choice: Some(ToolChoice::Required), store: None, @@ -906,65 +925,283 @@ fn test_helper_and_error_paths_cover_remaining_chat_branches() { assert_eq!(obj.get("temperature"), Some(&json!(0.2))); assert_eq!(obj.get("top_p"), Some(&json!(0.9))); assert_eq!(obj.get("stop"), Some(&json!(["END"]))); - assert_eq!(obj.get("max_tokens"), Some(&json!(64))); + assert_eq!(obj.get("max_completion_tokens"), Some(&json!(64))); assert!(obj.get("tools").unwrap().is_array()); assert_eq!(obj.get("tool_choice"), Some(&json!("required"))); match codec.encode(&annotated, &make_request(json!("still-not-an-object"))) { Err(FlowError::Internal(message)) => { - assert!(message.contains("original content is not an object")); + assert!(message.contains("not an object")); } other => panic!("unexpected encode result: {other:?}"), } } -// =================================================================== -// stream_options injection tests (for Phoenix / OpenInference usage) -// =================================================================== -/// On streaming requests (`stream: true`), the encoder forces -/// `stream_options.include_usage=true` when the caller did not provide -/// `stream_options`. Without this, OpenAI-compatible providers never emit -/// the terminal usage chunk and token counts are lost in Phoenix traces. #[test] -fn test_encode_injects_stream_options_on_streaming_request() { - let codec = OpenAIChatCodec; - let annotated = AnnotatedLlmRequest { - messages: vec![], - model: Some("gpt-4o".into()), - params: None, - tools: None, - tool_choice: None, - store: None, - previous_response_id: None, - truncation: None, - reasoning: None, - include: None, - user: None, - metadata: None, - service_tier: None, - parallel_tool_calls: None, - max_output_tokens: None, - max_tool_calls: None, - top_logprobs: None, - stream: None, - extra: serde_json::Map::new(), +fn chat_request_component_branch_matrix() { + assert!(decode_chat_content(&json!({})).is_err()); + for invalid in [ + json!(42), + json!({"type": "text"}), + json!({"type": "image_url"}), + json!({"type": "image_url", "image_url": {"detail": "high"}}), + json!({"type": "input_audio"}), + json!({"type": "file"}), + json!({"type": "refusal"}), + ] { + assert!(decode_chat_content_part(&invalid).is_err()); + } + let parts = MessageContent::Parts(vec![ + ContentPart::Text { + text: "hello".into(), + extra: serde_json::Map::from_iter([("future".into(), json!(1))]), + }, + ContentPart::ImageUrl { + image_url: OpenAiImageUrl { + url: "https://example.com/image.png".into(), + detail: Some("high".into()), + }, + extra: serde_json::Map::new(), + }, + ContentPart::Audio { + audio: json!({"data": "audio", "format": "wav"}), + extra: serde_json::Map::new(), + }, + ContentPart::File { + file: json!({"file_id": "file_1"}), + extra: serde_json::Map::new(), + }, + ContentPart::Refusal { + refusal: "no".into(), + extra: serde_json::Map::new(), + }, + ContentPart::ProviderNative { + provider: "openai_chat".into(), + kind: "future".into(), + value: json!({"type": "future", "value": 1}), + }, + ]); + assert_eq!( + encode_chat_content(&parts) + .unwrap() + .as_array() + .unwrap() + .len(), + 6 + ); + assert!( + encode_chat_content(&MessageContent::Parts(vec![ContentPart::Image { + image: json!({}), + extra: serde_json::Map::new(), + }])) + .is_err() + ); + + for invalid in [ + json!(42), + json!({"type": "function"}), + json!({"id": "call", "type": "function", "function": {"name": "lookup", "arguments": "{}"}, "future": 1}), + json!({"id": "call", "type": "function", "function": {"name": "lookup", "arguments": "{}", "future": 1}}), + json!({"type": "function", "function": {"name": "lookup", "arguments": "{}"}}), + json!({"id": "call", "type": "function", "function": {"arguments": "{}"}}), + json!({"id": "call", "type": "function", "function": {"name": "lookup"}}), + ] { + let result = decode_chat_tool_call(&invalid); + assert!(result.is_err() || result.unwrap().is_none()); + } + assert!( + decode_chat_tool_call(&json!({"type": "custom"})) + .unwrap() + .is_none() + ); + assert!( + decode_chat_tool_call(&json!({ + "id": "call", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"} + })) + .unwrap() + .is_some() + ); + + for invalid in [ + json!(42), + json!({"content": "x"}), + json!({"role": "user"}), + json!({"role": "user", "content": "x", "name": 7}), + json!({"role": "assistant", "tool_calls": 7}), + json!({"role": "tool", "content": "x"}), + json!({"role": "function", "content": null}), + ] { + assert!(decode_chat_message(&invalid).is_err()); + } + for native in [ + json!({"role": "user", "content": "x", "future": 1}), + json!({"role": "assistant", "content": null, "future": 1}), + json!({"role": "assistant", "content": null, "tool_calls": [{"type": "custom"}]}), + json!({"role": "future", "content": "x"}), + ] { + assert!(matches!( + decode_chat_message(&native).unwrap(), + Message::ProviderNative { .. } + )); + } + for portable in [ + json!({"role": "system", "content": "s", "name": null}), + json!({"role": "developer", "content": "d", "name": "dev"}), + json!({"role": "user", "content": "u"}), + json!({"role": "assistant", "content": null, "tool_calls": null, "name": "bot"}), + json!({"role": "tool", "content": "ok", "tool_call_id": "call"}), + json!({"role": "function", "content": null, "name": "legacy"}), + ] { + assert!(!matches!( + decode_chat_message(&portable).unwrap(), + Message::ProviderNative { .. } + )); + } + + let tool_call = ToolCall { + id: "call".into(), + call_type: "function".into(), + function: FunctionCall { + name: "lookup".into(), + arguments: "{}".into(), + }, }; - let encoded = codec - .encode( - &annotated, - &make_request(json!({ - "messages": [], - "model": "gpt-4o", - "stream": true, - })), - ) - .unwrap(); - let obj = encoded.content.as_object().unwrap(); + let messages = vec![ + Message::System { + content: MessageContent::Text("s".into()), + name: Some("system".into()), + }, + Message::Developer { + content: MessageContent::Text("d".into()), + name: None, + }, + Message::User { + content: MessageContent::Text("u".into()), + name: None, + }, + Message::Assistant { + content: Some(MessageContent::Text("a".into())), + tool_calls: Some(vec![tool_call]), + name: Some("bot".into()), + }, + Message::Tool { + content: MessageContent::Text("ok".into()), + tool_call_id: "call".into(), + }, + Message::Function { + content: None, + name: "legacy".into(), + }, + Message::ProviderNative { + provider: "openai_chat".into(), + kind: "future".into(), + value: json!({"role": "future"}), + }, + ]; + for message in &messages { + assert!(encode_chat_message(message).is_ok()); + } + assert!( + encode_chat_message(&Message::ToolCallItem { + id: None, + call_id: "call".into(), + name: "lookup".into(), + arguments: json!({}), + extra: serde_json::Map::new(), + }) + .is_err() + ); + + for invalid in [ + json!(42), + json!({"type": "function"}), + json!({"type": "function", "function": {}}), + ] { + assert!(decode_chat_tool(&invalid).is_err()); + } + assert!(matches!( + decode_chat_tool(&json!({"type": "custom", "name": "grammar"})).unwrap(), + ToolDefinition::ProviderNative { .. } + )); + let function_tool = ToolDefinition::Function { + function: FunctionDefinition { + name: "lookup".into(), + description: Some("Lookup".into()), + parameters: Some(json!({"type": "object"})), + strict: Some(true), + extra: serde_json::Map::from_iter([("future_function".into(), json!(1))]), + }, + extra: serde_json::Map::from_iter([("future_wrapper".into(), json!(2))]), + }; + assert_eq!( + encode_chat_tool(&function_tool).unwrap()["function"]["strict"], + json!(true) + ); + let native_tool = ToolDefinition::ProviderNative { + provider: "openai_chat".into(), + kind: "custom".into(), + value: json!({"type": "custom"}), + }; + assert_eq!( + encode_chat_tool(&native_tool).unwrap()["type"], + json!("custom") + ); + let mismatched_tool = ToolDefinition::ProviderNative { + provider: "openai_responses".into(), + kind: "custom".into(), + value: json!({"type": "custom"}), + }; + assert!(encode_chat_tool(&mismatched_tool).is_err()); + + let specific = ToolChoice::Specific(ToolChoiceFunction { + choice_type: "function".into(), + function: ToolChoiceFunctionName { + name: "lookup".into(), + }, + }); + assert_eq!( + encode_chat_tool_choice(&ToolChoice::None).unwrap(), + json!("none") + ); assert_eq!( - obj.get("stream_options"), - Some(&json!({"include_usage": true})), - "encoder must inject stream_options.include_usage on streaming requests when absent", + encode_chat_tool_choice(&specific).unwrap()["type"], + json!("function") ); + let native_choice = ToolChoice::ProviderNative(ProviderNativeComponent { + provider: "openai_chat".into(), + kind: "tool_choice".into(), + value: json!({"type": "custom"}), + }); + assert_eq!( + encode_chat_tool_choice(&native_choice).unwrap()["type"], + json!("custom") + ); + let mismatched_choice = ToolChoice::ProviderNative(ProviderNativeComponent { + provider: "anthropic_messages".into(), + kind: "tool_choice".into(), + value: json!({"type": "tool"}), + }); + assert!(encode_chat_tool_choice(&mismatched_choice).is_err()); +} +// =================================================================== +// stream_options identity tests +// =================================================================== + +/// Request codecs do not enrich streaming requests as a side effect. +#[test] +fn test_encode_does_not_inject_stream_options_on_streaming_request() { + let codec = OpenAIChatCodec; + let original = make_request(json!({ + "messages": [], + "model": "gpt-4o", + "stream": true + })); + let annotated = codec.decode(&original).unwrap(); + let encoded = codec.encode(&annotated, &original).unwrap(); + assert_eq!(encoded, original); + assert!(encoded.content.get("stream_options").is_none()); } /// Caller-supplied `stream_options` must be preserved verbatim. This matters @@ -973,40 +1210,15 @@ fn test_encode_injects_stream_options_on_streaming_request() { #[test] fn test_encode_preserves_caller_stream_options() { let codec = OpenAIChatCodec; - let annotated = AnnotatedLlmRequest { - messages: vec![], - model: Some("gpt-4o".into()), - params: None, - tools: None, - tool_choice: None, - store: None, - previous_response_id: None, - truncation: None, - reasoning: None, - include: None, - user: None, - metadata: None, - service_tier: None, - parallel_tool_calls: None, - max_output_tokens: None, - max_tool_calls: None, - top_logprobs: None, - stream: None, - extra: serde_json::Map::new(), - }; - let caller_set = json!({ + let original = make_request(json!({ "messages": [], "model": "gpt-4o", "stream": true, "stream_options": { "include_usage": false } - }); - let encoded = codec.encode(&annotated, &make_request(caller_set)).unwrap(); - let obj = encoded.content.as_object().unwrap(); - assert_eq!( - obj.get("stream_options"), - Some(&json!({"include_usage": false})), - "caller-provided stream_options must be preserved verbatim", - ); + })); + let annotated = codec.decode(&original).unwrap(); + let encoded = codec.encode(&annotated, &original).unwrap(); + assert_eq!(encoded, original); } /// Per the OpenAI Chat Completions spec, `stream_options` is only valid on @@ -1016,57 +1228,310 @@ fn test_encode_preserves_caller_stream_options() { #[test] fn test_encode_does_not_inject_stream_options_on_non_streaming() { let codec = OpenAIChatCodec; - let annotated = AnnotatedLlmRequest { - messages: vec![], - model: Some("gpt-4o".into()), - params: None, - tools: None, - tool_choice: None, - store: None, - previous_response_id: None, - truncation: None, - reasoning: None, - include: None, - user: None, - metadata: None, - service_tier: None, - parallel_tool_calls: None, - max_output_tokens: None, - max_tool_calls: None, - top_logprobs: None, - stream: None, - extra: serde_json::Map::new(), + for content in [ + json!({"messages": [], "model": "gpt-4o"}), + json!({"messages": [], "model": "gpt-4o", "stream": false}), + ] { + let original = make_request(content); + let annotated = codec.decode(&original).unwrap(); + let encoded = codec.encode(&annotated, &original).unwrap(); + assert_eq!(encoded, original); + assert!(encoded.content.get("stream_options").is_none()); + } +} + +#[test] +fn request_schema_fixture_round_trips_and_preserves_adjacent_native_messages() { + let codec = OpenAIChatCodec; + let original = make_request(json!({ + "model": "gpt-4.1", + "messages": [ + {"role": "developer", "content": "Be exact.", "name": "policy"}, + {"role": "user", "content": [ + {"type": "text", "text": "Describe this", "future_part": 1}, + {"type": "image_url", "image_url": {"url": "https://example.com/a.png", "detail": "high"}}, + {"type": "input_audio", "input_audio": {"data": "AAAA", "format": "wav"}}, + {"type": "file", "file": {"file_id": "file_1"}} + ]}, + {"role": "assistant", "content": null, "refusal": null, "audio": {"id": "audio_1"}, "tool_calls": [{ + "id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{ \"q\" : \"x\" }"} + }]}, + {"role": "tool", "tool_call_id": "call_1", "content": "result"}, + {"role": "function", "name": "legacy", "content": null} + ], + "audio": {"format": "wav", "voice": "alloy"}, + "frequency_penalty": 0.1, + "function_call": "auto", + "functions": [{"name": "legacy", "parameters": {"type": "object"}}], + "logit_bias": {"42": 1}, + "logprobs": true, + "max_completion_tokens": 128, + "metadata": {"tenant": "a"}, + "modalities": ["text", "audio"], + "moderation": {"mode": "auto"}, + "n": 1, + "parallel_tool_calls": true, + "prediction": {"type": "content", "content": "prefix"}, + "presence_penalty": 0.2, + "prompt_cache_key": "cache-key", + "prompt_cache_options": {"scope": "request"}, + "prompt_cache_retention": "24h", + "reasoning_effort": "medium", + "response_format": {"type": "json_object"}, + "safety_identifier": "safe-user", + "seed": 7, + "service_tier": "auto", + "stop": "END", + "store": true, + "stream": true, + "stream_options": {"include_usage": false}, + "temperature": 0.3, + "tool_choice": {"type": "function", "function": {"name": "lookup"}}, + "tools": [ + {"type": "function", "function": {"name": "lookup", "description": "Lookup", "parameters": {"type": "object"}, "strict": true}}, + {"type": "custom", "custom": {"name": "grammar", "format": {"type": "grammar", "grammar": "start: WORD"}}} + ], + "top_logprobs": 3, + "top_p": 0.8, + "user": "user_1", + "verbosity": "low", + "web_search_options": {"search_context_size": "low"}, + "future_field": null + })); + + let mut annotated = codec.decode(&original).unwrap(); + assert_eq!(codec.encode(&annotated, &original).unwrap(), original); + + let Message::User { + content: MessageContent::Parts(parts), + .. + } = &mut annotated.messages[1] + else { + panic!("expected portable user message parts"); + }; + let ContentPart::Text { text, extra } = &mut parts[0] else { + panic!("expected portable text part"); }; + *text = "Summarize this".into(); + assert_eq!(extra.get("future_part"), Some(&json!(1))); - let encoded = codec - .encode( - &annotated, - &make_request(json!({"messages": [], "model": "gpt-4o"})), - ) - .unwrap(); - let obj = encoded.content.as_object().unwrap(); - assert!( - !obj.contains_key("stream_options"), - "stream_options must not be injected when stream key is absent", + let encoded = codec.encode(&annotated, &original).unwrap(); + let mut expected = original.clone(); + expected.content["messages"][1]["content"][0]["text"] = json!("Summarize this"); + assert_eq!(encoded, expected); +} + +#[test] +fn chat_tool_edits_preserve_nested_unknown_fields_and_explicit_nulls() { + let codec = OpenAIChatCodec; + let original = make_request(json!({ + "model": "gpt-4.1", + "messages": [{"role": "user", "content": "hello"}], + "tools": [{ + "type": "function", + "function": { + "name": "lookup_before", + "description": null, + "parameters": {"type": "object"}, + "strict": null, + "future_function": {"mode": "keep"} + }, + "future_wrapper": null + }], + "tool_choice": { + "type": "function", + "function": {"name": "lookup_before", "future_nested": null}, + "future_choice": {"mode": "keep"} + } + })); + let mut annotated = codec.decode(&original).unwrap(); + let ToolDefinition::Function { function, .. } = &mut annotated.tools.as_mut().unwrap()[0] + else { + panic!("expected portable function tool"); + }; + function.name = "lookup_after".into(); + let Some(ToolChoice::Specific(choice)) = &mut annotated.tool_choice else { + panic!("expected specific tool choice"); + }; + choice.function.name = "lookup_after".into(); + + let encoded = codec.encode(&annotated, &original).unwrap(); + let mut expected = original; + expected.content["tools"][0]["function"]["name"] = json!("lookup_after"); + expected.content["tool_choice"]["function"]["name"] = json!("lookup_after"); + assert_eq!(encoded, expected); +} + +#[test] +fn chat_reordered_content_parts_keep_nested_provider_fields_with_their_items() { + let codec = OpenAIChatCodec; + let original = make_request(json!({ + "model": "gpt-4.1", + "messages": [{ + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": "a", "vendor_marker": "A"} + }, + { + "type": "image_url", + "image_url": {"url": "b", "vendor_marker": "B"} + } + ] + }] + })); + let mut annotated = codec.decode(&original).unwrap(); + let Message::User { + content: MessageContent::Parts(parts), + .. + } = &mut annotated.messages[0] + else { + panic!("expected portable user content parts"); + }; + parts.swap(0, 1); + + let encoded = codec.encode(&annotated, &original).unwrap(); + assert_eq!( + encoded.content["messages"][0]["content"], + json!([ + { + "type": "image_url", + "image_url": {"url": "b", "vendor_marker": "B"} + }, + { + "type": "image_url", + "image_url": {"url": "a", "vendor_marker": "A"} + } + ]) ); +} - let encoded = codec - .encode( - &annotated, - &make_request(json!({ - "messages": [], - "model": "gpt-4o", - "stream": false, - })), - ) - .unwrap(); - let obj = encoded.content.as_object().unwrap(); +#[test] +fn chat_reordered_and_edited_content_parts_are_rejected_without_provenance() { + let codec = OpenAIChatCodec; + let original = make_request(json!({ + "model": "gpt-4.1", + "messages": [{ + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": "a", "vendor_marker": "A"} + }, + { + "type": "image_url", + "image_url": {"url": "b", "vendor_marker": "B"} + } + ] + }] + })); + let mut annotated = codec.decode(&original).unwrap(); + let Message::User { + content: MessageContent::Parts(parts), + .. + } = &mut annotated.messages[0] + else { + panic!("expected portable user content parts"); + }; + parts.swap(0, 1); + for (part, url) in parts.iter_mut().zip(["b-edited", "a-edited"]) { + let ContentPart::ImageUrl { image_url, .. } = part else { + panic!("expected image URL part"); + }; + image_url.url = url.into(); + } + + let error = codec.encode(&annotated, &original).unwrap_err(); assert!( - !obj.contains_key("stream_options"), - "stream_options must not be injected when stream: false", + error + .to_string() + .contains("multiple edited array items without stable identities") ); } +#[test] +fn native_chat_tool_choices_round_trip_without_fallback_loss() { + let codec = OpenAIChatCodec; + for tool_choice in [ + json!({"type": "custom", "custom": {"name": "grammar"}}), + json!({"type": "future_tool", "name": "next"}), + ] { + let original = make_request(json!({ + "model": "gpt-4.1", + "messages": [], + "tool_choice": tool_choice + })); + let annotated = codec.decode(&original).unwrap(); + assert!(matches!( + annotated.tool_choice, + Some(ToolChoice::ProviderNative(_)) + )); + assert_eq!(codec.encode(&annotated, &original).unwrap(), original); + } +} + +#[test] +fn request_schema_rejects_malformed_known_chat_components() { + let codec = OpenAIChatCodec; + for message in [ + json!({"role": "function", "name": "legacy", "content": 7}), + json!({ + "role": "assistant", + "content": null, + "tool_calls": [{ + "id": "call_1", + "type": "function", + "function": {"name": "lookup"} + }] + }), + ] { + let error = codec + .decode(&make_request( + json!({"model": "gpt-4.1", "messages": [message]}), + )) + .unwrap_err(); + assert!(error.to_string().contains("OpenAI Chat")); + } + + for (field, malformed) in [ + ("frequency_penalty", json!("high")), + ("parallel_tool_calls", json!("yes")), + ("top_logprobs", json!(-1)), + ("prompt_cache_key", json!(7)), + ("seed", json!(1.5)), + ("metadata", json!("tenant")), + ("audio", json!("wav")), + ("logit_bias", json!("42")), + ("moderation", json!("auto")), + ("prediction", json!("content")), + ("prompt_cache_options", json!("request")), + ("response_format", json!("json_object")), + ("stream_options", json!("include_usage")), + ("web_search_options", json!("low")), + ] { + let mut content = json!({"model": "gpt-4.1", "messages": []}); + content[field] = malformed; + let error = codec.decode(&make_request(content)).unwrap_err(); + assert!( + error.to_string().contains(field), + "unexpected error: {error}" + ); + } + + let error = codec + .decode(&make_request(json!({ + "model": "gpt-4.1", + "messages": [], + "tools": [{ + "type": "function", + "function": {"name": "lookup", "strict": "yes"} + }] + }))) + .unwrap_err(); + assert!(error.to_string().contains("strict")); +} + // =================================================================== // Streaming codec tests // =================================================================== diff --git a/crates/core/tests/unit/codec/openai_responses_tests.rs b/crates/core/tests/unit/codec/openai_responses_tests.rs index 3feafeda6..8dadc2680 100644 --- a/crates/core/tests/unit/codec/openai_responses_tests.rs +++ b/crates/core/tests/unit/codec/openai_responses_tests.rs @@ -6,7 +6,10 @@ use super::*; use serde_json::json; -use super::super::request::MessageContent; +use super::super::request::{ + ContentPart, FunctionDefinition, Message, MessageContent, OpenAiImageUrl, + ProviderNativeComponent, ToolChoiceFunction, ToolChoiceFunctionName, +}; use super::super::response::{ApiSpecificResponse, FinishReason}; // ------------------------------------------------------------------- @@ -459,18 +462,21 @@ fn test_decode_request_with_input_array() { let annotated = codec.decode(&request).unwrap(); assert_eq!(annotated.model, Some("gpt-4o".into())); - // instructions becomes system message (first) - assert!(annotated.messages.len() >= 2); + assert_eq!( + annotated.instructions, + Some(MessageContent::Text("Be helpful.".into())) + ); assert_eq!(annotated.system_prompt(), Some("Be helpful.")); - // input items become messages (after system) - // System + 3 input items = 4 total messages - assert_eq!(annotated.messages.len(), 4); + assert_eq!(annotated.messages.len(), 3); // Tools present let tools = annotated.tools.unwrap(); assert_eq!(tools.len(), 1); - assert_eq!(tools[0].function.name, "calculate"); + let ToolDefinition::Function { function, .. } = &tools[0] else { + panic!("expected a portable function tool"); + }; + assert_eq!(function.name, "calculate"); } #[test] @@ -563,16 +569,16 @@ fn test_decode_request_input_array_preserves_unparsed_items_in_extra() { ] })); let annotated = codec.decode(&request).unwrap(); - // strict-first behavior: no partial message extraction on mixed arrays - assert!(annotated.messages.is_empty()); - assert_eq!( - annotated + assert!(matches!(annotated.messages[0], Message::User { .. })); + assert!(matches!( + &annotated.messages[1], + Message::ToolResultItem { call_id, output, .. } + if call_id == "call_1" && output == &json!("ok") + )); + assert!( + !annotated .extra - .get("_openai_responses_unparsed_input_items"), - Some(&json!([ - { "role": "user", "content": "hello" }, - { "type": "function_call_output", "call_id": "call_1", "output": "ok" } - ])) + .contains_key("_openai_responses_unparsed_input_items") ); } @@ -620,14 +626,12 @@ fn test_decode_request_litellm_reasoning_input_item_preserved_and_controls_extra "parallel_tool_calls": true })); let annotated = codec.decode(&request).unwrap(); - // strict-first parse: mixed input array preserved whole in extra - assert!(annotated.messages.is_empty()); - assert!( - annotated - .extra - .get("_openai_responses_unparsed_input_items") - .is_some() - ); + assert!(matches!( + &annotated.messages[0], + Message::ProviderNative { provider, kind, .. } + if provider == "openai_responses" && kind == "reasoning" + )); + assert!(matches!(annotated.messages[1], Message::User { .. })); // stable controls still extracted assert_eq!(annotated.store, Some(true)); assert_eq!(annotated.parallel_tool_calls, Some(true)); @@ -679,6 +683,241 @@ fn test_decode_request_sglang_extensions_preserved_in_extra() { ); } +#[test] +fn request_schema_fixture_round_trips_ordered_native_items_and_surgical_edits() { + let codec = OpenAIResponsesCodec; + let original = make_request(json!({ + "model": "gpt-5", + "instructions": "Be exact.", + "input": [ + {"type": "message", "role": "user", "content": [ + {"type": "input_text", "text": "Solve this", "future_part": 1}, + {"type": "input_image", "image_url": "https://example.com/a.png", "detail": "high"}, + {"type": "input_file", "file_id": "file_1", "filename": "notes.txt"} + ]}, + {"type": "function_call", "id": "fc_1", "call_id": "call_1", "name": "lookup", "arguments": "{ \"q\" : \"x\" }", "status": "completed"}, + {"type": "function_call_output", "id": "fco_1", "call_id": "call_1", "output": [{"type": "input_text", "text": "result"}], "status": "completed"}, + {"type": "reasoning", "id": "rs_1", "summary": [{"type": "summary_text", "text": "work"}], "encrypted_content": "cipher"}, + {"type": "mcp_call", "id": "mcp_1", "server_label": "docs", "name": "search", "arguments": "{}"}, + {"type": "computer_call", "id": "cmp_1", "call_id": "cmp_call", "action": {"type": "screenshot"}}, + {"type": "shell_call", "id": "sh_1", "call_id": "sh_call", "action": {"commands": ["pwd"]}}, + {"type": "apply_patch_call", "id": "patch_1", "call_id": "patch_call", "operation": {"type": "create_file", "path": "a.txt", "diff": "+a"}}, + {"type": "file_search_call", "id": "fs_1", "queries": ["docs"], "status": "completed"} + ], + "background": true, + "context_management": [{"type": "compaction", "compact_threshold": 1000}], + "conversation": "conv_1", + "include": ["reasoning.encrypted_content"], + "max_output_tokens": 256, + "max_tool_calls": 4, + "metadata": {"tenant": "a"}, + "moderation": {"mode": "auto"}, + "parallel_tool_calls": true, + "previous_response_id": "resp_prev", + "prompt": {"id": "pmpt_1", "variables": {"name": "Ada"}}, + "prompt_cache_key": "cache-key", + "prompt_cache_options": {"scope": "request"}, + "prompt_cache_retention": "24h", + "reasoning": {"effort": "high", "summary": "auto"}, + "safety_identifier": "safe-user", + "service_tier": "auto", + "store": true, + "stream": true, + "stream_options": {"include_obfuscation": true}, + "temperature": 0.2, + "text": {"format": {"type": "json_object"}, "verbosity": "low"}, + "tool_choice": {"type": "mcp", "server_label": "docs"}, + "tools": [ + {"type": "function", "name": "lookup", "description": "Lookup", "parameters": {"type": "object"}, "strict": true}, + {"type": "file_search", "vector_store_ids": ["vs_1"]}, + {"type": "web_search_preview", "search_context_size": "low"}, + {"type": "computer_use_preview", "display_width": 1024, "display_height": 768, "environment": "browser"}, + {"type": "mcp", "server_label": "docs", "server_url": "https://example.com/mcp"}, + {"type": "custom", "name": "grammar", "format": {"type": "text"}} + ], + "top_logprobs": 2, + "top_p": 0.9, + "truncation": "auto", + "user": "user_1", + "future_field": null + })); + + let mut annotated = codec.decode(&original).unwrap(); + assert_eq!(annotated.messages.len(), 9); + assert!(matches!( + &annotated.messages[3], + Message::ProviderNative { kind, .. } if kind == "reasoning" + )); + assert_eq!(codec.encode(&annotated, &original).unwrap(), original); + + let Message::User { + content: MessageContent::Parts(parts), + .. + } = &mut annotated.messages[0] + else { + panic!("expected portable Responses message"); + }; + let ContentPart::Text { text, extra } = &mut parts[0] else { + panic!("expected portable Responses text part"); + }; + *text = "Solve carefully".into(); + assert_eq!(extra.get("future_part"), Some(&json!(1))); + + let encoded = codec.encode(&annotated, &original).unwrap(); + let mut expected = original.clone(); + expected.content["input"][0]["content"][0]["text"] = json!("Solve carefully"); + assert_eq!(encoded, expected); + + let mut without_future = codec.decode(&original).unwrap(); + without_future.extra.remove("future_field"); + let encoded = codec.encode(&without_future, &original).unwrap(); + assert!(encoded.content.get("future_field").is_none()); +} + +#[test] +fn responses_tool_edits_preserve_unknown_fields_and_explicit_nulls() { + let codec = OpenAIResponsesCodec; + let original = make_request(json!({ + "model": "gpt-5", + "input": "hello", + "tools": [{ + "type": "function", + "name": "lookup_before", + "description": null, + "parameters": {"type": "object"}, + "strict": null, + "future_tool": {"mode": "keep"} + }], + "tool_choice": { + "type": "function", + "name": "lookup_before", + "future_choice": null + } + })); + let mut annotated = codec.decode(&original).unwrap(); + let ToolDefinition::Function { function, .. } = &mut annotated.tools.as_mut().unwrap()[0] + else { + panic!("expected portable function tool"); + }; + function.name = "lookup_after".into(); + let Some(ToolChoice::Specific(choice)) = &mut annotated.tool_choice else { + panic!("expected specific tool choice"); + }; + choice.function.name = "lookup_after".into(); + + let encoded = codec.encode(&annotated, &original).unwrap(); + let mut expected = original; + expected.content["tools"][0]["name"] = json!("lookup_after"); + expected.content["tool_choice"]["name"] = json!("lookup_after"); + assert_eq!(encoded, expected); +} + +#[test] +fn responses_wrapped_function_tool_edits_preserve_the_wrapped_representation() { + let codec = OpenAIResponsesCodec; + let original = make_request(json!({ + "model": "gpt-5", + "input": "hello", + "tools": [{ + "type": "function", + "function": { + "name": "lookup_before", + "description": null, + "strict": null, + "future_function": {"mode": "keep"} + }, + "future_wrapper": null + }] + })); + let mut annotated = codec.decode(&original).unwrap(); + let ToolDefinition::Function { function, .. } = &mut annotated.tools.as_mut().unwrap()[0] + else { + panic!("expected portable function tool"); + }; + function.name = "lookup_after".into(); + + let encoded = codec.encode(&annotated, &original).unwrap(); + let mut expected = original; + expected.content["tools"][0]["function"]["name"] = json!("lookup_after"); + assert_eq!(encoded, expected); + assert!(encoded.content["tools"][0].get("name").is_none()); +} + +#[test] +fn responses_string_input_and_native_surface_mismatch_are_explicit() { + let codec = OpenAIResponsesCodec; + let original = make_request(json!({ + "model": "gpt-5", + "input": "hello", + "stream_options": null + })); + let mut annotated = codec.decode(&original).unwrap(); + assert_eq!(codec.encode(&annotated, &original).unwrap(), original); + + annotated.messages.push(Message::ProviderNative { + provider: "anthropic_messages".into(), + kind: "thinking".into(), + value: json!({"type": "thinking", "thinking": "private"}), + }); + let error = codec.encode(&annotated, &original).unwrap_err(); + assert!( + error + .to_string() + .contains("cannot be encoded for OpenAI Responses") + ); +} + +#[test] +fn request_schema_rejects_malformed_known_responses_items() { + let codec = OpenAIResponsesCodec; + for item in [ + json!({ + "type": "function_call", + "call_id": "call_1", + "name": "lookup" + }), + json!({"type": "function_call_output", "call_id": "call_1"}), + ] { + let error = codec + .decode(&make_request(json!({"model": "gpt-5", "input": [item]}))) + .unwrap_err(); + assert!(error.to_string().contains("missing")); + } + + for (field, malformed) in [ + ("background", json!("yes")), + ("max_output_tokens", json!(-1)), + ("parallel_tool_calls", json!("yes")), + ("previous_response_id", json!(7)), + ("reasoning", json!("high")), + ("include", json!("reasoning.encrypted_content")), + ("metadata", json!("tenant")), + ("context_management", json!("compaction")), + ("moderation", json!("auto")), + ("prompt", json!("pmpt_1")), + ("prompt_cache_options", json!("request")), + ("stream_options", json!("obfuscation")), + ("text", json!("json_object")), + ] { + let mut content = json!({"model": "gpt-5", "input": []}); + content[field] = malformed; + let error = codec.decode(&make_request(content)).unwrap_err(); + assert!( + error.to_string().contains(field), + "unexpected error: {error}" + ); + } + + let error = codec + .decode(&make_request(json!({ + "model": "gpt-5", + "input": [], + "tools": [{"type": "function", "name": "lookup", "strict": "yes"}] + }))) + .unwrap_err(); + assert!(error.to_string().contains("strict")); +} + // =================================================================== // Request encode tests // =================================================================== @@ -836,13 +1075,15 @@ fn test_helper_and_error_paths_cover_remaining_responses_branches() { }))) .unwrap_err() { - FlowError::Internal(message) => { - assert!(message.contains("OpenAI Responses tools decode")); + FlowError::InvalidArgument(message) => { + assert!(message.contains("tools must be an array")); } other => panic!("unexpected tools decode error: {other}"), } let annotated = AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![super::super::request::Message::User { content: MessageContent::Text("hello".into()), name: None, @@ -854,13 +1095,15 @@ fn test_helper_and_error_paths_cover_remaining_responses_branches() { top_p: Some(0.95), stop: None, }), - tools: Some(vec![ToolDefinition { - tool_type: "function".into(), + tools: Some(vec![ToolDefinition::Function { function: super::super::request::FunctionDefinition { name: "lookup".into(), description: Some("Look up data".into()), parameters: Some(json!({"type": "object"})), + strict: None, + extra: serde_json::Map::new(), }, + extra: serde_json::Map::new(), }]), tool_choice: Some(ToolChoice::Auto), store: None, @@ -899,12 +1142,279 @@ fn test_helper_and_error_paths_cover_remaining_responses_branches() { match codec.encode(&annotated, &make_request(json!("still-not-an-object"))) { Err(FlowError::Internal(message)) => { - assert!(message.contains("original content is not an object")); + assert!(message.contains("not an object")); } other => panic!("unexpected encode result: {other:?}"), } } +#[test] +fn responses_request_component_branch_matrix() { + assert!(decode_responses_content(&json!({})).is_err()); + for invalid in [ + json!(42), + json!({"type": "input_text"}), + json!({"type": "refusal"}), + ] { + assert!(decode_responses_content_part(&invalid).is_err()); + } + for valid in [ + json!({"type": "input_text", "text": "hello", "future": 1}), + json!({"type": "output_text", "text": "hello"}), + json!({"type": "input_image", "image_url": "https://example.com/a.png", "detail": "high", "future": 1}), + json!({"type": "input_file", "file_id": "file_1", "filename": "a.txt", "future": 1}), + json!({"type": "refusal", "refusal": "no", "future": 1}), + json!({"type": "future_part", "payload": 1}), + ] { + assert!(decode_responses_content_part(&valid).is_ok()); + } + + for invalid in [ + json!(42), + json!({"role": "user"}), + json!({"type": "function_call", "name": "lookup", "arguments": "{}"}), + json!({"type": "function_call", "call_id": "call", "arguments": "{}"}), + json!({"type": "function_call", "call_id": "call", "name": "lookup"}), + json!({"type": "function_call", "id": 7, "call_id": "call", "name": "lookup", "arguments": "{}"}), + json!({"type": "function_call_output", "output": "ok"}), + json!({"type": "function_call_output", "call_id": "call"}), + json!({"type": "function_call_output", "id": 7, "call_id": "call", "output": "ok"}), + ] { + assert!(decode_responses_input_item(&invalid).is_err()); + } + for portable in [ + json!({"type": "message", "role": "user", "content": "u"}), + json!({"type": "message", "role": "system", "content": "s"}), + json!({"type": "message", "role": "developer", "content": "d"}), + json!({"type": "message", "role": "assistant", "content": "a"}), + json!({"type": "function_call", "id": null, "call_id": "call", "name": "lookup", "arguments": "raw"}), + json!({"type": "function_call_output", "id": null, "call_id": "call", "output": "ok"}), + ] { + assert!(!matches!( + decode_responses_input_item(&portable).unwrap(), + Message::ProviderNative { .. } + )); + } + for native in [ + json!({"type": "message", "role": "user", "content": "u", "future": 1}), + json!({"type": "message", "role": "future", "content": "x"}), + json!({"type": "reasoning", "summary": []}), + ] { + assert!(matches!( + decode_responses_input_item(&native).unwrap(), + Message::ProviderNative { .. } + )); + } + + let content = MessageContent::Parts(vec![ + ContentPart::Text { + text: "hello".into(), + extra: serde_json::Map::from_iter([("future".into(), json!(1))]), + }, + ContentPart::ImageUrl { + image_url: OpenAiImageUrl { + url: "https://example.com/a.png".into(), + detail: Some("high".into()), + }, + extra: serde_json::Map::new(), + }, + ContentPart::Image { + image: json!({"file_id": "file_1"}), + extra: serde_json::Map::from_iter([("detail".into(), json!("low"))]), + }, + ContentPart::File { + file: json!({"file_id": "file_2"}), + extra: serde_json::Map::from_iter([("filename".into(), json!("a.txt"))]), + }, + ContentPart::ProviderNative { + provider: "openai_responses".into(), + kind: "future".into(), + value: json!({"type": "future"}), + }, + ]); + assert_eq!( + encode_responses_content(&content, false) + .unwrap() + .as_array() + .unwrap() + .len(), + 5 + ); + let assistant_content = MessageContent::Parts(vec![ + ContentPart::Text { + text: "answer".into(), + extra: serde_json::Map::new(), + }, + ContentPart::Refusal { + refusal: "no".into(), + extra: serde_json::Map::new(), + }, + ]); + assert_eq!( + encode_responses_content(&assistant_content, true).unwrap()[0]["type"], + json!("output_text") + ); + for invalid in [ + ContentPart::Image { + image: json!("bad"), + extra: serde_json::Map::new(), + }, + ContentPart::File { + file: json!("bad"), + extra: serde_json::Map::new(), + }, + ContentPart::Refusal { + refusal: "not user input".into(), + extra: serde_json::Map::new(), + }, + ] { + assert!(encode_responses_content(&MessageContent::Parts(vec![invalid]), false).is_err()); + } + + let items = vec![ + Message::User { + content: MessageContent::Text("u".into()), + name: None, + }, + Message::System { + content: MessageContent::Text("s".into()), + name: None, + }, + Message::Developer { + content: MessageContent::Text("d".into()), + name: None, + }, + Message::Assistant { + content: Some(assistant_content), + tool_calls: None, + name: None, + }, + Message::ToolCallItem { + id: Some("fc_1".into()), + call_id: "call_1".into(), + name: "lookup".into(), + arguments: json!({"q": "x"}), + extra: serde_json::Map::from_iter([("status".into(), json!("completed"))]), + }, + Message::ToolCallItem { + id: None, + call_id: "call_2".into(), + name: "lookup".into(), + arguments: json!("{ raw }"), + extra: serde_json::Map::new(), + }, + Message::ToolResultItem { + id: Some("fco_1".into()), + call_id: "call_1".into(), + output: json!({"ok": true}), + extra: serde_json::Map::new(), + }, + Message::ProviderNative { + provider: "openai_responses".into(), + kind: "reasoning".into(), + value: json!({"type": "reasoning"}), + }, + ]; + for item in &items { + assert!(encode_responses_input_item(item).is_ok()); + } + assert!( + encode_responses_input_item(&Message::Assistant { + content: None, + tool_calls: None, + name: None, + }) + .is_err() + ); + + for invalid in [ + json!(42), + json!({"type": "function"}), + json!({"type": "function", "function": {}}), + ] { + assert!(decode_responses_tool(&invalid).is_err()); + } + assert!(matches!( + decode_responses_tool(&json!({"type": "web_search_preview"})).unwrap(), + ToolDefinition::ProviderNative { .. } + )); + let function_tool = ToolDefinition::Function { + function: FunctionDefinition { + name: "lookup".into(), + description: Some("Lookup".into()), + parameters: Some(json!({"type": "object"})), + strict: Some(true), + extra: serde_json::Map::from_iter([("future_function".into(), json!(1))]), + }, + extra: serde_json::Map::from_iter([("future_wrapper".into(), json!(2))]), + }; + assert_eq!( + encode_responses_tool(&function_tool).unwrap()["strict"], + json!(true) + ); + let native_tool = ToolDefinition::ProviderNative { + provider: "openai_responses".into(), + kind: "web_search".into(), + value: json!({"type": "web_search_preview"}), + }; + assert_eq!( + encode_responses_tool(&native_tool).unwrap()["type"], + json!("web_search_preview") + ); + let mismatched_tool = ToolDefinition::ProviderNative { + provider: "openai_chat".into(), + kind: "custom".into(), + value: json!({"type": "custom"}), + }; + assert!(encode_responses_tool(&mismatched_tool).is_err()); + + for (wire, expected) in [ + (json!("required"), ToolChoice::Required), + (json!({"type": "any"}), ToolChoice::Required), + (json!({"type": "none"}), ToolChoice::None), + ( + json!({"type": "tool", "name": "lookup"}), + ToolChoice::Specific(ToolChoiceFunction { + choice_type: "function".into(), + function: ToolChoiceFunctionName { + name: "lookup".into(), + }, + }), + ), + ] { + assert_eq!(decode_openai_or_anthropic_tool_choice(&wire), expected); + } + let specific = ToolChoice::Specific(ToolChoiceFunction { + choice_type: "function".into(), + function: ToolChoiceFunctionName { + name: "lookup".into(), + }, + }); + assert_eq!( + encode_responses_tool_choice(&ToolChoice::None).unwrap(), + json!("none") + ); + assert_eq!( + encode_responses_tool_choice(&specific).unwrap()["name"], + json!("lookup") + ); + let native_choice = ToolChoice::ProviderNative(ProviderNativeComponent { + provider: "openai_responses".into(), + kind: "mcp".into(), + value: json!({"type": "mcp"}), + }); + assert_eq!( + encode_responses_tool_choice(&native_choice).unwrap()["type"], + json!("mcp") + ); + let mismatched_choice = ToolChoice::ProviderNative(ProviderNativeComponent { + provider: "anthropic_messages".into(), + kind: "tool".into(), + value: json!({"type": "tool"}), + }); + assert!(encode_responses_tool_choice(&mismatched_choice).is_err()); +} + // =================================================================== // Streaming codec tests // =================================================================== diff --git a/crates/core/tests/unit/codec/parity_tests.rs b/crates/core/tests/unit/codec/parity_tests.rs index 25a015a3b..955b38e1e 100644 --- a/crates/core/tests/unit/codec/parity_tests.rs +++ b/crates/core/tests/unit/codec/parity_tests.rs @@ -680,8 +680,19 @@ fn test_request_normalization_parity() { name: None, }, ]; + assert_eq!(chat.messages, expected_messages); + let expected_user_messages = vec![Message::User { + content: MessageContent::Text("Summarize the docs.".to_string()), + name: None, + }]; + assert_eq!(anthropic.messages, expected_user_messages); + assert_eq!(responses.messages, expected_user_messages); + assert_eq!( + anthropic.instructions, + Some(MessageContent::Text("You are terse.".to_string())) + ); + assert_eq!(anthropic.instructions, responses.instructions); for decoded in [&chat, &anthropic, &responses] { - assert_eq!(decoded.messages, expected_messages); assert_eq!(decoded.model.as_deref(), Some("parity-request-model")); assert_eq!(decoded.system_prompt(), Some("You are terse.")); assert_eq!(decoded.last_user_message(), Some("Summarize the docs.")); @@ -715,3 +726,106 @@ fn test_request_normalization_parity() { assert_eq!(chat.max_output_tokens, None); assert_eq!(anthropic.max_output_tokens, None); } + +#[test] +fn baseline_patching_keeps_raw_array_fields_with_reordered_logical_items() { + let original = json!([ + {"type": "text", "text": "a", "provider_marker": "A"}, + {"type": "text", "text": "b", "provider_marker": "B"}, + {"type": "text", "text": "c", "provider_marker": "C"} + ]); + let baseline = json!([ + {"type": "text", "text": "a"}, + {"type": "text", "text": "b"}, + {"type": "text", "text": "c"} + ]); + let edited = json!([ + {"type": "text", "text": "b"}, + {"type": "text", "text": "a"}, + {"type": "text", "text": "c"} + ]); + + assert_eq!( + super::patch_changed_json(&original, &baseline, &edited).unwrap(), + json!([ + {"type": "text", "text": "b", "provider_marker": "B"}, + {"type": "text", "text": "a", "provider_marker": "A"}, + {"type": "text", "text": "c", "provider_marker": "C"} + ]) + ); +} + +#[test] +fn baseline_patching_does_not_pair_reordered_deletion_with_unrelated_insertion() { + let original = json!([ + {"type": "text", "text": "a", "provider_marker": "A"}, + {"type": "text", "text": "b", "provider_marker": "B"} + ]); + let baseline = json!([ + {"type": "text", "text": "a"}, + {"type": "text", "text": "b"} + ]); + let edited = json!([ + {"type": "text", "text": "new"}, + {"type": "text", "text": "a"} + ]); + + assert_eq!( + super::patch_changed_json(&original, &baseline, &edited).unwrap(), + json!([ + {"type": "text", "text": "new"}, + {"type": "text", "text": "a", "provider_marker": "A"} + ]) + ); +} + +#[test] +fn baseline_patching_handles_insert_delete_and_duplicate_reorders() { + let duplicate_a = json!({"type": "text", "text": "a"}); + let b = json!({"type": "text", "text": "b"}); + let original = json!([ + {"type": "text", "text": "a", "provider_marker": "A1"}, + {"type": "text", "text": "a", "provider_marker": "A2"}, + {"type": "text", "text": "b", "provider_marker": "B"} + ]); + let baseline = Json::Array(vec![duplicate_a.clone(), duplicate_a.clone(), b.clone()]); + let edited = Json::Array(vec![ + duplicate_a.clone(), + b, + duplicate_a, + json!({"type": "text", "text": "new"}), + ]); + + assert_eq!( + super::patch_changed_json(&original, &baseline, &edited).unwrap(), + json!([ + {"type": "text", "text": "a", "provider_marker": "A1"}, + {"type": "text", "text": "b", "provider_marker": "B"}, + {"type": "text", "text": "a", "provider_marker": "A2"}, + {"type": "text", "text": "new"} + ]) + ); +} + +#[test] +fn baseline_patching_rejects_multiple_reordered_and_edited_items_without_provenance() { + let original = json!([ + {"type": "text", "text": "a", "provider_marker": "A"}, + {"type": "text", "text": "b", "provider_marker": "B"} + ]); + let baseline = json!([ + {"type": "text", "text": "a"}, + {"type": "text", "text": "b"} + ]); + let edited = json!([ + {"type": "text", "text": "b-edited"}, + {"type": "text", "text": "a-edited"} + ]); + + let error = super::patch_changed_json(&original, &baseline, &edited).unwrap_err(); + assert!( + error + .to_string() + .contains("multiple edited array items without stable identities") + ); +} diff --git a/crates/core/tests/unit/codec/request_tests.rs b/crates/core/tests/unit/codec/request_tests.rs index c35c60e24..ec164bc7e 100644 --- a/crates/core/tests/unit/codec/request_tests.rs +++ b/crates/core/tests/unit/codec/request_tests.rs @@ -13,6 +13,8 @@ use serde_json::json; #[test] fn test_annotated_llm_request_round_trip() { let req = AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![Message::User { content: MessageContent::Text("Hello".into()), name: None, @@ -137,6 +139,7 @@ fn test_message_content_text_serialization() { fn test_message_content_parts_serialization() { let content = MessageContent::Parts(vec![ContentPart::Text { text: "Hello world".into(), + extra: serde_json::Map::new(), }]); let json_val = serde_json::to_value(&content).unwrap(); assert_eq!(json_val, json!([{"type": "text", "text": "Hello world"}])); @@ -173,13 +176,15 @@ fn test_tool_call_serialization() { #[test] fn test_tool_definition_serialization() { - let td = ToolDefinition { - tool_type: "function".into(), + let td = ToolDefinition::Function { function: FunctionDefinition { name: "get_weather".into(), description: Some("Get current weather".into()), parameters: Some(json!({"type": "object", "properties": {"city": {"type": "string"}}})), + strict: None, + extra: serde_json::Map::new(), }, + extra: serde_json::Map::new(), }; let json_val = serde_json::to_value(&td).unwrap(); assert_eq!( @@ -273,6 +278,8 @@ fn test_annotated_llm_request_extra_flatten() { #[test] fn test_all_types_clone() { let req = AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![ Message::System { content: MessageContent::Text("system".into()), @@ -281,6 +288,7 @@ fn test_all_types_clone() { Message::User { content: MessageContent::Parts(vec![ContentPart::Text { text: "user part".into(), + extra: serde_json::Map::new(), }]), name: Some("alice".into()), }, @@ -292,13 +300,15 @@ fn test_all_types_clone() { top_p: Some(0.9), stop: Some(vec!["END".into()]), }), - tools: Some(vec![ToolDefinition { - tool_type: "function".into(), + tools: Some(vec![ToolDefinition::Function { function: FunctionDefinition { name: "test".into(), description: None, parameters: None, + strict: None, + extra: serde_json::Map::new(), }, + extra: serde_json::Map::new(), }]), tool_choice: Some(ToolChoice::Auto), store: None, @@ -355,6 +365,8 @@ fn test_all_types_partial_eq() { #[test] fn test_system_prompt_returns_text() { let req = AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![ Message::System { content: MessageContent::Text("Be helpful".into()), @@ -390,6 +402,8 @@ fn test_system_prompt_returns_text() { #[test] fn test_system_prompt_returns_none_when_absent() { let req = AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![Message::User { content: MessageContent::Text("Hi".into()), name: None, @@ -419,9 +433,12 @@ fn test_system_prompt_returns_none_when_absent() { #[test] fn test_system_prompt_from_parts() { let req = AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![Message::System { content: MessageContent::Parts(vec![ContentPart::Text { text: "Be concise".into(), + extra: serde_json::Map::new(), }]), name: None, }], @@ -454,6 +471,8 @@ fn test_system_prompt_from_parts() { #[test] fn test_last_user_message_returns_last() { let req = AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![ Message::User { content: MessageContent::Text("first".into()), @@ -494,6 +513,8 @@ fn test_last_user_message_returns_last() { #[test] fn test_last_user_message_returns_none_when_absent() { let req = AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![Message::System { content: MessageContent::Text("system".into()), name: None, @@ -523,6 +544,8 @@ fn test_last_user_message_returns_none_when_absent() { #[test] fn test_last_user_message_from_parts() { let req = AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![ Message::Assistant { content: None, @@ -532,6 +555,7 @@ fn test_last_user_message_from_parts() { Message::User { content: MessageContent::Parts(vec![ContentPart::Text { text: "from parts".into(), + extra: serde_json::Map::new(), }]), name: None, }, @@ -565,6 +589,8 @@ fn test_last_user_message_from_parts() { #[test] fn test_has_tool_calls_true() { let req = AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![Message::Assistant { content: None, tool_calls: Some(vec![ToolCall { @@ -602,6 +628,8 @@ fn test_has_tool_calls_true() { #[test] fn test_has_tool_calls_false_no_assistant() { let req = AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![Message::User { content: MessageContent::Text("hi".into()), name: None, @@ -631,6 +659,8 @@ fn test_has_tool_calls_false_no_assistant() { #[test] fn test_has_tool_calls_false_empty_vec() { let req = AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![Message::Assistant { content: Some(MessageContent::Text("hello".into())), tool_calls: Some(vec![]), diff --git a/crates/core/tests/unit/codec/response_tests.rs b/crates/core/tests/unit/codec/response_tests.rs index b028264c3..1dfe9047f 100644 --- a/crates/core/tests/unit/codec/response_tests.rs +++ b/crates/core/tests/unit/codec/response_tests.rs @@ -1632,6 +1632,7 @@ fn test_response_text_extracts_first_text_from_parts() { let resp = AnnotatedLlmResponse { message: Some(MessageContent::Parts(vec![ContentPart::Text { text: "Part text".into(), + extra: serde_json::Map::new(), }])), ..minimal_response() }; diff --git a/crates/core/tests/unit/dynamic_worker_tests.rs b/crates/core/tests/unit/dynamic_worker_tests.rs index 6e31a69dc..1e28358ad 100644 --- a/crates/core/tests/unit/dynamic_worker_tests.rs +++ b/crates/core/tests/unit/dynamic_worker_tests.rs @@ -399,7 +399,7 @@ async fn callback_helpers_cover_worker_response_edges() { "llm_intercept_invalid_request" => InvokeResponse { result: Some(InvokeResult::LlmRequest(LlmRequestInterceptResult { outcome: Some(JsonEnvelope { - schema: "nemo.relay.LlmRequestInterceptOutcome@1".into(), + schema: "nemo.relay.LlmRequestInterceptOutcome@2".into(), json: br#"{"request":null}"#.to_vec(), }), })), @@ -415,7 +415,7 @@ async fn callback_helpers_cover_worker_response_edges() { "llm_intercept_invalid_annotated" => InvokeResponse { result: Some(InvokeResult::LlmRequest(LlmRequestInterceptResult { outcome: Some(JsonEnvelope { - schema: "nemo.relay.LlmRequestInterceptOutcome@1".into(), + schema: "nemo.relay.LlmRequestInterceptOutcome@2".into(), json: serde_json::to_vec(&json!({ "request": valid_llm_request(), "annotated_request": 3, diff --git a/crates/core/tests/unit/observability/atof_tests.rs b/crates/core/tests/unit/observability/atof_tests.rs index d18be0989..8452ee407 100644 --- a/crates/core/tests/unit/observability/atof_tests.rs +++ b/crates/core/tests/unit/observability/atof_tests.rs @@ -75,6 +75,8 @@ fn make_scope_start_event(name: &str) -> Event { fn make_annotated_llm_event(name: &str) -> Event { let request = AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![Message::User { content: MessageContent::Text("hello".into()), name: None, diff --git a/crates/core/tests/unit/observability/openinference_tests.rs b/crates/core/tests/unit/observability/openinference_tests.rs index a7b52ea17..31352353c 100644 --- a/crates/core/tests/unit/observability/openinference_tests.rs +++ b/crates/core/tests/unit/observability/openinference_tests.rs @@ -15,8 +15,8 @@ use crate::api::scope::{event, pop_scope, push_scope}; use crate::api::tool::ToolAttributes; use crate::codec::model_pricing::pricing_test_mutex; use crate::codec::request::{ - AnnotatedLlmRequest, FunctionDefinition, GenerationParams, Message, MessageContent, - ToolDefinition, + AnnotatedLlmRequest, ContentPart, FunctionDefinition, GenerationParams, Message, + MessageContent, ToolDefinition, }; use crate::codec::response::{ AnnotatedLlmResponse, CostEstimate, CostSource, FinishReason, PricingCatalog, PricingResolver, @@ -98,6 +98,8 @@ fn assert_no_attr_contains(attributes: &HashMap, expected: &str) fn empty_annotated_request() -> AnnotatedLlmRequest { AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: Vec::new(), model: None, params: None, @@ -367,6 +369,8 @@ fn openai_chat_provider_response(model_id: &str) -> Json { fn sample_openinference_annotated_request() -> AnnotatedLlmRequest { AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![ Message::System { content: MessageContent::Text("Use concise answers.".to_string()), @@ -384,8 +388,7 @@ fn sample_openinference_annotated_request() -> AnnotatedLlmRequest { top_p: None, stop: None, }), - tools: Some(vec![ToolDefinition { - tool_type: "function".to_string(), + tools: Some(vec![ToolDefinition::Function { function: FunctionDefinition { name: "search_docs".to_string(), description: Some("Search the docs corpus.".to_string()), @@ -393,7 +396,10 @@ fn sample_openinference_annotated_request() -> AnnotatedLlmRequest { "type": "object", "properties": {"query": {"type": "string"}} })), + strict: None, + extra: serde_json::Map::new(), }, + extra: serde_json::Map::new(), }]), ..empty_annotated_request() } @@ -4225,6 +4231,95 @@ fn annotated_llm_payloads_emit_flattened_openinference_message_and_tool_attribut assert_attr(&attributes, "llm.finish_reason", "tool_use"); } +#[test] +fn annotated_input_projection_covers_extended_roles_and_native_text() { + let messages = vec![ + Message::Developer { + content: MessageContent::Text("developer".into()), + name: None, + }, + Message::Function { + content: Some("legacy".into()), + name: "legacy".into(), + }, + Message::ToolCallItem { + id: None, + call_id: "call_1".into(), + name: "lookup".into(), + arguments: json!({}), + extra: serde_json::Map::new(), + }, + Message::ToolResultItem { + id: None, + call_id: "call_1".into(), + output: json!("ok"), + extra: serde_json::Map::new(), + }, + Message::ProviderNative { + provider: "openai_responses".into(), + kind: "message".into(), + value: json!({"role": "critic", "content": "native"}), + }, + Message::ProviderNative { + provider: "openai_responses".into(), + kind: "reasoning".into(), + value: json!({"type": "reasoning"}), + }, + ]; + let mut attributes = Vec::new(); + push_annotated_input_messages(&mut attributes, &messages); + let attributes = attr_map(&attributes); + assert_attr( + &attributes, + "llm.input_messages.0.message.role", + "developer", + ); + assert_attr(&attributes, "llm.input_messages.1.message.role", "function"); + assert_attr( + &attributes, + "llm.input_messages.1.message.content", + "legacy", + ); + assert_attr( + &attributes, + "llm.input_messages.2.message.role", + "assistant", + ); + assert_attr(&attributes, "llm.input_messages.3.message.role", "tool"); + assert_attr(&attributes, "llm.input_messages.4.message.role", "critic"); + assert_attr( + &attributes, + "llm.input_messages.4.message.content", + "native", + ); + assert_attr( + &attributes, + "llm.input_messages.5.message.role", + "provider_native", + ); + + let content = MessageContent::Parts(vec![ + ContentPart::Refusal { + refusal: "portable refusal".into(), + extra: serde_json::Map::new(), + }, + ContentPart::ProviderNative { + provider: "openai_responses".into(), + kind: "output_text".into(), + value: json!({"text": "native text"}), + }, + ContentPart::ProviderNative { + provider: "openai_responses".into(), + kind: "refusal".into(), + value: json!({"refusal": "native refusal"}), + }, + ]); + assert_eq!( + message_content_text(&content).as_deref(), + Some("portable refusal\nnative text\nnative refusal") + ); +} + #[test] fn hermes_exact_api_payloads_emit_openinference_text_usage_and_metadata() { let (provider, exporter) = make_provider(); diff --git a/crates/core/tests/unit/plugin_dynamic_tests.rs b/crates/core/tests/unit/plugin_dynamic_tests.rs index 9a64bde7d..fb5d721a0 100644 --- a/crates/core/tests/unit/plugin_dynamic_tests.rs +++ b/crates/core/tests/unit/plugin_dynamic_tests.rs @@ -1837,3 +1837,16 @@ fn registry_reconstruction_and_raw_native_validation_cover_rejections() { registry.add(disabled).unwrap(); assert!(!registry.disable("acme.guardrails.pii").unwrap()); } + +#[test] +fn annotated_request_consumers_must_exclude_relay_zero_five() { + let error = validate_annotated_request_consumer_compatibility( + ">=0.5,<1.0", + "example.request_interceptor", + ) + .unwrap_err(); + assert!(error.to_string().contains(">=0.6,<1.0")); + + validate_annotated_request_consumer_compatibility(">=0.6,<1.0", "example.request_interceptor") + .unwrap(); +} diff --git a/crates/core/tests/unit/shared_tests.rs b/crates/core/tests/unit/shared_tests.rs index 019f49bcb..9d17ce188 100644 --- a/crates/core/tests/unit/shared_tests.rs +++ b/crates/core/tests/unit/shared_tests.rs @@ -26,6 +26,8 @@ struct SharedTestCodec; impl LlmCodec for SharedTestCodec { fn decode(&self, request: &LlmRequest) -> Result { Ok(AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![Message::User { content: MessageContent::Text( request.content["prompt"] diff --git a/crates/core/tests/unit/skill_load_tests.rs b/crates/core/tests/unit/skill_load_tests.rs index 10ce0f9b2..e7eead528 100644 --- a/crates/core/tests/unit/skill_load_tests.rs +++ b/crates/core/tests/unit/skill_load_tests.rs @@ -77,6 +77,26 @@ fn structured_readers_accept_every_supported_tool_and_path_field() { "mcp__filesystem__read_multiple_files", json!({"paths": ["/skills/review/SKILL.md"]}), ), + ( + "mcp__resources__read_resource", + json!({"uri": "file:///skills/review/SKILL.md"}), + ), + ( + "mcp__resources__read_resource", + json!({"path": "/skills/review/SKILL.md"}), + ), + ( + "mcp__filesystem__get_file_contents", + json!({"path": "/skills/review/SKILL.md"}), + ), + ( + "mcp__filesystem__read_file_content", + json!({"absolute_path": "/skills/review/SKILL.md"}), + ), + ( + "read_file", + json!({"uri": "skill://catalog/review/SKILL.md"}), + ), ] { assert_detects(tool, args, &["review"], SkillLoadSource::StructuredRead); } @@ -113,6 +133,9 @@ fn structured_readers_allow_zero_offset_but_reject_every_partial_read_control() for args in [ json!({"file_path": "/skills/review/SKILL.md", "offset": 1}), json!({"file_path": "/skills/review/SKILL.md", "offset": -1}), + json!({"file_path": "/skills/review/SKILL.md", "offset": 1.0}), + json!({"file_path": "/skills/review/SKILL.md", "offset": "1"}), + json!({"file_path": "/skills/review/SKILL.md", "offset": null}), json!({"file_path": "/skills/review/SKILL.md", "limit": 2000}), json!({"file_path": "/skills/review/SKILL.md", "range": "1:20"}), json!({"file_path": "/skills/review/SKILL.md", "head": true}), @@ -130,6 +153,25 @@ fn structured_readers_allow_zero_offset_but_reject_every_partial_read_control() } } +#[test] +fn structured_readers_require_paths_to_end_exactly_in_skill_md() { + for (tool, args) in [ + ("Read", json!({"file_path": "/skills/review/SKILL.md "})), + ("Read", json!({"file_path": "'/skills/review/SKILL.md"})), + ("Read", json!({"file_path": "/skills/review/SKILL.md'"})), + ( + "read_file", + json!({"absolute_path": "/skills/review/SKILL.md/"}), + ), + ( + "mcp__resources__read_resource", + json!({"uri": "file:///skills/review/SKILL.md/"}), + ), + ] { + assert_rejected(tool, args); + } +} + #[test] fn structured_readers_reject_non_skill_paths_missing_parents_and_non_read_tools() { for (tool, args) in [ @@ -148,6 +190,46 @@ fn structured_readers_reject_non_skill_paths_missing_parents_and_non_read_tools( ("thread", json!({"path": "/skills/review/SKILL.md"})), ("spread", json!({"path": "/skills/review/SKILL.md"})), ("unread", json!({"path": "/skills/review/SKILL.md"})), + ( + "mcp__resources__list_resource", + json!({"uri": "file:///skills/review/SKILL.md"}), + ), + ( + "mcp__filesystem__download_file_contents", + json!({"absolute_path": "/skills/review/SKILL.md"}), + ), + ( + "mcp__filesystem__read_file_contents", + json!({"path": "/skills/review/SKILL.md"}), + ), + ] { + assert_rejected(tool, args); + } +} + +#[test] +fn new_structured_reader_schemas_reject_partial_controls_and_malformed_uris() { + for (tool, args) in [ + ( + "mcp__resources__read_resource", + json!({"uri": "file:///skills/review/SKILL.md", "limit": 2000}), + ), + ( + "read_file", + json!({"absolute_path": "/skills/review/SKILL.md", "offset": 1}), + ), + ( + "mcp__filesystem__get_file_contents", + json!({"path": "/skills/review/SKILL.md", "range": "1:20"}), + ), + ( + "mcp__resources__read_resource", + json!({"uri": "file:///skills/review/SKILL.md?raw=true"}), + ), + ( + "mcp__resources__read_resource", + json!({"uri": "file:///SKILL.md"}), + ), ] { assert_rejected(tool, args); } @@ -158,7 +240,7 @@ fn shell_detection_accepts_complete_cat_bat_and_batcat_commands() { for (tool, command, names) in [ ( "Bash", - "cat -n '/workspace/skills/review/SKILL.md'", + "cat -u '/workspace/skills/review/SKILL.md'", &["review"][..], ), ( @@ -173,7 +255,7 @@ fn shell_detection_accepts_complete_cat_bat_and_batcat_commands() { ), ( "run_shell_command", - "batcat /skills/review/SKILL.md /skills/testing/SKILL.md /skills/review/SKILL.md", + "batcat -p /skills/review/SKILL.md /skills/testing/SKILL.md /skills/review/SKILL.md", &["review", "testing"][..], ), ] { @@ -191,6 +273,8 @@ fn shell_detection_accepts_complete_powershell_get_content_forms() { for command in [ "Get-Content -Raw -LiteralPath 'C:\\skills\\review\\SKILL.md'", "Get-Content -Encoding utf8 -Path C:\\skills\\review\\SKILL.md", + "Get-Content -AsByteStream -ReadCount 0 -Path C:\\skills\\review\\SKILL.md", + "Get-Content -Force -Path C:\\skills\\review\\SKILL.md", "C:\\Windows\\System32\\Get-Content.exe C:\\skills\\review\\SKILL.md", ] { assert_detects( @@ -202,6 +286,56 @@ fn shell_detection_accepts_complete_powershell_get_content_forms() { } } +#[test] +fn shell_detection_accepts_exact_complete_reader_wrappers() { + for (tool, command) in [ + ( + "exec_command", + "sh -c 'cat /workspace/skills/review/SKILL.md'", + ), + ( + "shell", + "bash -c 'bat --plain /workspace/skills/review/SKILL.md'", + ), + ("Bash", "bash -lc 'cat /workspace/skills/review/SKILL.md'"), + ( + "exec_command", + "zsh -c 'cat /workspace/skills/review/SKILL.md'", + ), + ( + "terminal", + "zsh -lc 'batcat --plain /workspace/skills/review/SKILL.md'", + ), + ("shell", "fish -c 'cat /workspace/skills/review/SKILL.md'"), + ( + "exec_command", + "fish --command='bat --plain /workspace/skills/review/SKILL.md'", + ), + ( + "terminal", + "powershell -Command Get-Content -Raw -Path C:\\skills\\review\\SKILL.md", + ), + ( + "exec_command", + r#"pwsh -Command "Get-Content -Raw -LiteralPath 'C:\skills\review\SKILL.md'""#, + ), + ("sh", "cat /workspace/skills/review/SKILL.md"), + ("zsh", "cat /workspace/skills/review/SKILL.md"), + ("fish", "cat /workspace/skills/review/SKILL.md"), + ( + "pwsh", + "Get-Content -Raw -Path C:\\skills\\review\\SKILL.md", + ), + ] { + assert_detects( + tool, + json!({"command": command}), + &["review"], + SkillLoadSource::ShellRead, + ); + } +} + #[test] fn shell_detection_rejects_partial_transformed_and_compound_commands() { for command in [ @@ -216,12 +350,48 @@ fn shell_detection_rejects_partial_transformed_and_compound_commands() { "Get-Content -Head 20 /skills/review/SKILL.md", "Get-Content -First 20 /skills/review/SKILL.md", "Get-Content -Last 20 /skills/review/SKILL.md", + "Get-Content -Ta 20 /skills/review/SKILL.md", + "Get-Content -To 20 /skills/review/SKILL.md", + "Get-Content -? C:\\skills\\review\\SKILL.md", "cat /skills/review/SKILL.md | head", + "cat /skills/review/SKILL.md|head", "cat /skills/review/SKILL.md > /tmp/copy", + "cat /skills/review/SKILL.md>out", "cat /skills/review/SKILL.md < /tmp/input", + "cat /skills/review/SKILL.md 2>/tmp/error", + "cat /skills/review/SKILL.md << /tmp/copy'", + "bash -lc 'cat /skills/review/SKILL.md|head'", + "bash -lc 'cat --version /skills/review/SKILL.md'", + "bash -lc 'cat -s /skills/review/SKILL.md'", + "bash -lc 'cat -nEv /skills/review/SKILL.md'", + "fish -c 'cat --squeeze-blank /skills/review/SKILL.md'", + "fish -c 'bat --style=numbers /skills/review/SKILL.md'", + "bash -lc 'cat $(find /skills -name SKILL.md)'", + "bash -lc 'cat /skills/review/SKILL.md' ignored", + "powershell -NoProfile -Command Get-Content C:\\skills\\review\\SKILL.md", + "powershell -Command Get-Content -Ta 20 C:\\skills\\review\\SKILL.md", + "pwsh -Command Get-Content -To 20 C:\\skills\\review\\SKILL.md", + "powershell -Command Get-Content -? C:\\skills\\review\\SKILL.md", + "powershell -Command Get-Content -Stream Zone.Identifier C:\\skills\\review\\SKILL.md", + "powershell -Command \"Get-Content C:\\skills\\review\\SKILL.md; Write-Host done\"", + "powershell -Command gc C:\\skills\\review\\SKILL.md", + "pwsh -Command cat C:\\skills\\review\\SKILL.md", + "cmd /c type C:\\skills\\review\\SKILL.md", + "cat '/skills/review/SKILL.md/'", + "cat '/skills/review/SKILL.md '", + ] { + assert_rejected("shell", json!({"command": command})); + } +} + #[test] fn shell_detection_rejects_malformed_unknown_and_non_command_inputs() { for (tool, args) in [ @@ -242,6 +446,14 @@ fn shell_detection_rejects_malformed_unknown_and_non_command_inputs() { ("shell", json!({"command": 7})), ("shell", json!({"script": "cat /skills/review/SKILL.md"})), ("python", json!({"command": "cat /skills/review/SKILL.md"})), + ( + "powershell", + json!({"cmd": "cat C:\\skills\\review\\SKILL.md"}), + ), + ( + "pwsh", + json!({"command": "gc C:\\skills\\review\\SKILL.md"}), + ), ] { assert_rejected(tool, args); } diff --git a/crates/core/tests/unit/types_tests.rs b/crates/core/tests/unit/types_tests.rs index 557a3286a..720651f60 100644 --- a/crates/core/tests/unit/types_tests.rs +++ b/crates/core/tests/unit/types_tests.rs @@ -93,6 +93,8 @@ crate::editor_config! { fn annotated_request(model: &str, text: &str) -> AnnotatedLlmRequest { AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![Message::User { content: MessageContent::Text(text.into()), name: None, diff --git a/crates/ffi/nemo_relay.h b/crates/ffi/nemo_relay.h index af178bac7..f5acfa076 100644 --- a/crates/ffi/nemo_relay.h +++ b/crates/ffi/nemo_relay.h @@ -1150,6 +1150,10 @@ NemoRelayStatus nemo_relay_deregister_subscriber(const char *name); /** * Wait for subscriber callbacks queued before this call to finish. + * + * Call this function outside native subscriber callbacks. A re-entrant call returns without + * waiting to avoid blocking the dispatcher, so callbacks later in the same dispatch snapshot can + * still run. */ NemoRelayStatus nemo_relay_flush_subscribers(void); @@ -1297,7 +1301,11 @@ NemoRelayStatus nemo_relay_atof_exporter_register(const struct FfiAtofExporter * NemoRelayStatus nemo_relay_atof_exporter_deregister(const char *name); /** - * Flushes the ATOF exporter output file. + * Outside a native subscriber callback, waits for queued subscriber delivery, then flushes the + * configured file sink or asks the configured stream sink to drain for up to its timeout. + * + * A re-entrant call does not establish the delivery barrier. A stream timeout is logged and does + * not by itself return an error. * * # Safety * `exporter` must be a valid, non-null pointer. @@ -1305,7 +1313,11 @@ NemoRelayStatus nemo_relay_atof_exporter_deregister(const char *name); NemoRelayStatus nemo_relay_atof_exporter_force_flush(const struct FfiAtofExporter *exporter); /** - * Shuts down the ATOF exporter by flushing output. + * Outside a native subscriber callback, waits for queued subscriber delivery, then flushes the + * configured file sink or asks the configured stream sink to drain and close up to its timeout. + * + * A re-entrant call does not establish the delivery barrier. A stream timeout is logged and does + * not by itself return an error. * * # Safety * `exporter` must be a valid, non-null pointer. diff --git a/crates/ffi/src/api/llm_registry.rs b/crates/ffi/src/api/llm_registry.rs index 543758355..1d53e4ce4 100644 --- a/crates/ffi/src/api/llm_registry.rs +++ b/crates/ffi/src/api/llm_registry.rs @@ -390,6 +390,10 @@ pub unsafe extern "C" fn nemo_relay_deregister_subscriber(name: *const c_char) - } /// Wait for subscriber callbacks queued before this call to finish. +/// +/// Call this function outside native subscriber callbacks. A re-entrant call returns without +/// waiting to avoid blocking the dispatcher, so callbacks later in the same dispatch snapshot can +/// still run. #[unsafe(no_mangle)] pub extern "C" fn nemo_relay_flush_subscribers() -> NemoRelayStatus { clear_last_error(); diff --git a/crates/ffi/src/api/observability.rs b/crates/ffi/src/api/observability.rs index 989b4f127..65c397dba 100644 --- a/crates/ffi/src/api/observability.rs +++ b/crates/ffi/src/api/observability.rs @@ -431,7 +431,11 @@ pub unsafe extern "C" fn nemo_relay_atof_exporter_deregister( } } -/// Flushes the ATOF exporter output file. +/// Outside a native subscriber callback, waits for queued subscriber delivery, then flushes the +/// configured file sink or asks the configured stream sink to drain for up to its timeout. +/// +/// A re-entrant call does not establish the delivery barrier. A stream timeout is logged and does +/// not by itself return an error. /// /// # Safety /// `exporter` must be a valid, non-null pointer. @@ -450,7 +454,11 @@ pub unsafe extern "C" fn nemo_relay_atof_exporter_force_flush( } } -/// Shuts down the ATOF exporter by flushing output. +/// Outside a native subscriber callback, waits for queued subscriber delivery, then flushes the +/// configured file sink or asks the configured stream sink to drain and close up to its timeout. +/// +/// A re-entrant call does not establish the delivery barrier. A stream timeout is logged and does +/// not by itself return an error. /// /// # Safety /// `exporter` must be a valid, non-null pointer. diff --git a/crates/ffi/tests/integration/callable_extra_tests.rs b/crates/ffi/tests/integration/callable_extra_tests.rs index b7f68845f..ca7eaf50d 100644 --- a/crates/ffi/tests/integration/callable_extra_tests.rs +++ b/crates/ffi/tests/integration/callable_extra_tests.rs @@ -257,6 +257,8 @@ fn test_callable_extra_request_intercept_and_codec_paths() { assert!(decode_err.to_string().contains("invalid JSON")); let annotated = AnnotatedLLMRequest { + instructions: None, + api_specific: None, model: Some("test-model".into()), messages: vec![], params: None, diff --git a/crates/ffi/tests/unit/callable_tests.rs b/crates/ffi/tests/unit/callable_tests.rs index b6eda8347..c127dd132 100644 --- a/crates/ffi/tests/unit/callable_tests.rs +++ b/crates/ffi/tests/unit/callable_tests.rs @@ -409,6 +409,8 @@ fn test_wrap_llm_request_intercept_with_annotated_input() { let request_intercept = wrap_llm_request_intercept_fn(llm_request_intercept_cb, std::ptr::null_mut(), None); let annotated = nemo_relay::codec::request::AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![], model: Some("test-model".into()), params: None, diff --git a/crates/ffi/tests/unit/types_tests.rs b/crates/ffi/tests/unit/types_tests.rs index 6ba6a0cfe..d8ef8bce7 100644 --- a/crates/ffi/tests/unit/types_tests.rs +++ b/crates/ffi/tests/unit/types_tests.rs @@ -559,6 +559,24 @@ fn test_llm_request_and_event_accessors() { #[test] fn test_annotated_event_accessors_and_codec_handles() { let annotated_request = nemo_relay::codec::request::AnnotatedLlmRequest { + instructions: Some(nemo_relay::codec::request::MessageContent::Text( + "follow policy".into(), + )), + api_specific: Some( + nemo_relay::codec::request::ApiSpecificRequest::OpenAIResponses { + background: Some(true), + context_management: None, + conversation: None, + moderation: None, + prompt: None, + prompt_cache_key: Some("cache-key".into()), + prompt_cache_options: None, + prompt_cache_retention: None, + safety_identifier: None, + stream_options: None, + text: None, + }, + ), messages: vec![nemo_relay::codec::request::Message::User { content: nemo_relay::codec::request::MessageContent::Text("hello".into()), name: Some("tester".into()), @@ -604,6 +622,18 @@ fn test_annotated_event_accessors_and_codec_handles() { let annotated_request_value: serde_json::Value = serde_json::from_str(&annotated_request_json).unwrap(); assert_eq!(annotated_request_value["model"], json!("gpt-test")); + assert_eq!( + annotated_request_value["instructions"], + json!("follow policy") + ); + assert_eq!( + annotated_request_value["api_specific"], + json!({ + "api": "openai_responses", + "background": true, + "prompt_cache_key": "cache-key" + }) + ); assert_eq!(annotated_request_value["provider"], json!("ffi")); assert!(unsafe { nemo_relay_event_annotated_response(&ffi_start) }.is_null()); diff --git a/crates/node/README.md b/crates/node/README.md index c3820a132..60a496506 100644 --- a/crates/node/README.md +++ b/crates/node/README.md @@ -59,7 +59,7 @@ The Node.js package provides the following capabilities: Install the npm package in a Node.js 24 or newer project: ```bash -npm install nemo-relay-node +npm install nemo-relay-node@0.6.0 ``` ## Getting Started @@ -87,6 +87,7 @@ async function main() { }); flushSubscribers(); + await new Promise((resolve) => setImmediate(resolve)); deregisterSubscriber("printer"); } @@ -96,6 +97,10 @@ main().catch((error) => { }); ``` +Native subscriber delivery is asynchronous. `flushSubscribers()` drains the +native dispatcher. The extra event-loop turn lets queued JavaScript callback +side effects complete before deregistration or exit. + The main runtime API is exported from `nemo-relay-node`. Additional entry points are available at `nemo-relay-node/typed`, `nemo-relay-node/plugin`, `nemo-relay-node/adaptive`, and `nemo-relay-node/observability`. diff --git a/crates/node/src/api/mod.rs b/crates/node/src/api/mod.rs index 4c8943a7e..12cc22b96 100644 --- a/crates/node/src/api/mod.rs +++ b/crates/node/src/api/mod.rs @@ -1924,7 +1924,7 @@ pub fn tool_call_execute( env: Env, name: String, args: Json, - func: ThreadsafeFunction, + #[napi(ts_arg_type = "(arg: Json) => any")] func: JsFunction, handle: Option<&ScopeHandle>, attributes: Option, data: Option, @@ -1934,7 +1934,8 @@ pub fn tool_call_execute( let parent = handle .map(|h| h.inner.clone()) .unwrap_or_else(task_scope_top); - let exec_fn = callable::wrap_js_tool_exec_fn(func); + let callback = callable::safe_execution_callback(&env, &func)?; + let exec_fn = callable::wrap_js_tool_exec_fn(json_callback_tsfn(&env, &callback)?); let default_fn: ToolExecutionNextFn = std::sync::Arc::new(move |args| exec_fn(args)); let scope_stack = current_scope_stack_handle(); @@ -2116,7 +2117,7 @@ pub fn llm_call_execute( env: Env, name: String, request: Json, - func: ThreadsafeFunction, + #[napi(ts_arg_type = "(arg: Json) => any")] func: JsFunction, handle: Option<&ScopeHandle>, attributes: Option, data: Option, @@ -2132,7 +2133,8 @@ pub fn llm_call_execute( .unwrap_or_else(task_scope_top); let llm_request: LlmRequest = serde_json::from_value(request) .map_err(|e| napi::Error::from_reason(format!("invalid LlmRequest: {e}")))?; - let exec_fn = callable::wrap_js_llm_exec_fn(func); + let callback = callable::safe_execution_callback(&env, &func)?; + let exec_fn = callable::wrap_js_llm_exec_fn(json_callback_tsfn(&env, &callback)?); let default_fn: LlmExecutionNextFn = std::sync::Arc::new(move |req| exec_fn(req)); let codec = match (codec_decode, codec_encode) { (Some(d), Some(e)) => Some(callable::wrap_js_codec(d, e)), @@ -2849,6 +2851,10 @@ pub fn deregister_subscriber(name: String) -> Result { /// Wait for native subscriber callbacks queued before this call to finish. /// +/// Call this function outside native subscriber callbacks. A re-entrant call returns without +/// waiting to avoid blocking the dispatcher, so callbacks later in the same dispatch snapshot can +/// still run. +/// /// JavaScript subscribers are queued through Node's `ThreadsafeFunction`; callers that /// need JS callback side effects should await an event-loop tick after this returns. #[napi] @@ -3676,7 +3682,10 @@ impl AtofExporter { .map_err(|e| napi::Error::from_reason(e.to_string())) } - /// Flush the output file. + /// Outside a native subscriber callback, wait for queued subscriber delivery, then flush the + /// file sink or ask the stream sink to drain for up to its timeout. A re-entrant call does not + /// establish the delivery barrier. A stream timeout is logged and does not by itself return an + /// error. #[napi] pub fn force_flush(&self) -> napi::Result<()> { self.inner @@ -3684,7 +3693,10 @@ impl AtofExporter { .map_err(|e| napi::Error::from_reason(e.to_string())) } - /// Shut down the exporter by flushing output. + /// Outside a native subscriber callback, wait for queued subscriber delivery, then flush the + /// file sink or ask the stream sink to drain and close up to its timeout. A re-entrant call + /// does not establish the delivery barrier. A stream timeout is logged and does not by itself + /// return an error. #[napi] pub fn shutdown(&self) -> napi::Result<()> { self.inner diff --git a/crates/node/src/callable.rs b/crates/node/src/callable.rs index 4d21d4ade..6937e8dd5 100644 --- a/crates/node/src/callable.rs +++ b/crates/node/src/callable.rs @@ -36,6 +36,7 @@ use nemo_relay::codec::response::AnnotatedLlmResponse; use nemo_relay::codec::traits::{LlmCodec, LlmResponseCodec}; use nemo_relay::error::{FlowError, Result}; +use crate::callback_factory; use crate::convert::{callback_json, record_callback_error}; use crate::promise_call::{JsonNextFn, JsonStreamNextFn, PromiseAwareFn}; use crate::types::{EventSanitizeFields, JsEvent, event_sanitize_fields_from_json}; @@ -119,6 +120,15 @@ pub(crate) fn safe_middleware_callback(env: &Env, func: &JsFunction) -> napi::Re Ok(unsafe { wrapper_unknown.cast::() }) } +/// Wrap a synchronous execution callback so failures cross the N-API boundary as data. +/// +/// The return value is validated before NAPI-RS attempts to convert it to JSON. This +/// prevents unsupported values such as `BigInt` from reaching conversion paths that +/// abort the Node process. +pub(crate) fn safe_execution_callback(env: &Env, func: &JsFunction) -> napi::Result { + callback_factory::wrap_execution_callback(env, func) +} + pub(crate) fn unwrap_middleware_result(value: Json, error_prefix: &str) -> Result { let result: MiddlewareCallbackResult = serde_json::from_value(value).map_err(|error| { FlowError::Internal(format!( @@ -314,7 +324,8 @@ pub fn wrap_js_tool_exec_fn( "failed to queue JS tool execution callback: {status:?}", ))); } - rx.await.map_err(|e| FlowError::Internal(e.to_string())) + let result = rx.await.map_err(|e| FlowError::Internal(e.to_string()))?; + unwrap_middleware_result(result, "JS tool execution callback failed") }) }) } @@ -507,7 +518,8 @@ pub fn wrap_js_llm_exec_fn( "failed to queue JS LLM execution callback: {status:?}", ))); } - rx.await.map_err(|e| FlowError::Internal(e.to_string())) + let result = rx.await.map_err(|e| FlowError::Internal(e.to_string()))?; + unwrap_middleware_result(result, "JS LLM execution callback failed") }) }) } diff --git a/crates/node/src/callback_factory.rs b/crates/node/src/callback_factory.rs new file mode 100644 index 000000000..891373dbf --- /dev/null +++ b/crates/node/src/callback_factory.rs @@ -0,0 +1,124 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Cached JavaScript callback wrapper factories for the Node binding. + +use napi::{Env, JsFunction, JsObject, JsUnknown, NapiRaw, NapiValue}; + +const CALLBACK_FACTORIES_PROPERTY: &str = "__nemo_relay_callback_factories_v1"; + +const CALLBACK_FACTORIES_SOURCE: &str = r#"(() => { + function jsonValue(value, seen = new Set()) { + if (value === null || typeof value === 'string' || typeof value === 'boolean') { + return value; + } + if (typeof value === 'number') { + if (!Number.isFinite(value)) { + throw new TypeError('JavaScript callback returned a non-finite number that cannot be converted to JSON'); + } + return value; + } + if (typeof value !== 'object') { + throw new TypeError(`JavaScript callback returned an unsupported ${typeof value} value that cannot be converted to JSON`); + } + if (seen.has(value)) { + throw new TypeError('JavaScript callback returned a circular value that cannot be converted to JSON'); + } + seen.add(value); + if (Array.isArray(value)) { + const length = value.length; + const result = new Array(length); + for (let index = 0; index < length; index += 1) { + result[index] = jsonValue(value[index], seen); + } + seen.delete(value); + return result; + } + + const prototype = Object.getPrototypeOf(value); + if (prototype !== Object.prototype && prototype !== null) { + seen.delete(value); + throw new TypeError('JavaScript callback returned an unsupported object value that cannot be converted to JSON'); + } + + const result = Object.create(null); + for (const key of Object.keys(value)) { + result[key] = jsonValue(value[key], seen); + } + seen.delete(value); + return result; + } + + return { + execution(fn) { + return function __nemo_relay_execution_wrapper(...args) { + try { + const value = fn(...args); + return { ok: true, value: jsonValue(value === undefined ? null : value) }; + } catch (error) { + let message = 'JavaScript callback failed'; + try { + message = String(error?.message ?? error); + } catch {} + return { ok: false, error: message }; + } + }; + }, + + promise(fn) { + return function __nemo_relay_promise_wrapper(error, arg0, next, resolve, reject) { + if (error != null) { + reject(error); + return; + } + Promise.resolve().then(() => ( + next === undefined ? fn(arg0) : fn(arg0, next) + )).then((value) => jsonValue(value === undefined ? null : value)).then(resolve, reject); + }; + }, + }; +})()"#; + +fn as_unknown(env: &Env, value: &T) -> JsUnknown { + unsafe { JsUnknown::from_raw_unchecked(env.raw(), value.raw()) } +} + +fn callback_factories(env: &Env) -> napi::Result { + let global = env.get_global()?; + if global.has_own_property(CALLBACK_FACTORIES_PROPERTY)? { + return global.get_named_property(CALLBACK_FACTORIES_PROPERTY); + } + + let factories: JsObject = env.run_script(CALLBACK_FACTORIES_SOURCE)?; + let object: JsFunction = global.get_named_property("Object")?; + let object = unsafe { JsObject::from_raw_unchecked(env.raw(), object.raw()) }; + let define_property: JsFunction = object.get_named_property("defineProperty")?; + let property = env.create_string(CALLBACK_FACTORIES_PROPERTY)?; + let mut descriptor = env.create_object()?; + descriptor.set_named_property("value", factories)?; + define_property.call( + None, + &[ + as_unknown(env, &global), + as_unknown(env, &property), + as_unknown(env, &descriptor), + ], + )?; + + global.get_named_property(CALLBACK_FACTORIES_PROPERTY) +} + +fn wrap_callback(env: &Env, func: &JsFunction, factory_name: &str) -> napi::Result { + let factories = callback_factories(env)?; + let factory: JsFunction = factories.get_named_property(factory_name)?; + let wrapper = factory.call(None, &[as_unknown(env, func)])?; + Ok(unsafe { wrapper.cast::() }) +} + +pub(crate) fn wrap_execution_callback(env: &Env, func: &JsFunction) -> napi::Result { + wrap_callback(env, func, "execution") +} + +pub(crate) fn wrap_promise_callback(env: &Env, func: &JsFunction) -> napi::Result { + wrap_callback(env, func, "promise") +} diff --git a/crates/node/src/lib.rs b/crates/node/src/lib.rs index 689c5f99a..c89f2f791 100644 --- a/crates/node/src/lib.rs +++ b/crates/node/src/lib.rs @@ -16,6 +16,7 @@ mod api; mod callable; +mod callback_factory; mod convert; mod promise_call; mod stream; diff --git a/crates/node/src/promise_call.rs b/crates/node/src/promise_call.rs index 92c51cc6e..bb5435207 100644 --- a/crates/node/src/promise_call.rs +++ b/crates/node/src/promise_call.rs @@ -22,6 +22,8 @@ use serde_json::Value as Json; use nemo_relay::error::{FlowError, Result as FlowResult}; +use crate::callback_factory; + pub type JsonNextFn = Arc Pin> + Send>> + Send + Sync>; pub type JsonStreamNextFn = @@ -156,8 +158,12 @@ fn build_completion_unknowns( ) -> napi::Result<(JsUnknown, JsUnknown)> { let resolve_completion = completion.clone(); let resolve = env.create_function_from_closure("__nemo_relay_resolve", move |ctx| { - let value = ctx.get::(0).unwrap_or(Json::Null); - resolve_completion.send(Ok(value)); + let value = ctx.get::(0).map_err(|error| { + FlowError::Internal(format!( + "JavaScript callback result could not be converted to JSON: {error}" + )) + }); + resolve_completion.send(value); ctx.env.get_undefined() })?; @@ -178,22 +184,6 @@ fn build_completion_unknowns( )) } -fn create_promise_wrapper(env: &Env, callable: &JsFunction) -> napi::Result { - let factory: JsFunction = env.run_script( - r#"((fn) => function __nemo_relay_promise_wrapper(error, arg0, next, resolve, reject) { - if (error != null) { - reject(error); - return; - } - Promise.resolve().then(() => ( - next === undefined ? fn(arg0) : fn(arg0, next) - )).then(resolve, reject); -})"#, - )?; - let wrapper_unknown: JsUnknown = factory.call(None, &[function_to_unknown(env, callable)])?; - Ok(unsafe { wrapper_unknown.cast::() }) -} - /// A wrapper around a JS function that can be called from any thread and /// transparently handles both synchronous and Promise return values. pub struct PromiseAwareFn { @@ -205,7 +195,7 @@ impl PromiseAwareFn { /// /// Must be called on the JS main thread (i.e., in a sync `#[napi]` function). pub fn new(env: &Env, func: &JsFunction) -> napi::Result { - let wrapper = create_promise_wrapper(env, func)?; + let wrapper = callback_factory::wrap_promise_callback(env, func)?; let mut tsfn = env.create_threadsafe_function(&wrapper, 0, |ctx: ThreadSafeCallContext| { let next = match ctx.value.next { diff --git a/crates/node/tests/codec_tests.mjs b/crates/node/tests/codec_tests.mjs index 400ad62d6..83554b24e 100644 --- a/crates/node/tests/codec_tests.mjs +++ b/crates/node/tests/codec_tests.mjs @@ -92,7 +92,7 @@ describe('Codec pipeline integration', () => { await execWithCodec( 'test-llm', makeRequest(), - async (req) => ({ + (req) => ({ choices: [ { message: { @@ -135,7 +135,7 @@ describe('Codec pipeline integration', () => { await execWithCodec( 'test-llm', makeRequest(), - async (req) => { + (req) => { executedContent = req.content; return { choices: [], @@ -222,7 +222,7 @@ describe('Codec pipeline integration', () => { await execWithCodec( 'test-llm', makeRequest(), - async (req) => ({ + (req) => ({ choices: [], }), decodeB, diff --git a/crates/node/tests/event_sanitizers_tests.mjs b/crates/node/tests/event_sanitizers_tests.mjs index 2af2b18ab..fdd67520d 100644 --- a/crates/node/tests/event_sanitizers_tests.mjs +++ b/crates/node/tests/event_sanitizers_tests.mjs @@ -135,7 +135,7 @@ describe('event sanitizer registries', () => { metadata: { background: true }, })); try { - await lib.toolCallExecute('background-tool', { raw: true }, async (args) => args); + await lib.toolCallExecute('background-tool', { raw: true }, (args) => args); lib.flushSubscribers(); await waitFor(events, 2); } finally { @@ -168,7 +168,7 @@ describe('event sanitizer registries', () => { })); lib.registerScopeSanitizeStartGuardrail(name, 0, sanitizer); try { - await lib.toolCallExecute(name, { kept: kind }, async (args) => args); + await lib.toolCallExecute(name, { kept: kind }, (args) => args); lib.flushSubscribers(); await waitFor(events, (Object.keys(invalidResults).indexOf(kind) + 1) * 2); } finally { @@ -199,7 +199,7 @@ describe('event sanitizer registries', () => { throw new Error('background sanitizer boom'); }); try { - await lib.toolCallExecute('background-throw-tool', { kept: true }, async (args) => args); + await lib.toolCallExecute('background-throw-tool', { kept: true }, (args) => args); lib.flushSubscribers(); await waitFor(events, 2); const start = events.find( diff --git a/crates/node/tests/llm_tests.mjs b/crates/node/tests/llm_tests.mjs index 2500c0a24..05f85aaf6 100644 --- a/crates/node/tests/llm_tests.mjs +++ b/crates/node/tests/llm_tests.mjs @@ -184,6 +184,42 @@ describe('LLM execute', () => { assert.equal(result, null); }); + it('sync execute rejects thrown callbacks without terminating Node', async () => { + await assert.rejects( + () => + llmCallExecute('exec_llm_throw', makeNative(), () => { + throw new Error('sync llm callback failed'); + }), + /sync llm callback failed/, + ); + + const result = await llmCallExecute('exec_llm_after_throw', makeNative(), () => ({ + ok: true, + })); + assert.deepEqual(result, { + ok: true, + }); + }); + + it('sync and async execute reject invalid JSON results without terminating Node', async () => { + const cases = [ + ['sync_bigint', () => llmCallExecute('exec_llm_bigint', makeNative(), () => ({ value: 1n }))], + ['async_bigint', () => llmCallExecuteAsync('exec_async_llm_bigint', makeNative(), async () => 1n)], + ['sync_sparse_array', () => llmCallExecute('exec_llm_sparse_array', makeNative(), () => [, 1])], + ['async_sparse_array', () => llmCallExecuteAsync('exec_async_llm_sparse_array', makeNative(), async () => [, 1])], + ]; + + for (const [kind, execute] of cases) { + await assert.rejects(execute, /bigint|undefined|json/i); + const result = await llmCallExecute(`exec_llm_after_${kind}`, makeNative(), () => ({ + ok: true, + })); + assert.deepEqual(result, { + ok: true, + }); + } + }); + it('execute records OTEL status metadata on end events', async () => { const events = []; registerSubscriber('node_llm_status_metadata_sub', (e) => events.push(e)); diff --git a/crates/node/tests/tools_tests.mjs b/crates/node/tests/tools_tests.mjs index fcc40fe0c..75c170a73 100644 --- a/crates/node/tests/tools_tests.mjs +++ b/crates/node/tests/tools_tests.mjs @@ -238,6 +238,107 @@ describe('Tool execute', () => { assert.equal(result, null); }); + it('sync execute rejects thrown callbacks without terminating Node', async () => { + await assert.rejects( + () => + toolCallExecute('exec_tool_throw', {}, () => { + throw new Error('sync tool callback failed'); + }), + /sync tool callback failed/, + ); + + const result = await toolCallExecute('exec_tool_after_throw', {}, () => ({ + ok: true, + })); + assert.deepEqual(result, { + ok: true, + }); + }); + + it('sync and async execute reject invalid JSON results without terminating Node', async () => { + const cases = [ + ['sync_bigint', () => toolCallExecute('exec_tool_bigint', {}, () => 1n)], + ['async_bigint', () => toolCallExecuteAsync('exec_async_tool_bigint', {}, async () => ({ value: 1n }))], + ['sync_date', () => toolCallExecute('exec_tool_date', {}, () => new Date())], + ['async_map', () => toolCallExecuteAsync('exec_async_tool_map', {}, async () => new Map([['value', 1]]))], + ['sync_sparse_array', () => toolCallExecute('exec_tool_sparse_array', {}, () => [, 1])], + ['async_sparse_array', () => toolCallExecuteAsync('exec_async_tool_sparse_array', {}, async () => [, 1])], + ]; + + for (const [kind, execute] of cases) { + await assert.rejects(execute, /bigint|undefined|json/i); + const result = await toolCallExecute(`exec_tool_after_${kind}`, {}, () => ({ + ok: true, + })); + assert.deepEqual(result, { + ok: true, + }); + } + }); + + it('sync and async execute reject circular results without terminating Node', async () => { + const circularResult = () => { + const result = {}; + result.self = result; + return result; + }; + + const cases = [ + ['sync', () => toolCallExecute('exec_tool_circular', {}, circularResult)], + ['async', () => toolCallExecuteAsync('exec_async_tool_circular', {}, async () => circularResult())], + ]; + + for (const [kind, execute] of cases) { + await assert.rejects(execute, /circular|json/i); + + const result = await toolCallExecute(`exec_tool_after_circular_${kind}`, {}, () => ({ + ok: true, + })); + assert.deepEqual(result, { ok: true }); + } + }); + + it('sync and async execute read stateful getter results once', async () => { + const statefulResult = () => { + let reads = 0; + return { + get value() { + reads += 1; + return reads; + }, + }; + }; + const cases = [ + () => toolCallExecute('exec_tool_stateful_getter', {}, statefulResult), + () => toolCallExecuteAsync('exec_async_tool_stateful_getter', {}, async () => statefulResult()), + ]; + + for (const execute of cases) { + const result = await execute(); + assert.deepEqual(result, { + value: 1, + }); + } + }); + + it('sync and async execute serialize arrays by index instead of their iterator', async () => { + const statefulArray = () => { + const result = [1]; + result[Symbol.iterator] = function* iterator() { + yield 99; + }; + return result; + }; + const cases = [ + () => toolCallExecute('exec_tool_iterator', {}, statefulArray), + () => toolCallExecuteAsync('exec_async_tool_iterator', {}, async () => statefulArray()), + ]; + + for (const execute of cases) { + assert.deepEqual(await execute(), [1]); + } + }); + it('execute with attributes', async () => { const result = await toolCallExecute( 'exec_attr_tool', diff --git a/crates/node/tests/typed_tests.mjs b/crates/node/tests/typed_tests.mjs index 843675a7b..6ebfc0e2f 100644 --- a/crates/node/tests/typed_tests.mjs +++ b/crates/node/tests/typed_tests.mjs @@ -814,8 +814,8 @@ describe('typedLlmStreamExecute', () => { assert.equal(collected.length, 2); assert.notEqual(interceptedAnnotated, null); assert.equal(interceptedAnnotated.model, 'gpt-4.1-mini'); - assert.equal(interceptedAnnotated.messages[0].role, 'system'); - assert.equal(interceptedAnnotated.messages[1].role, 'user'); + assert.equal(interceptedAnnotated.instructions, 'Be terse.'); + assert.equal(interceptedAnnotated.messages[0].role, 'user'); const deadline = Date.now() + 2000; while ( diff --git a/crates/pii-redaction/README.md b/crates/pii-redaction/README.md index 731cfd499..911dc2ddc 100644 --- a/crates/pii-redaction/README.md +++ b/crates/pii-redaction/README.md @@ -34,6 +34,8 @@ NeMo Relay PII Redaction allows you to: structured secrets, and cloud credentials. - Handle codec-aware LLMs with overlay support for `openai_chat`, `openai_responses`, and `anthropic_messages`. +- Remove conversational trajectory content while preserving event structure, + tool-call identity, model attribution, routing, usage, and cost analytics. - Use the `local_model` config contract and provider registration surface for future model-backed implementations. @@ -108,6 +110,54 @@ required. Profile-array mode covers marks, LLM and tool observability, and scope metadata automatically. The original single-policy surface flags remain available for backward compatibility but cannot be combined with `profiles`. +### Structure-Preserving Trajectory Export + +Use the `trajectory_context` preset when exported trajectories must retain +their analytical structure without retaining chat, reasoning, tool, or +multimodal content. Pair it with a later email profile so email addresses are +also removed from otherwise-preserved metadata and custom marks: + +```toml +[[components]] +kind = "pii_redaction" +enabled = true + +[components.config] +version = 1 +codec = "anthropic_messages" + +[[components.config.profiles]] +mode = "builtin" +priority = 80 + +[components.config.profiles.builtin] +preset = "trajectory_context" +custom_mark_payload_policy = "redact_all_leaves" + +[[components.config.profiles]] +mode = "builtin" +priority = 90 + +[components.config.profiles.builtin] +action = "redact" +detector = "email" +``` + +`custom_mark_payload_policy = "preserve"` is the default and leaves unknown +plugin mark payloads intact for analysis. Use `redact_all_leaves` when opaque +plugins may emit content: scalar leaves in data, metadata, and opaque category +profile fields are replaced while typed category identity remains valid. +Strings become `[REDACTED]`, numbers become `0`, booleans become `false`, and +nulls, keys, arrays, and object shape are retained. +Known Relay marks are sanitized semantically so their structural and analytical +fields remain usable. This choice affects canonical event fields before +subscriber fan-out; exporter-owned resource attributes are outside this +boundary. + +The preset defines its own action and therefore cannot be combined with +`action`, `detector`, `pattern`, `target_paths`, or mask-specific fields. Its +optional `replacement` defaults to `[REDACTED]`. + ## Built-In Backend The shipped `builtin` backend supports these actions: diff --git a/crates/pii-redaction/src/builtin.rs b/crates/pii-redaction/src/builtin.rs index d1b88dd4e..43cc9cdb8 100644 --- a/crates/pii-redaction/src/builtin.rs +++ b/crates/pii-redaction/src/builtin.rs @@ -23,6 +23,7 @@ use nemo_relay::plugin::{PluginError, Result as PluginResult}; use super::component::BuiltinBackendConfig; use super::detectors::BuiltinDetector; use super::overlay::BuiltinCodecName; +use super::trajectory::{CustomMarkPayloadPolicy, TrajectorySanitizer}; #[derive(Clone)] pub(super) struct CompiledBuiltinBackend { @@ -31,6 +32,7 @@ pub(super) struct CompiledBuiltinBackend { request_codec: Option>, response_codec: Option>, codec_name: Option, + trajectory: Option, } #[derive(Clone)] @@ -71,6 +73,48 @@ impl CompiledBuiltinBackend { config: BuiltinBackendConfig, codec_name: Option, ) -> PluginResult { + let trajectory = match config.preset.as_deref() { + Some("trajectory_context") => { + if config.detector.is_some() + || config.pattern.is_some() + || !config.target_paths.is_empty() + || config.mask_char.is_some() + || config.unmasked_prefix.is_some() + || config.unmasked_suffix.is_some() + { + return Err(PluginError::InvalidConfig( + "builtin.preset cannot be combined with matcher, target-path, or mask fields" + .to_string(), + )); + } + let policy = CustomMarkPayloadPolicy::parse(&config.custom_mark_payload_policy) + .ok_or_else(|| { + PluginError::InvalidConfig(format!( + "unsupported custom-mark payload policy '{}'", + config.custom_mark_payload_policy + )) + })?; + Some(TrajectorySanitizer::new( + config + .replacement + .clone() + .unwrap_or_else(|| "[REDACTED]".to_string()), + policy, + )) + } + Some(other) => { + return Err(PluginError::InvalidConfig(format!( + "unsupported builtin preset '{other}'" + ))); + } + None => None, + }; + if trajectory.is_none() && config.custom_mark_payload_policy != "preserve" { + return Err(PluginError::InvalidConfig( + "builtin.custom_mark_payload_policy requires builtin.preset = 'trajectory_context'" + .to_string(), + )); + } let detector = config .detector .as_deref() @@ -127,6 +171,7 @@ impl CompiledBuiltinBackend { request_codec: surface.map(build_request_codec), response_codec: surface.map(build_response_codec), codec_name: surface.map(BuiltinCodecName::from_provider_surface), + trajectory, }) } @@ -258,7 +303,12 @@ impl CompiledBuiltinBackend { } pub(super) fn tool_sanitize_callback(backend: CompiledBuiltinBackend) -> ToolSanitizeFn { - Arc::new(move |_name: &str, payload: Json| backend.sanitize_json_preorder_dfs(payload)) + Arc::new( + move |_name: &str, payload: Json| match backend.trajectory.as_ref() { + Some(trajectory) => trajectory.sanitize_tool_payload(payload), + None => backend.sanitize_json_preorder_dfs(payload), + }, + ) } pub(super) fn event_sanitize_callback(backend: CompiledBuiltinBackend) -> EventSanitizeFn { @@ -291,6 +341,9 @@ fn event_sanitize_callback_with_scope_categories( return fields; } + if let Some(trajectory) = backend.trajectory.as_ref() { + return trajectory.sanitize_event_fields(event, fields); + } let specialized_scope = matches!(event, Event::Scope(_)) && event .category() @@ -316,6 +369,10 @@ pub(super) fn llm_sanitize_request_callback( backend: CompiledBuiltinBackend, ) -> LlmSanitizeRequestFn { Arc::new(move |mut request: LlmRequest| { + if let Some(trajectory) = backend.trajectory.as_ref() { + request.content = trajectory.sanitize_provider_payload(request.content); + return request; + } if let Some(encoded) = backend.sanitize_request_with_codec(&request) { return encoded; } @@ -328,6 +385,9 @@ pub(super) fn llm_sanitize_response_callback( backend: CompiledBuiltinBackend, ) -> LlmSanitizeResponseFn { Arc::new(move |payload: Json| { + if let Some(trajectory) = backend.trajectory.as_ref() { + return trajectory.sanitize_provider_payload(payload); + } if backend.target_paths.is_empty() { return backend.sanitize_json_preorder_dfs(payload); } diff --git a/crates/pii-redaction/src/component.rs b/crates/pii-redaction/src/component.rs index 275bf2ebe..592c157c3 100644 --- a/crates/pii-redaction/src/component.rs +++ b/crates/pii-redaction/src/component.rs @@ -175,8 +175,15 @@ impl Default for PiiRedactionProfile { #[derive(Debug, Clone, Serialize, Deserialize)] #[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] pub struct BuiltinBackendConfig { + /// Optional semantic sanitization preset. + #[serde(default, skip_serializing_if = "Option::is_none")] + #[cfg_attr(feature = "schema", schemars(schema_with = "builtin_preset_schema"))] + pub preset: Option, /// Action applied to matching string leaves. - #[serde(default = "default_builtin_action")] + #[serde( + default = "default_builtin_action", + skip_serializing_if = "is_default_builtin_action" + )] #[cfg_attr(feature = "schema", schemars(schema_with = "builtin_action_schema"))] pub action: String, /// Exact JSON-pointer paths to sanitize. Empty means every string leaf. @@ -200,11 +207,22 @@ pub struct BuiltinBackendConfig { /// Number of trailing characters to keep when `action = "mask"`. #[serde(default, skip_serializing_if = "Option::is_none")] pub unmasked_suffix: Option, + /// How the trajectory preset handles opaque custom-mark payloads. + #[serde( + default = "default_custom_mark_payload_policy", + skip_serializing_if = "is_default_custom_mark_payload_policy" + )] + #[cfg_attr( + feature = "schema", + schemars(schema_with = "custom_mark_payload_policy_schema") + )] + pub custom_mark_payload_policy: String, } impl Default for BuiltinBackendConfig { fn default() -> Self { Self { + preset: None, action: default_builtin_action(), target_paths: Vec::new(), pattern: None, @@ -213,6 +231,7 @@ impl Default for BuiltinBackendConfig { mask_char: None, unmasked_prefix: None, unmasked_suffix: None, + custom_mark_payload_policy: default_custom_mark_payload_policy(), } } } @@ -327,6 +346,12 @@ static PII_REDACTION_PROFILE_LIST_ITEM: nemo_relay::config_editor::EditorListIte nemo_relay::editor_config! { impl BuiltinBackendConfig { + preset => { + label: "preset", + kind: Enum, + values: ["trajectory_context"], + optional: true, + }, action => { label: "action", kind: Enum, @@ -359,6 +384,11 @@ nemo_relay::editor_config! { mask_char => { label: "mask_char", kind: String, optional: true }, unmasked_prefix => { label: "unmasked_prefix", kind: Integer, optional: true }, unmasked_suffix => { label: "unmasked_suffix", kind: Integer, optional: true }, + custom_mark_payload_policy => { + label: "custom_mark_payload_policy", + kind: Enum, + values: ["preserve", "redact_all_leaves"], + }, } } @@ -439,6 +469,24 @@ fn builtin_action_schema( ) } +#[cfg(feature = "schema")] +fn builtin_preset_schema( + generator: &mut schemars::r#gen::SchemaGenerator, +) -> schemars::schema::Schema { + string_enum_schema(generator, &["trajectory_context"], None) +} + +#[cfg(feature = "schema")] +fn custom_mark_payload_policy_schema( + generator: &mut schemars::r#gen::SchemaGenerator, +) -> schemars::schema::Schema { + string_enum_schema( + generator, + &["preserve", "redact_all_leaves"], + Some("preserve"), + ) +} + #[cfg(feature = "schema")] fn codec_schema(generator: &mut schemars::r#gen::SchemaGenerator) -> schemars::schema::Schema { string_enum_schema( @@ -584,6 +632,7 @@ fn validate_pii_redaction_plugin_config( plugin_config, "builtin", &[ + "preset", "action", "target_paths", "pattern", @@ -592,6 +641,7 @@ fn validate_pii_redaction_plugin_config( "mask_char", "unmasked_prefix", "unmasked_suffix", + "custom_mark_payload_policy", ], ); validate_section_fields( @@ -612,7 +662,7 @@ fn validate_pii_redaction_plugin_config( validate_surface_selection(&mut diagnostics, &config.policy, &config); validate_codec_requirements(&mut diagnostics, &config.policy, &config); validate_builtin_mode_requirements(&mut diagnostics, &config.policy, plugin_config, &config); - validate_builtin_action_requirements(&mut diagnostics, &config.policy, &config); + validate_builtin_action_requirements(&mut diagnostics, &config.policy, plugin_config, &config); validate_local_mode_requirements(&mut diagnostics, &config.policy, plugin_config, &config); diagnostics @@ -684,6 +734,7 @@ fn validate_profile_configuration( raw_profile, "builtin", &[ + "preset", "action", "target_paths", "pattern", @@ -692,6 +743,7 @@ fn validate_profile_configuration( "mask_char", "unmasked_prefix", "unmasked_suffix", + "custom_mark_payload_policy", ], ); validate_section_fields( @@ -717,6 +769,7 @@ fn validate_profile_configuration( validate_builtin_action_requirements( &mut profile_diagnostics, &config.policy, + raw_profile, &profile_config, ); validate_local_mode_requirements( @@ -848,12 +901,78 @@ fn validate_builtin_mode_requirements( fn validate_builtin_action_requirements( diagnostics: &mut Vec, policy: &ConfigPolicy, + plugin_config: &Map, config: &PiiRedactionConfig, ) { let Some(builtin) = config.builtin.as_ref() else { return; }; + if let Some(preset) = builtin.preset.as_deref() { + if preset != "trajectory_context" { + push_policy_diag( + diagnostics, + policy.unsupported_value, + "pii_redaction.unsupported_value", + Some(PII_REDACTION_PLUGIN_KIND.to_string()), + Some("builtin.preset".to_string()), + "builtin.preset must be 'trajectory_context'".to_string(), + ); + } + if !matches!( + builtin.custom_mark_payload_policy.as_str(), + "preserve" | "redact_all_leaves" + ) { + push_policy_diag( + diagnostics, + policy.unsupported_value, + "pii_redaction.unsupported_value", + Some(PII_REDACTION_PLUGIN_KIND.to_string()), + Some("builtin.custom_mark_payload_policy".to_string()), + "builtin.custom_mark_payload_policy must be 'preserve' or 'redact_all_leaves'" + .to_string(), + ); + } + let raw_builtin = plugin_config.get("builtin").and_then(Json::as_object); + for field in [ + "action", + "detector", + "pattern", + "target_paths", + "mask_char", + "unmasked_prefix", + "unmasked_suffix", + ] { + if raw_builtin.is_some_and(|raw| raw.contains_key(field)) { + push_policy_diag( + diagnostics, + policy.unsupported_value, + "pii_redaction.unsupported_value", + Some(PII_REDACTION_PLUGIN_KIND.to_string()), + Some(format!("builtin.{field}")), + format!("builtin.{field} cannot be combined with builtin.preset"), + ); + } + } + return; + } + + let custom_policy_configured = plugin_config + .get("builtin") + .and_then(Json::as_object) + .is_some_and(|builtin| builtin.contains_key("custom_mark_payload_policy")); + if custom_policy_configured { + push_policy_diag( + diagnostics, + policy.unsupported_value, + "pii_redaction.unsupported_value", + Some(PII_REDACTION_PLUGIN_KIND.to_string()), + Some("builtin.custom_mark_payload_policy".to_string()), + "builtin.custom_mark_payload_policy requires builtin.preset = 'trajectory_context'" + .to_string(), + ); + } + if !matches!( builtin.action.as_str(), "remove" | "redact" | "regex_replace" | "hash" | "mask" @@ -1223,6 +1342,10 @@ fn default_builtin_action() -> String { "remove".to_string() } +fn default_custom_mark_payload_policy() -> String { + "preserve".to_string() +} + fn default_true() -> bool { true } @@ -1235,6 +1358,14 @@ fn is_default_mode(mode: &str) -> bool { mode == "builtin" } +fn is_default_builtin_action(action: &str) -> bool { + action == "remove" +} + +fn is_default_custom_mark_payload_policy(policy: &str) -> bool { + policy == "preserve" +} + fn is_true(value: &bool) -> bool { *value } diff --git a/crates/pii-redaction/src/lib.rs b/crates/pii-redaction/src/lib.rs index b9acb28fb..0abb15432 100644 --- a/crates/pii-redaction/src/lib.rs +++ b/crates/pii-redaction/src/lib.rs @@ -13,6 +13,7 @@ pub mod component; pub(crate) mod detectors; pub(crate) mod local; pub(crate) mod overlay; +pub(crate) mod trajectory; #[cfg(test)] pub(crate) fn test_mutex() -> &'static Mutex<()> { diff --git a/crates/pii-redaction/src/overlay.rs b/crates/pii-redaction/src/overlay.rs index 0f97caef4..2cab7c35f 100644 --- a/crates/pii-redaction/src/overlay.rs +++ b/crates/pii-redaction/src/overlay.rs @@ -299,8 +299,18 @@ fn annotated_message_text(message: Option<&MessageContent>) -> Option { let text_parts: Vec<&str> = parts .iter() .filter_map(|part| match part { - ContentPart::Text { text } => Some(text.as_str()), - ContentPart::ImageUrl { .. } => None, + ContentPart::Text { text, .. } => Some(text.as_str()), + ContentPart::Refusal { refusal, .. } => Some(refusal.as_str()), + ContentPart::ProviderNative { value, .. } => value + .get("text") + .and_then(Json::as_str) + .or_else(|| value.get("refusal").and_then(Json::as_str)), + ContentPart::ImageUrl { .. } + | ContentPart::Image { .. } + | ContentPart::Audio { .. } + | ContentPart::File { .. } + | ContentPart::ToolUse { .. } + | ContentPart::ToolResult { .. } => None, }) .collect(); (!text_parts.is_empty()).then(|| text_parts.join("\n")) diff --git a/crates/pii-redaction/src/trajectory.rs b/crates/pii-redaction/src/trajectory.rs new file mode 100644 index 000000000..05039cc47 --- /dev/null +++ b/crates/pii-redaction/src/trajectory.rs @@ -0,0 +1,453 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Structure-preserving removal of conversational trajectory content. + +use std::sync::Arc; + +use serde_json::Value as Json; + +use nemo_relay::api::event::{CategoryProfile, Event}; +use nemo_relay::codec::request::AnnotatedLlmRequest; +use nemo_relay::codec::response::AnnotatedLlmResponse; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(super) enum CustomMarkPayloadPolicy { + Preserve, + RedactAllLeaves, +} + +impl CustomMarkPayloadPolicy { + pub(super) fn parse(value: &str) -> Option { + match value { + "preserve" => Some(Self::Preserve), + "redact_all_leaves" => Some(Self::RedactAllLeaves), + _ => None, + } + } +} + +#[derive(Clone)] +pub(super) struct TrajectorySanitizer { + replacement: Arc, + custom_mark_payload_policy: CustomMarkPayloadPolicy, +} + +impl TrajectorySanitizer { + pub(super) fn new(replacement: String, policy: CustomMarkPayloadPolicy) -> Self { + Self { + replacement: Arc::new(replacement), + custom_mark_payload_policy: policy, + } + } + + pub(super) fn sanitize_tool_payload(&self, value: Json) -> Json { + redact_all_leaves(value, &self.replacement) + } + + pub(super) fn sanitize_provider_payload(&self, value: Json) -> Json { + redact_semantic_content(value, &self.replacement, None) + } + + pub(super) fn sanitize_annotated_request( + &self, + request: AnnotatedLlmRequest, + ) -> Option { + let mut value = serde_json::to_value(request).ok()?; + let preserved = take_root_fields( + &value, + &[ + "model", + "tool_choice", + "store", + "previous_response_id", + "truncation", + "include", + "service_tier", + "parallel_tool_calls", + "max_output_tokens", + "max_tool_calls", + "top_logprobs", + "stream", + ], + ); + value = redact_semantic_content(value, &self.replacement, None); + restore_root_fields(&mut value, preserved); + serde_json::from_value(value).ok() + } + + pub(super) fn sanitize_annotated_response( + &self, + response: AnnotatedLlmResponse, + ) -> Option { + let mut value = serde_json::to_value(response).ok()?; + let mut preserved = take_root_fields( + &value, + &[ + "id", + "model", + "finish_reason", + "usage", + "optimization_summary", + ], + ); + if let Some((_, summary)) = preserved + .iter_mut() + .find(|(field, _)| field == "optimization_summary") + { + sanitize_optimization_payloads(summary, &self.replacement); + } + value = redact_semantic_content(value, &self.replacement, None); + restore_root_fields(&mut value, preserved); + serde_json::from_value(value).ok() + } + + pub(super) fn sanitize_event_fields( + &self, + event: &Event, + mut fields: nemo_relay::api::event::EventSanitizeFields, + ) -> nemo_relay::api::event::EventSanitizeFields { + let category = event.category().map(|category| category.as_str()); + let specialized_scope = + matches!(event, Event::Scope(_)) && matches!(category, Some("llm" | "tool")); + let unknown_custom_mark = matches!(event, Event::Mark(_)) + && category == Some("custom") + && !is_known_content_bearing_mark(event.name()); + + if unknown_custom_mark { + if self.custom_mark_payload_policy == CustomMarkPayloadPolicy::RedactAllLeaves { + fields.data = fields + .data + .map(|value| redact_all_leaves(value, &self.replacement)); + fields.metadata = fields + .metadata + .map(|value| redact_all_leaves(value, &self.replacement)); + fields.category_profile = fields + .category_profile + .and_then(|profile| redact_custom_category_profile(profile, self)); + } + return fields; + } + + if !specialized_scope { + fields.data = fields + .data + .map(|value| redact_semantic_content(value, &self.replacement, None)); + } + fields.metadata = fields + .metadata + .map(|value| redact_semantic_content(value, &self.replacement, None)); + fields.category_profile = fields + .category_profile + .and_then(|profile| sanitize_category_profile(profile, &self.replacement)); + fields + } +} + +fn is_known_content_bearing_mark(name: &str) -> bool { + matches!( + name, + "llm.chunk" | "nemo_relay.llm.optimization" | "skill.load" + ) +} + +fn sanitize_category_profile( + mut profile: CategoryProfile, + replacement: &str, +) -> Option { + profile.annotated_request = profile.annotated_request.as_ref().and_then(|request| { + TrajectorySanitizer::new(replacement.to_string(), CustomMarkPayloadPolicy::Preserve) + .sanitize_annotated_request((**request).clone()) + .map(Arc::new) + }); + profile.annotated_response = profile.annotated_response.as_ref().and_then(|response| { + TrajectorySanitizer::new(replacement.to_string(), CustomMarkPayloadPolicy::Preserve) + .sanitize_annotated_response((**response).clone()) + .map(Arc::new) + }); + profile.extra = profile + .extra + .into_iter() + .map(|(key, value)| { + let value = redact_semantic_content(value, replacement, Some(&key)); + (key, value) + }) + .collect(); + Some(profile) +} + +fn redact_custom_category_profile( + mut profile: CategoryProfile, + sanitizer: &TrajectorySanitizer, +) -> Option { + profile.annotated_request = profile.annotated_request.as_ref().and_then(|request| { + sanitizer + .sanitize_annotated_request((**request).clone()) + .map(Arc::new) + }); + profile.annotated_response = profile.annotated_response.as_ref().and_then(|response| { + sanitizer + .sanitize_annotated_response((**response).clone()) + .map(Arc::new) + }); + profile.extra = profile + .extra + .into_iter() + .map(|(key, value)| { + let value = redact_all_leaves(value, &sanitizer.replacement); + (key, value) + }) + .collect(); + Some(profile) +} + +fn take_root_fields(value: &Json, fields: &[&str]) -> Vec<(String, Json)> { + let Some(object) = value.as_object() else { + return Vec::new(); + }; + fields + .iter() + .filter_map(|field| { + object + .get(*field) + .cloned() + .map(|value| ((*field).to_string(), value)) + }) + .collect() +} + +fn restore_root_fields(value: &mut Json, fields: Vec<(String, Json)>) { + let Some(object) = value.as_object_mut() else { + return; + }; + object.extend(fields); +} + +fn sanitize_optimization_payloads(value: &mut Json, replacement: &str) { + let Some(contributions) = value.get_mut("contributions").and_then(Json::as_array_mut) else { + return; + }; + for contribution in contributions { + let Some(contribution) = contribution.as_object_mut() else { + continue; + }; + if let Some(payload) = contribution.get_mut("payload") { + *payload = redact_all_leaves(payload.take(), replacement); + } + let known = [ + "id", + "sequence", + "producer", + "kind", + "applied", + "model_transition", + "token_impact", + "payload_schema", + "payload", + ]; + for (key, value) in contribution.iter_mut() { + if !known.contains(&key.as_str()) { + *value = redact_all_leaves(value.take(), replacement); + } + } + } +} + +pub(super) fn redact_all_leaves(value: Json, replacement: &str) -> Json { + match value { + Json::Null => Json::Null, + Json::Bool(_) => Json::Bool(false), + Json::Number(_) => Json::from(0), + Json::String(_) => Json::String(replacement.to_string()), + Json::Array(values) => Json::Array( + values + .into_iter() + .map(|value| redact_all_leaves(value, replacement)) + .collect(), + ), + Json::Object(values) => Json::Object( + values + .into_iter() + .map(|(key, value)| (key, redact_all_leaves(value, replacement))) + .collect(), + ), + } +} + +fn redact_semantic_content(value: Json, replacement: &str, field: Option<&str>) -> Json { + match value { + Json::Null => Json::Null, + Json::Bool(value) => { + if field.is_some_and(preserve_analytical_bool) { + Json::Bool(value) + } else { + Json::Bool(false) + } + } + Json::Number(value) => { + if field.is_some_and(preserve_analytical_number) { + Json::Number(value) + } else { + Json::from(0) + } + } + Json::String(value) => { + if field.is_some_and(preserve_analytical_string) { + Json::String(value) + } else if field == Some("arguments") { + redact_stringified_json(value, replacement) + } else { + Json::String(replacement.to_string()) + } + } + Json::Array(values) => Json::Array( + values + .into_iter() + .map(|value| redact_semantic_content(value, replacement, field)) + .collect(), + ), + Json::Object(values) => { + let preserve_tool_name = preserves_tool_or_function_name(field, &values); + Json::Object( + values + .into_iter() + .map(|(key, value)| { + let value = if key == "name" && preserve_tool_name && value.is_string() { + value + } else { + redact_semantic_content(value, replacement, Some(&key)) + }; + (key, value) + }) + .collect(), + ) + } + } +} + +fn redact_stringified_json(value: String, replacement: &str) -> Json { + let Ok(parsed) = serde_json::from_str::(&value) else { + return Json::String(replacement.to_string()); + }; + let scrubbed = redact_all_leaves(parsed, replacement); + Json::String(serde_json::to_string(&scrubbed).unwrap_or_else(|_| replacement.to_string())) +} + +fn preserve_analytical_bool(key: &str) -> bool { + matches!( + key, + "applied" + | "enabled" + | "store" + | "stream" + | "parallel_tool_calls" + | "required" + | "additionalProperties" + ) +} + +fn preserve_analytical_number(key: &str) -> bool { + matches!( + key, + "chunk_index" + | "index" + | "attempt" + | "sequence" + | "priority" + | "version" + | "temperature" + | "top_p" + | "top_logprobs" + | "max_tokens" + | "max_output_tokens" + | "max_tool_calls" + | "total" + | "input" + | "output" + | "cache_read" + | "cache_write" + | "confidence" + | "logprob" + ) || key.ends_with("_tokens") + || key.ends_with("_count") + || key.ends_with("_index") + || key.ends_with("_indices") + || key.ends_with("_cost") + || key.ends_with("_latency") + || key.ends_with("_millis") + || key.ends_with("_ms") + || key.ends_with("_seconds") + || key.ends_with("_timestamp") +} + +fn preserve_analytical_string(key: &str) -> bool { + if matches!(key, "token" | "token_id") { + return false; + } + matches!( + key, + "role" + | "type" + | "api" + | "kind" + | "producer" + | "subtype" + | "model" + | "model_name" + | "provider" + | "protocol" + | "backend" + | "tier" + | "status" + | "finish_reason" + | "stop_reason" + | "service_tier" + | "mode" + | "quality" + | "estimation_method" + | "currency" + | "pricing_source" + | "pricing_provider" + | "pricing_model" + | "pricing_as_of" + | "required" + | "version" + | "event_type" + | "object" + | "object_type" + | "system_fingerprint" + | "truncation" + | "include" + | "detail" + | "media_type" + | "selected_model" + | "selected_backend" + | "selected_tier" + | "selected_protocol" + | "selected_route" + | "selected_endpoint" + | "baseline_model" + | "baseline_backend" + | "baseline_tier" + | "baseline_protocol" + | "baseline_route" + | "effective_model" + | "effective_backend" + | "effective_tier" + | "effective_protocol" + | "effective_route" + ) || key == "id" + || key.ends_with("_id") + || key.ends_with("_uuid") +} + +fn preserves_tool_or_function_name( + container: Option<&str>, + object: &serde_json::Map, +) -> bool { + matches!(container, Some("function" | "tools")) + || object + .get("type") + .and_then(Json::as_str) + .is_some_and(|kind| matches!(kind, "function" | "function_call" | "tool_use")) +} diff --git a/crates/pii-redaction/tests/coverage/overlay_tests.rs b/crates/pii-redaction/tests/coverage/overlay_tests.rs index e22d04643..00e877c21 100644 --- a/crates/pii-redaction/tests/coverage/overlay_tests.rs +++ b/crates/pii-redaction/tests/coverage/overlay_tests.rs @@ -57,6 +57,32 @@ fn openai_chat_overlay_removes_tool_calls_when_typed_entry_has_wrong_shape() { assert!(!message.contains_key("tool_calls")); } +#[test] +fn annotated_message_text_includes_provider_native_text_and_refusal_parts() { + let content = MessageContent::Parts(vec![ + ContentPart::ProviderNative { + provider: "openai_responses".into(), + kind: "output_text".into(), + value: json!({"text": "redacted text"}), + }, + ContentPart::ProviderNative { + provider: "openai_responses".into(), + kind: "refusal".into(), + value: json!({"refusal": "redacted refusal"}), + }, + ContentPart::ProviderNative { + provider: "openai_responses".into(), + kind: "reasoning".into(), + value: json!({"summary": []}), + }, + ]); + + assert_eq!( + annotated_message_text(Some(&content)).as_deref(), + Some("redacted text\nredacted refusal") + ); +} + #[test] fn openai_responses_overlay_removes_extra_function_calls() { let mut items = vec![ diff --git a/crates/pii-redaction/tests/unit/component_tests.rs b/crates/pii-redaction/tests/unit/component_tests.rs index a9c423771..9e79a56f1 100644 --- a/crates/pii-redaction/tests/unit/component_tests.rs +++ b/crates/pii-redaction/tests/unit/component_tests.rs @@ -10,11 +10,12 @@ use crate::api::event::{ ScopeCategory, ScopeEvent, }; use crate::api::llm::{ - LlmCallExecuteParams, LlmCallParams, LlmRequest, llm_call, llm_call_execute, + LlmCallExecuteParams, LlmCallParams, LlmRequest, LlmStreamCallExecuteParams, llm_call, + llm_call_execute, llm_stream_call_execute, }; use crate::api::runtime::{ - LlmExecutionNextFn, NemoRelayContextState, create_scope_stack, global_context, - set_thread_scope_stack, + LlmExecutionNextFn, LlmJsonStream, LlmStreamExecutionNextFn, NemoRelayContextState, + create_scope_stack, global_context, set_thread_scope_stack, }; use crate::api::scope::{ EmitMarkEventParams, PopScopeParams, PushScopeParams, ScopeType, event, pop_scope, push_scope, @@ -23,12 +24,13 @@ use crate::api::subscriber::{deregister_subscriber, register_subscriber}; use crate::api::tool::{ToolCallEndParams, ToolCallParams, tool_call, tool_call_end}; use crate::codec::openai_chat::OpenAIChatCodec; use crate::codec::openai_responses::OpenAIResponsesCodec; -use crate::codec::traits::LlmResponseCodec; +use crate::codec::traits::{LlmCodec, LlmResponseCodec}; use crate::plugin::{ PluginComponentSpec, PluginConfig, PluginError, PluginRegistrationContext, clear_plugin_configuration, ensure_builtin_plugins_registered, initialize_plugins, list_plugin_kinds, rollback_registrations, validate_plugin_config, }; +use futures::StreamExt; use nemo_relay::observability::atif::{AtifAgentInfo, AtifExporter}; use nemo_relay::observability::atof::{AtofExporter, AtofExporterConfig}; use nemo_relay::observability::openinference::OpenInferenceSubscriber; @@ -128,12 +130,582 @@ fn builtin_registry_includes_pii_redaction_component() { fn builtin_backend_config_default_matches_documented_action_default() { let config = BuiltinBackendConfig::default(); + assert!(config.preset.is_none()); assert_eq!(config.action, "remove"); + assert_eq!(config.custom_mark_payload_policy, "preserve"); assert!(config.target_paths.is_empty()); assert!(config.pattern.is_none()); assert!(config.detector.is_none()); } +#[test] +fn trajectory_preset_validates_without_an_action_and_rejects_matcher_fields() { + let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); + reset_runtime(); + let valid = validate_plugin_config(&plugin_config(json!({ + "codec": "openai_chat", + "profiles": [{ + "mode": "builtin", + "builtin": { + "preset": "trajectory_context", + "custom_mark_payload_policy": "redact_all_leaves" + } + }] + }))); + assert!(!valid.has_errors(), "{:#?}", valid.diagnostics); + + let invalid = validate_plugin_config(&plugin_config(json!({ + "codec": "openai_chat", + "profiles": [{ + "mode": "builtin", + "builtin": { + "preset": "trajectory_context", + "action": "redact", + "detector": "email" + } + }] + }))); + assert!(invalid.has_errors()); + assert!( + invalid.diagnostics.iter().any(|diagnostic| { + diagnostic.field.as_deref() == Some("profiles[0].builtin.action") + }) + ); + assert!( + invalid.diagnostics.iter().any(|diagnostic| { + diagnostic.field.as_deref() == Some("profiles[0].builtin.detector") + }) + ); + + let policy_without_preset = validate_plugin_config(&plugin_config(json!({ + "codec": "openai_chat", + "profiles": [{ + "mode": "builtin", + "builtin": {"custom_mark_payload_policy": "preserve"} + }] + }))); + assert!(policy_without_preset.diagnostics.iter().any(|diagnostic| { + diagnostic.field.as_deref() == Some("profiles[0].builtin.custom_mark_payload_policy") + })); +} + +fn trajectory_backend(codec: Option<&str>, policy: &str) -> crate::builtin::CompiledBuiltinBackend { + crate::builtin::CompiledBuiltinBackend::new( + BuiltinBackendConfig { + preset: Some("trajectory_context".into()), + custom_mark_payload_policy: policy.into(), + ..BuiltinBackendConfig::default() + }, + codec.map(str::to_string), + ) + .unwrap() +} + +#[test] +fn trajectory_preset_redacts_chat_content_without_erasing_request_structure() { + let callback = crate::builtin::llm_sanitize_request_callback(trajectory_backend( + Some("openai_chat"), + "preserve", + )); + let request = callback(LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "model": "claude-sonnet-4-6", + "messages": [ + {"role": "system", "content": "private system prompt"}, + {"role": "user", "content": [ + {"type": "text", "text": "private user prompt"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,secret", "detail": "high"}} + ]}, + {"role": "assistant", "tool_calls": [{ + "id": "call_1", + "type": "function", + "function": {"name": "search", "arguments": "{\"query\":\"private query\",\"limit\":5}"} + }]}, + {"role": "tool", "tool_call_id": "call_1", "content": "private result"} + ], + "tools": [{"type": "function", "function": { + "name": "search", + "description": "private description", + "parameters": {"type": "object", "properties": { + "query": {"type": "string", "description": "private schema text", "default": "private default"} + }, "required": ["query"]} + }}], + "temperature": 0.2, + "stop": ["private stop sequence"], + "participant": {"name": "Alice Example", "username": "alice"}, + "person_name": "Alice Example" + }), + }); + + assert_eq!(request.content["model"], "claude-sonnet-4-6"); + assert_eq!(request.content["temperature"], 0.2); + assert_eq!(request.content["stop"][0], "[REDACTED]"); + assert_eq!(request.content["participant"]["name"], "[REDACTED]"); + assert_eq!(request.content["participant"]["username"], "[REDACTED]"); + assert_eq!(request.content["person_name"], "[REDACTED]"); + assert_eq!(request.content["messages"][0]["role"], "system"); + assert_eq!(request.content["messages"][0]["content"], "[REDACTED]"); + assert_eq!(request.content["messages"][1]["content"][0]["type"], "text"); + assert_eq!( + request.content["messages"][1]["content"][0]["text"], + "[REDACTED]" + ); + assert_eq!( + request.content["messages"][1]["content"][1]["image_url"]["url"], + "[REDACTED]" + ); + assert_eq!( + request.content["messages"][2]["tool_calls"][0]["id"], + "call_1" + ); + assert_eq!( + request.content["messages"][2]["tool_calls"][0]["function"]["name"], + "search" + ); + let arguments: Json = serde_json::from_str( + request.content["messages"][2]["tool_calls"][0]["function"]["arguments"] + .as_str() + .unwrap(), + ) + .unwrap(); + assert_eq!(arguments, json!({"query": "[REDACTED]", "limit": 0})); + assert_eq!(request.content["messages"][3]["tool_call_id"], "call_1"); + assert_eq!(request.content["messages"][3]["content"], "[REDACTED]"); + assert_eq!(request.content["tools"][0]["function"]["name"], "search"); + assert_eq!( + request.content["tools"][0]["function"]["description"], + "[REDACTED]" + ); + assert_eq!( + request.content["tools"][0]["function"]["parameters"]["required"][0], + "query" + ); +} + +#[test] +fn trajectory_preset_preserves_response_analytics_and_redacts_response_content() { + let callback = crate::builtin::llm_sanitize_response_callback(trajectory_backend( + Some("openai_chat"), + "preserve", + )); + let sanitized = callback(json!({ + "id": "chatcmpl_1", + "model": "claude-opus-4-6", + "choices": [{"index": 0, "finish_reason": "tool_calls", "message": { + "role": "assistant", + "content": "private answer", + "tool_calls": [{"id": "call_1", "type": "function", "function": { + "name": "terminal", "arguments": "{\"command\":\"cat secret.txt\"}" + }}] + }, "logprobs": {"content": [{"token": "secret", "logprob": -0.5}]}}], + "usage": {"prompt_tokens": 20, "completion_tokens": 5, "total_tokens": 25}, + "cost": {"total": 1.25} + })); + + assert_eq!(sanitized["id"], "chatcmpl_1"); + assert_eq!(sanitized["model"], "claude-opus-4-6"); + assert_eq!(sanitized["choices"][0]["finish_reason"], "tool_calls"); + assert_eq!(sanitized["choices"][0]["message"]["role"], "assistant"); + assert_eq!(sanitized["choices"][0]["message"]["content"], "[REDACTED]"); + assert_eq!( + sanitized["choices"][0]["message"]["tool_calls"][0]["id"], + "call_1" + ); + assert_eq!( + sanitized["choices"][0]["message"]["tool_calls"][0]["function"]["name"], + "terminal" + ); + let arguments: Json = serde_json::from_str( + sanitized["choices"][0]["message"]["tool_calls"][0]["function"]["arguments"] + .as_str() + .unwrap(), + ) + .unwrap(); + assert_eq!(arguments, json!({"command": "[REDACTED]"})); + assert_eq!(sanitized["usage"]["total_tokens"], 25); + assert_eq!(sanitized["cost"]["total"], 1.25); +} + +#[test] +fn trajectory_preset_covers_responses_and_anthropic_provider_shapes() { + let responses_request = crate::builtin::llm_sanitize_request_callback(trajectory_backend( + Some("openai_responses"), + "preserve", + ))(LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "model": "gpt-5", + "input": [{"role": "user", "content": [{"type": "input_text", "text": "private input"}]}], + "reasoning": {"effort": "high", "summary": "private reasoning"}, + "max_output_tokens": 100 + }), + }); + assert_eq!(responses_request.content["model"], "gpt-5"); + assert_eq!(responses_request.content["input"][0]["role"], "user"); + assert_eq!( + responses_request.content["input"][0]["content"][0]["text"], + "[REDACTED]" + ); + assert_eq!(responses_request.content["max_output_tokens"], 100); + + let responses_response = crate::builtin::llm_sanitize_response_callback(trajectory_backend( + Some("openai_responses"), + "preserve", + ))(json!({ + "id": "resp_1", + "model": "gpt-5", + "status": "completed", + "output": [{"id": "msg_1", "type": "message", "role": "assistant", "content": [ + {"type": "output_text", "text": "private output"} + ]}], + "usage": {"input_tokens": 10, "output_tokens": 4, "total_tokens": 14} + })); + assert_eq!(responses_response["id"], "resp_1"); + assert_eq!(responses_response["status"], "completed"); + assert_eq!(responses_response["output"][0]["id"], "msg_1"); + assert_eq!( + responses_response["output"][0]["content"][0]["text"], + "[REDACTED]" + ); + assert_eq!(responses_response["usage"]["total_tokens"], 14); + + let anthropic_request = crate::builtin::llm_sanitize_request_callback(trajectory_backend( + Some("anthropic_messages"), + "preserve", + ))(LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "model": "claude-sonnet-4-6", + "system": "private system", + "messages": [{"role": "user", "content": [ + {"type": "text", "text": "private user"}, + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "private-base64"}} + ]}], + "max_tokens": 128 + }), + }); + assert_eq!(anthropic_request.content["model"], "claude-sonnet-4-6"); + assert_eq!(anthropic_request.content["system"], "[REDACTED]"); + assert_eq!( + anthropic_request.content["messages"][0]["content"][0]["text"], "[REDACTED]", + "{}", + anthropic_request.content + ); + assert_eq!( + anthropic_request.content["messages"][0]["content"][1]["source"]["data"], + "[REDACTED]" + ); + assert_eq!(anthropic_request.content["max_tokens"], 128); + + let anthropic_response = crate::builtin::llm_sanitize_response_callback(trajectory_backend( + Some("anthropic_messages"), + "preserve", + ))(json!({ + "id": "msg_1", + "model": "claude-sonnet-4-6", + "type": "message", + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "private chain of thought", "signature": "private-signature"}, + {"type": "text", "text": "private answer"} + ], + "stop_reason": "end_turn", + "usage": {"input_tokens": 12, "output_tokens": 6, "cache_read_input_tokens": 8} + })); + assert_eq!(anthropic_response["id"], "msg_1"); + assert_eq!(anthropic_response["role"], "assistant"); + assert_eq!(anthropic_response["content"][0]["type"], "thinking"); + assert_eq!(anthropic_response["content"][0]["thinking"], "[REDACTED]"); + assert_eq!(anthropic_response["content"][1]["text"], "[REDACTED]"); + assert_eq!(anthropic_response["usage"]["input_tokens"], 12); + assert_eq!(anthropic_response["usage"]["cache_read_input_tokens"], 8); +} + +#[test] +fn trajectory_preset_redacts_known_marks_and_nested_scope_content() { + let callback = crate::builtin::event_sanitize_callback(trajectory_backend(None, "preserve")); + let chunk = Event::Mark(MarkEvent::new( + BaseEvent::builder().name("llm.chunk").build(), + Some(EventCategory::custom()), + Some(CategoryProfile::builder().subtype("llm.chunk").build()), + )); + let sanitized = callback( + &chunk, + EventSanitizeFields { + data: Some(json!({ + "chunk_index": 2, + "event_type": "content_block_delta", + "delta": {"type": "text_delta", "text": "private delta"} + })), + category_profile: chunk.category_profile().cloned(), + metadata: None, + }, + ); + assert_eq!(sanitized.data.as_ref().unwrap()["chunk_index"], 2); + assert_eq!( + sanitized.data.as_ref().unwrap()["event_type"], + "content_block_delta" + ); + assert_eq!( + sanitized.data.as_ref().unwrap()["delta"]["text"], + "[REDACTED]" + ); + + let optimization = Event::Mark(MarkEvent::new( + BaseEvent::builder() + .name("nemo_relay.llm.optimization") + .build(), + Some(EventCategory::custom()), + Some( + CategoryProfile::builder() + .subtype("nemo_relay.llm.optimization") + .build(), + ), + )); + let sanitized = callback( + &optimization, + EventSanitizeFields { + data: Some(json!({ + "producer": "neutral.router", + "kind": "model_routing", + "applied": true, + "model_transition": { + "baseline": {"model": "claude-opus-4-6"}, + "effective": {"model": "claude-sonnet-4-6"} + }, + "token_impact": {"saved": {"prompt_tokens": 40, "total_tokens": 40}}, + "payload": {"private_excerpt": "private content"} + })), + category_profile: optimization.category_profile().cloned(), + metadata: None, + }, + ); + assert_eq!( + sanitized.data.as_ref().unwrap()["producer"], + "neutral.router" + ); + assert_eq!( + sanitized.data.as_ref().unwrap()["model_transition"]["baseline"]["model"], + "claude-opus-4-6" + ); + assert_eq!( + sanitized.data.as_ref().unwrap()["token_impact"]["saved"]["total_tokens"], + 40 + ); + assert_eq!( + sanitized.data.as_ref().unwrap()["payload"]["private_excerpt"], + "[REDACTED]" + ); + + let nested_agent = Event::Scope(ScopeEvent::new( + BaseEvent::builder().name("worker-agent").build(), + ScopeCategory::Start, + Vec::new(), + EventCategory::agent(), + None, + )); + let sanitized = callback( + &nested_agent, + EventSanitizeFields { + data: Some(json!({ + "request_id": "request-1", + "instruction": "private delegated task", + "history": [{"role": "user", "content": "private history"}] + })), + category_profile: None, + metadata: Some(json!({"parent_scope_id": "scope-1", "note": "private note"})), + }, + ); + assert_eq!(sanitized.data.as_ref().unwrap()["request_id"], "request-1"); + assert_eq!( + sanitized.data.as_ref().unwrap()["instruction"], + "[REDACTED]" + ); + assert_eq!( + sanitized.data.as_ref().unwrap()["history"][0]["role"], + "user" + ); + assert_eq!( + sanitized.data.as_ref().unwrap()["history"][0]["content"], + "[REDACTED]" + ); + assert_eq!( + sanitized.metadata.as_ref().unwrap()["parent_scope_id"], + "scope-1" + ); + assert_eq!(sanitized.metadata.as_ref().unwrap()["note"], "[REDACTED]"); +} + +#[test] +fn trajectory_custom_mark_policy_is_explicit_and_shape_preserving() { + let event = Event::Mark(MarkEvent::new( + BaseEvent::builder().name("neutral.plugin.evidence").build(), + Some(EventCategory::custom()), + Some(CategoryProfile { + subtype: Some("neutral.plugin".into()), + extra: BTreeMap::from([("opaque".into(), json!({"label": "private"}))]), + ..CategoryProfile::default() + }), + )); + let fields = EventSanitizeFields { + data: Some(json!({"text": "private", "score": 0.75, "nested": [true, null]})), + category_profile: event.category_profile().cloned(), + metadata: Some(json!({"owner": "private owner"})), + }; + + let preserve = crate::builtin::event_sanitize_callback(trajectory_backend(None, "preserve")); + assert_eq!(preserve(&event, fields.clone()), fields); + + let redact = + crate::builtin::event_sanitize_callback(trajectory_backend(None, "redact_all_leaves")); + let sanitized = redact(&event, fields); + assert_eq!( + sanitized.data.unwrap(), + json!({ + "text": "[REDACTED]", "score": 0, "nested": [false, null] + }) + ); + assert_eq!(sanitized.metadata.unwrap(), json!({"owner": "[REDACTED]"})); + let profile = sanitized.category_profile.unwrap(); + assert_eq!(profile.subtype.as_deref(), Some("neutral.plugin")); + assert_eq!(profile.extra["opaque"]["label"], "[REDACTED]"); +} + +#[test] +fn trajectory_profile_preserves_typed_llm_accounting_while_redacting_annotations() { + let callback = crate::builtin::event_sanitize_callback(trajectory_backend(None, "preserve")); + let annotated_response: nemo_relay::codec::response::AnnotatedLlmResponse = + serde_json::from_value(json!({ + "model": "claude-sonnet-4-6", + "message": "private answer", + "finish_reason": "complete", + "usage": { + "prompt_tokens": 100, + "completion_tokens": 20, + "total_tokens": 120, + "cost": { + "total": 0.42, + "currency": "USD", + "source": "provider_reported" + } + }, + "optimization_summary": { + "schema_version": "1", + "calculation_version": "1", + "status": "complete", + "baseline_model": {"model": "claude-opus-4-6"}, + "effective_model": {"model": "claude-sonnet-4-6"}, + "effective_usage": {"prompt_tokens": 100, "completion_tokens": 20, "total_tokens": 120}, + "baseline_usage": {"prompt_tokens": 140, "completion_tokens": 20, "total_tokens": 160}, + "tokens_saved": {"prompt_tokens": 40, "total_tokens": 40}, + "estimated_cost_saved": 0.8, + "currency": "USD", + "contributions": [{ + "producer": "neutral.optimizer", + "kind": "input_compression", + "applied": true, + "token_impact": { + "saved": {"prompt_tokens": 40, "total_tokens": 40}, + "quality": "estimated", + "estimation_method": "neutral_counter" + }, + "payload_schema": {"name": "neutral.evidence", "version": "1"}, + "payload": {"private_excerpt": "private tool output", "strategy": "head_tail"} + }] + } + })) + .unwrap(); + let event = Event::Scope(ScopeEvent::new( + BaseEvent::builder().name("llm").build(), + ScopeCategory::End, + Vec::new(), + EventCategory::llm(), + None, + )); + let sanitized = callback( + &event, + EventSanitizeFields { + data: Some(json!({"already": "sanitized by the response callback"})), + category_profile: Some( + CategoryProfile::builder() + .model_name("claude-sonnet-4-6") + .annotated_response(Arc::new(annotated_response)) + .build(), + ), + metadata: None, + }, + ); + + let profile = sanitized.category_profile.unwrap(); + assert_eq!(profile.model_name.as_deref(), Some("claude-sonnet-4-6")); + let response = profile.annotated_response.unwrap(); + assert_eq!(response.response_text(), Some("[REDACTED]")); + assert_eq!(response.usage.as_ref().unwrap().total_tokens, Some(120)); + assert_eq!( + response + .usage + .as_ref() + .unwrap() + .cost + .as_ref() + .unwrap() + .total, + Some(0.42) + ); + let summary = response.optimization_summary.as_ref().unwrap(); + assert_eq!(summary.tokens_saved.prompt_tokens, Some(40)); + assert_eq!(summary.estimated_cost_saved, Some(0.8)); + assert_eq!(summary.contributions[0].producer, "neutral.optimizer"); + assert_eq!( + summary.contributions[0].payload.as_ref().unwrap()["private_excerpt"], + "[REDACTED]" + ); + assert_eq!( + summary.contributions[0].payload.as_ref().unwrap()["strategy"], + "[REDACTED]" + ); + assert_eq!( + sanitized.data.unwrap()["already"], + "sanitized by the response callback", + "specialized LLM data must not be processed twice" + ); +} + +#[test] +fn preserved_custom_marks_remain_eligible_for_a_later_email_profile() { + let event = Event::Mark(MarkEvent::new( + BaseEvent::builder().name("neutral.plugin.evidence").build(), + Some(EventCategory::custom()), + Some(CategoryProfile::builder().subtype("neutral.plugin").build()), + )); + let fields = EventSanitizeFields { + data: Some(json!({"owner": "alice@example.com", "score": 0.9})), + category_profile: event.category_profile().cloned(), + metadata: Some(json!({"contact": "bob@example.com"})), + }; + let trajectory = crate::builtin::event_sanitize_callback(trajectory_backend(None, "preserve")); + let email = crate::builtin::event_sanitize_callback( + crate::builtin::CompiledBuiltinBackend::new( + BuiltinBackendConfig { + action: "redact".into(), + detector: Some("email".into()), + ..BuiltinBackendConfig::default() + }, + None, + ) + .unwrap(), + ); + + let sanitized = email(&event, trajectory(&event, fields)); + assert_eq!(sanitized.data.as_ref().unwrap()["owner"], "[REDACTED]"); + assert_eq!(sanitized.data.as_ref().unwrap()["score"], 0.9); + assert_eq!( + sanitized.metadata.as_ref().unwrap()["contact"], + "[REDACTED]" + ); +} + #[test] fn pii_redaction_defaults_enable_mark_sanitization() { let config = PiiRedactionConfig::default(); @@ -183,6 +755,18 @@ fn typed_profile_config_serializes_without_conflicting_legacy_defaults() { assert!(!report.has_errors(), "{:?}", report.diagnostics); } +#[test] +fn typed_trajectory_preset_omits_the_legacy_default_action() { + let config = BuiltinBackendConfig { + preset: Some("trajectory_context".into()), + ..BuiltinBackendConfig::default() + }; + let serialized = serde_json::to_value(config).unwrap(); + assert_eq!(serialized["preset"], "trajectory_context"); + assert!(serialized.get("action").is_none()); + assert!(serialized.get("custom_mark_payload_policy").is_none()); +} + #[test] fn profile_array_executes_every_profile_in_stable_array_order() { let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); @@ -980,22 +1564,28 @@ fn mark_false_preserves_mark_fields() { } #[test] -fn sanitized_pii_never_reaches_subscribers_or_exporters() { +fn sanitized_trajectory_content_never_reaches_subscribers_or_exporters() { let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); reset_runtime(); setup_isolated_thread(); futures::executor::block_on(initialize_plugins(plugin_config(json!({ - "mode": "builtin", "codec": "openai_chat", - "input": true, - "output": true, - "tool_input": false, - "tool_output": false, - "builtin": { - "action": "redact", - "detector": "email" - } + "profiles": [ + { + "mode": "builtin", + "priority": 80, + "builtin": { + "preset": "trajectory_context", + "custom_mark_payload_policy": "redact_all_leaves" + } + }, + { + "mode": "builtin", + "priority": 90, + "builtin": {"action": "redact", "detector": "email"} + } + ] })))) .unwrap(); @@ -1038,11 +1628,12 @@ fn sanitized_pii_never_reaches_subscribers_or_exporters() { .unwrap(); let raw_pii = "person@example.com"; + let raw_context = "private trajectory context canary"; let agent = push_scope( PushScopeParams::builder() .name("hermes-agent") .scope_type(ScopeType::Agent) - .input(json!({"prompt": raw_pii})) + .input(json!({"prompt": raw_context, "request_id": "request-1"})) .metadata(json!({"session_owner": raw_pii})) .build(), ) @@ -1050,7 +1641,7 @@ fn sanitized_pii_never_reaches_subscribers_or_exporters() { event( EmitMarkEventParams::builder() .name("hermes.checkpoint") - .data(json!({"email": raw_pii})) + .data(json!({"content": raw_context, "email": raw_pii, "score": 0.95})) .metadata(json!({"reviewer": raw_pii})) .build(), ) @@ -1058,7 +1649,7 @@ fn sanitized_pii_never_reaches_subscribers_or_exporters() { pop_scope( PopScopeParams::builder() .handle_uuid(&agent.uuid) - .output(json!({"answer": raw_pii})) + .output(json!({"answer": raw_context})) .metadata(json!({"approver": raw_pii})) .build(), ) @@ -1085,8 +1676,28 @@ fn sanitized_pii_never_reaches_subscribers_or_exporters() { !output.contains(raw_pii), "raw PII leaked through {surface}: {output}" ); + assert!( + !output.contains(raw_context), + "trajectory context leaked through {surface}: {output}" + ); } + let captured = captured_events_snapshot(&captured); + let agent_start = captured + .iter() + .find(|event| { + event.name() == "hermes-agent" && event.scope_category() == Some(ScopeCategory::Start) + }) + .unwrap(); + assert_eq!(agent_start.data().unwrap()["request_id"], "request-1"); + assert_eq!(agent_start.data().unwrap()["prompt"], "[REDACTED]"); + let custom_mark = captured + .iter() + .find(|event| event.name() == "hermes.checkpoint") + .unwrap(); + assert_eq!(custom_mark.data().unwrap()["content"], "[REDACTED]"); + assert_eq!(custom_mark.data().unwrap()["score"], 0); + deregister_subscriber("pii-regression-subscriber").unwrap(); atof.deregister("pii-regression-atof").unwrap(); deregister_subscriber("pii-regression-atif").unwrap(); @@ -1097,6 +1708,100 @@ fn sanitized_pii_never_reaches_subscribers_or_exporters() { clear_plugin_configuration().unwrap(); } +#[tokio::test] +async fn trajectory_preset_sanitizes_stream_finalization_without_changing_client_chunks() { + let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); + reset_runtime(); + setup_isolated_thread(); + + initialize_plugins(plugin_config(json!({ + "codec": "openai_chat", + "profiles": [{ + "mode": "builtin", + "builtin": {"preset": "trajectory_context"} + }] + }))) + .await + .unwrap(); + + let events = capture_events("pii-trajectory-stream"); + let raw_delta = "private streaming delta"; + let provider: LlmStreamExecutionNextFn = Arc::new(move |_request| { + Box::pin(async move { + Ok(LlmJsonStream::new(futures::stream::iter(vec![Ok(json!({ + "id": "chatcmpl-stream", + "object": "chat.completion.chunk", + "model": "gpt-4o-mini", + "choices": [{"index": 0, "delta": {"content": raw_delta}, "finish_reason": null}] + }))]))) + }) + }); + let request_codec: Arc = Arc::new(OpenAIChatCodec); + let response_codec: Arc = Arc::new(OpenAIChatCodec); + let mut stream = llm_stream_call_execute( + LlmStreamCallExecuteParams::builder() + .name("openai") + .request(LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "private stream prompt"}] + }), + }) + .func(provider) + .collector(Box::new(|_| Ok(()))) + .finalizer(Box::new(|| { + json!({ + "id": "chatcmpl-stream", + "model": "gpt-4o-mini", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "private final answer"}, + "finish_reason": "stop" + }], + "usage": {"prompt_tokens": 8, "completion_tokens": 3, "total_tokens": 11} + }) + })) + .codec(request_codec) + .response_codec(response_codec) + .build(), + ) + .await + .unwrap(); + + let chunk = stream.next().await.unwrap().unwrap(); + assert_eq!(chunk["choices"][0]["delta"]["content"], raw_delta); + assert!(stream.next().await.is_none()); + + let captured = captured_events_snapshot(&events); + let start = captured + .iter() + .find(|event| event.scope_category() == Some(ScopeCategory::Start)) + .unwrap(); + assert_eq!( + start.input().unwrap()["content"]["messages"][0]["content"], + "[REDACTED]" + ); + let chunk_mark = captured + .iter() + .find(|event| event.name() == "llm.chunk") + .unwrap(); + assert_eq!(chunk_mark.data().unwrap()["chunk_index"], 0); + assert!(!chunk_mark.to_json_string().unwrap().contains(raw_delta)); + let end = captured + .iter() + .find(|event| event.scope_category() == Some(ScopeCategory::End)) + .unwrap(); + assert_eq!( + end.output().unwrap()["choices"][0]["message"]["content"], + "[REDACTED]" + ); + assert_eq!(end.output().unwrap()["usage"]["total_tokens"], 11); + + deregister_subscriber("pii-trajectory-stream").unwrap(); + clear_plugin_configuration().unwrap(); +} + #[test] fn builtin_backend_sanitizes_tool_start_and_end_payloads_with_preorder_targets() { let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); diff --git a/crates/python/src/convert.rs b/crates/python/src/convert.rs index 273397bdb..e0354a27c 100644 --- a/crates/python/src/convert.rs +++ b/crates/python/src/convert.rs @@ -7,12 +7,85 @@ //! the required/optional × to-json/from-json matrix used throughout the PyO3 //! binding layer. +use std::collections::HashSet; + use chrono::{DateTime, Utc}; -use pyo3::prelude::*; +use pyo3::types::{ + PyByteArray, PyBytes, PyDict, PyFrozenSet, PyMapping, PySequence, PySet, PyString, +}; +use pyo3::{intern, prelude::*}; use serde_json::Value as Json; +fn validate_acyclic( + value: &Bound<'_, PyAny>, + active_containers: &mut HashSet, +) -> PyResult<()> { + if value.is_instance_of::() + || value.is_instance_of::() + || value.is_instance_of::() + { + return Ok(()); + } + + let dataclass_fields = value + .getattr_opt(intern!(value.py(), "__dataclass_fields__")) + .ok() + .flatten(); + let is_container = value.cast::().is_ok() + || value.cast::().is_ok() + || value.cast::().is_ok() + || value.cast::().is_ok() + || dataclass_fields.is_some(); + if !is_container { + return Ok(()); + } + + let identity = value.as_ptr() as usize; + if !active_containers.insert(identity) { + return Err(pyo3::exceptions::PyValueError::new_err( + "Failed to convert to JSON: circular reference detected", + )); + } + + let result = if let Ok(set) = value.cast::() { + set.iter() + .try_for_each(|child| validate_acyclic(&child, active_containers)) + } else if let Ok(set) = value.cast::() { + set.iter() + .try_for_each(|child| validate_acyclic(&child, active_containers)) + } else if let Ok(sequence) = value.cast::() { + (0..sequence.len()?).try_for_each(|index| { + let child = sequence.get_item(index)?; + validate_acyclic(&child, active_containers) + }) + } else if let Ok(mapping) = value.cast::() { + let keys = mapping.keys()?; + let values = mapping.values()?; + keys.iter() + .chain(values.iter()) + .try_for_each(|child| validate_acyclic(&child, active_containers)) + } else if let Some(fields) = dataclass_fields { + let fields = fields.cast::()?.keys(); + let attributes = value + .getattr(intern!(value.py(), "__dict__"))? + .cast_into::()?; + fields.iter().try_for_each(|field| { + if let Some(child) = attributes.get_item(&field)? { + validate_acyclic(&child, active_containers)?; + } + Ok(()) + }) + } else { + Ok(()) + }; + + active_containers.remove(&identity); + result +} + /// Convert a Python object to serde_json::Value via pythonize. pub fn py_to_json(obj: &Bound<'_, PyAny>) -> PyResult { + validate_acyclic(obj, &mut HashSet::new())?; pythonize::depythonize(obj).map_err(|e| { PyErr::new::(format!("Failed to convert to JSON: {e}")) }) diff --git a/crates/python/src/py_api/mod.rs b/crates/python/src/py_api/mod.rs index 5abbdd933..bc026d76e 100644 --- a/crates/python/src/py_api/mod.rs +++ b/crates/python/src/py_api/mod.rs @@ -1328,6 +1328,10 @@ fn deregister_subscriber(name: &str) -> PyResult { } /// Wait for subscriber callbacks queued before this call to finish. +/// +/// Call this function outside native subscriber callbacks. A re-entrant call returns without +/// waiting to avoid blocking the dispatcher, so callbacks later in the same dispatch snapshot can +/// still run. #[pyfunction] fn flush_subscribers(py: Python<'_>) -> PyResult<()> { py.detach(core_subscriber_api::flush_subscribers) diff --git a/crates/python/src/py_types/codecs.rs b/crates/python/src/py_types/codecs.rs index 636f82fb1..73efdf5a8 100644 --- a/crates/python/src/py_types/codecs.rs +++ b/crates/python/src/py_types/codecs.rs @@ -6,12 +6,14 @@ use serde::de::DeserializeOwned; use super::core::PyLLMRequest; use super::{ - AnnotatedLLMRequest, AnnotatedLLMResponse, Arc, Bound, GenerationParams, LlmCodec, - LlmResponseCodec, Message, PyAny, PyResult, Python, ToolChoice, ToolDefinition, json_to_py, - py_to_json, to_python_json_value, + AnnotatedLLMRequest, AnnotatedLLMResponse, ApiSpecificRequest, Arc, Bound, GenerationParams, + LlmCodec, LlmResponseCodec, Message, MessageContent, PyAny, PyResult, Python, ToolChoice, + ToolDefinition, json_to_py, py_to_json, to_python_json_value, }; #[cfg(test)] use super::{ + FORCE_ANNOTATED_REQUEST_API_SPECIFIC_SERIALIZATION_ERROR, + FORCE_ANNOTATED_REQUEST_INSTRUCTIONS_SERIALIZATION_ERROR, FORCE_ANNOTATED_REQUEST_MESSAGES_SERIALIZATION_ERROR, FORCE_ANNOTATED_REQUEST_PARAMS_SERIALIZATION_ERROR, FORCE_ANNOTATED_REQUEST_TOOL_CHOICE_SERIALIZATION_ERROR, @@ -31,15 +33,18 @@ use nemo_relay::codec::response::FinishReason; /// A structured view of an LLM request produced by a Codec. /// /// Provides typed access to conversation messages, model name, generation -/// parameters, tool definitions, tool choice, and extensible extra fields. +/// parameters, tool definitions, tool choice, tagged provider-specific fields, +/// and extensible unknown top-level fields. /// /// Properties: /// messages (list): Parsed conversation messages (list of dicts with a ``role`` key). +/// instructions (str | list | None): Provider-level system instructions. /// model (str | None): Model identifier (e.g., ``"gpt-4"``). /// params (dict | None): Normalized generation parameters. /// tools (list | None): Tool definitions (function schemas). /// tool_choice (Any | None): Tool choice control. -/// extra (dict): Provider-specific extra fields. +/// api_specific (dict | None): Tagged provider-specific request fields. +/// extra (dict): Unknown future top-level fields. /// /// Helper methods: /// system_prompt() -> str | None: Text of the first system message. @@ -79,19 +84,27 @@ impl PyAnnotatedLLMRequest { /// /// Args: /// messages: A list of message dicts, each with a ``role`` key. + /// instructions: Optional provider-level instructions. /// model: Optional model identifier. /// params: Optional generation parameters dict. /// tools: Optional list of tool definition dicts. /// tool_choice: Optional tool choice control. - /// extra: Optional dict of provider-specific extra fields. + /// api_specific: Optional tagged provider-specific request fields. + /// extra: Optional dict of unknown future top-level fields. #[new] - #[pyo3(signature = (messages, *, model=None, params=None, tools=None, tool_choice=None, extra=None))] + #[pyo3(signature = (messages, *, instructions=None, model=None, params=None, tools=None, tool_choice=None, api_specific=None, extra=None))] + #[allow( + clippy::too_many_arguments, + reason = "the Python constructor mirrors the public annotation fields as keywords" + )] pub(crate) fn new( messages: &Bound<'_, PyAny>, + instructions: Option<&Bound<'_, PyAny>>, model: Option, params: Option<&Bound<'_, PyAny>>, tools: Option<&Bound<'_, PyAny>>, tool_choice: Option<&Bound<'_, PyAny>>, + api_specific: Option<&Bound<'_, PyAny>>, extra: Option<&Bound<'_, PyAny>>, ) -> PyResult { let msgs: Vec = pythonize::depythonize(messages).map_err(|e| { @@ -105,6 +118,16 @@ impl PyAnnotatedLLMRequest { })?), _ => None, }; + let request_instructions: Option = match instructions { + Some(value) if !value.is_none() => { + Some(pythonize::depythonize(value).map_err(|e| { + pyo3::exceptions::PyValueError::new_err(format!( + "invalid instructions: expected a string or content-part list: {e}" + )) + })?) + } + _ => None, + }; let tool_defs: Option> = match tools { Some(t) if !t.is_none() => Some(pythonize::depythonize(t).map_err(|e| { pyo3::exceptions::PyValueError::new_err(format!("invalid tools: {e}")) @@ -119,6 +142,16 @@ impl PyAnnotatedLLMRequest { } _ => None, }; + let provider_fields: Option = match api_specific { + Some(value) if !value.is_none() => { + Some(pythonize::depythonize(value).map_err(|e| { + pyo3::exceptions::PyValueError::new_err(format!( + "invalid api_specific: expected a tagged request object: {e}" + )) + })?) + } + _ => None, + }; let extra_map: serde_json::Map = match extra { Some(e) if !e.is_none() => pythonize::depythonize(e).map_err(|e| { pyo3::exceptions::PyValueError::new_err(format!("invalid extra: {e}")) @@ -128,6 +161,7 @@ impl PyAnnotatedLLMRequest { Ok(Self { inner: AnnotatedLLMRequest { messages: msgs, + instructions: request_instructions, model, params: gen_params, tools: tool_defs, @@ -145,6 +179,7 @@ impl PyAnnotatedLLMRequest { max_tool_calls: None, top_logprobs: None, stream: None, + api_specific: provider_fields, extra: extra_map, }, }) @@ -171,6 +206,36 @@ impl PyAnnotatedLLMRequest { Ok(()) } + #[getter] + pub(crate) fn instructions(&self, py: Python<'_>) -> PyResult> { + match &self.inner.instructions { + Some(value) => { + let value = to_python_json_value( + value, + "serialization error", + #[cfg(test)] + FORCE_ANNOTATED_REQUEST_INSTRUCTIONS_SERIALIZATION_ERROR, + )?; + json_to_py(py, &value) + } + None => Ok(py.None()), + } + } + + #[setter] + pub(crate) fn set_instructions(&mut self, value: &Bound<'_, PyAny>) -> PyResult<()> { + if value.is_none() { + self.inner.instructions = None; + } else { + self.inner.instructions = Some(pythonize::depythonize(value).map_err(|e| { + pyo3::exceptions::PyValueError::new_err(format!( + "invalid instructions: expected a string or content-part list: {e}" + )) + })?); + } + Ok(()) + } + #[getter] pub(crate) fn model(&self) -> Option { self.inner.model.clone() @@ -395,6 +460,34 @@ impl PyAnnotatedLLMRequest { self.inner.stream = value; } + #[getter] + pub(crate) fn api_specific(&self, py: Python<'_>) -> PyResult> { + match &self.inner.api_specific { + Some(value) => { + let value = to_python_json_value( + value, + "serialization error", + #[cfg(test)] + FORCE_ANNOTATED_REQUEST_API_SPECIFIC_SERIALIZATION_ERROR, + )?; + json_to_py(py, &value) + } + None => Ok(py.None()), + } + } + + #[setter] + pub(crate) fn set_api_specific(&mut self, value: &Bound<'_, PyAny>) -> PyResult<()> { + if value.is_none() { + self.inner.api_specific = None; + } else { + self.inner.api_specific = Some(pythonize::depythonize(value).map_err(|e| { + pyo3::exceptions::PyValueError::new_err(format!("invalid api_specific: {e}")) + })?); + } + Ok(()) + } + #[getter] pub(crate) fn extra(&self, py: Python<'_>) -> PyResult> { let value = serde_json::Value::Object(self.inner.extra.clone()); diff --git a/crates/python/src/py_types/mod.rs b/crates/python/src/py_types/mod.rs index b713d630a..046b580ff 100644 --- a/crates/python/src/py_types/mod.rs +++ b/crates/python/src/py_types/mod.rs @@ -19,8 +19,8 @@ use nemo_relay::api::runtime::ScopeStackHandle; use nemo_relay::api::scope::{ScopeAttributes, ScopeHandle, ScopeType as CoreScopeType}; use nemo_relay::api::tool::{ToolAttributes, ToolHandle}; use nemo_relay::codec::request::{ - AnnotatedLlmRequest as AnnotatedLLMRequest, GenerationParams, Message, ToolChoice, - ToolDefinition, + AnnotatedLlmRequest as AnnotatedLLMRequest, ApiSpecificRequest, GenerationParams, Message, + MessageContent, ToolChoice, ToolDefinition, }; use nemo_relay::codec::response::AnnotatedLlmResponse as AnnotatedLLMResponse; use nemo_relay::codec::traits::{LlmCodec, LlmResponseCodec}; @@ -55,6 +55,10 @@ pub(crate) const FORCE_ANNOTATED_RESPONSE_USAGE_SERIALIZATION_ERROR: u64 = 1 << pub(crate) const FORCE_ANNOTATED_RESPONSE_API_SPECIFIC_SERIALIZATION_ERROR: u64 = 1 << 9; #[cfg(test)] pub(crate) const FORCE_ANNOTATED_RESPONSE_OPTIMIZATION_SUMMARY_SERIALIZATION_ERROR: u64 = 1 << 10; +#[cfg(test)] +pub(crate) const FORCE_ANNOTATED_REQUEST_INSTRUCTIONS_SERIALIZATION_ERROR: u64 = 1 << 11; +#[cfg(test)] +pub(crate) const FORCE_ANNOTATED_REQUEST_API_SPECIFIC_SERIALIZATION_ERROR: u64 = 1 << 12; #[cfg(test)] pub(crate) fn set_forced_serialization_mask_for_tests(mask: u64) { diff --git a/crates/python/src/py_types/observability.rs b/crates/python/src/py_types/observability.rs index 85a24bb23..f021f8e49 100644 --- a/crates/python/src/py_types/observability.rs +++ b/crates/python/src/py_types/observability.rs @@ -391,13 +391,19 @@ impl PyAtofExporter { .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string())) } - /// Flush the output file. + /// Outside a native subscriber callback, wait for queued subscriber delivery, then flush the + /// file sink or ask the stream sink to drain for up to its timeout. A re-entrant call does not + /// establish the delivery barrier. A stream timeout is logged and does not by itself return an + /// error. pub(crate) fn force_flush(&self, py: Python<'_>) -> PyResult<()> { py.detach(|| self.inner.force_flush()) .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string())) } - /// Shut down the exporter by flushing output. + /// Outside a native subscriber callback, wait for queued subscriber delivery, then flush the + /// file sink or ask the stream sink to drain and close up to its timeout. A re-entrant call + /// does not establish the delivery barrier. A stream timeout is logged and does not by itself + /// return an error. pub(crate) fn shutdown(&self, py: Python<'_>) -> PyResult<()> { py.detach(|| self.inner.shutdown()) .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string())) diff --git a/crates/python/tests/coverage/py_types_coverage_tests.rs b/crates/python/tests/coverage/py_types_coverage_tests.rs index 0a8ae2452..11d2fdd9d 100644 --- a/crates/python/tests/coverage/py_types_coverage_tests.rs +++ b/crates/python/tests/coverage/py_types_coverage_tests.rs @@ -572,6 +572,8 @@ fn test_stream_request_event_and_handle_wrappers_cover_remaining_methods() { assert_eq!(request.__repr__(), "LLMRequest(...)"); let annotated_request = AnnotatedLLMRequest { + instructions: None, + api_specific: None, messages: vec![ Message::System { content: MessageContent::Text("system".into()), @@ -1070,19 +1072,31 @@ fn test_annotated_llm_types_and_builtin_codecs_cover_mutators_and_codecs() { &json!({"type": "function", "function": {"name": "lookup"}}), ) .unwrap(); + let instructions = json_to_py(py, &json!("Initial policy")).unwrap(); + let api_specific = json_to_py(py, &json!({"api": "openai_chat", "seed": 7})).unwrap(); let extra = json_to_py(py, &json!({"provider": "test"})).unwrap(); let mut annotated = PyAnnotatedLLMRequest::new( messages.bind(py), + Some(instructions.bind(py)), Some("demo-model".into()), Some(params.bind(py)), Some(tools.bind(py)), Some(tool_choice.bind(py)), + Some(api_specific.bind(py)), Some(extra.bind(py)), ) .unwrap(); assert_eq!(annotated.model(), Some("demo-model".into())); - assert_eq!(annotated.system_prompt(), Some("You are terse.".into())); + assert_eq!( + py_to_json(annotated.instructions(py).unwrap().bind(py)).unwrap(), + json!("Initial policy") + ); + assert_eq!( + py_to_json(annotated.api_specific(py).unwrap().bind(py)).unwrap(), + json!({"api": "openai_chat", "seed": 7}) + ); + assert_eq!(annotated.system_prompt(), Some("Initial policy".into())); assert_eq!( annotated.last_user_message(), Some("Where is the weather?".into()) @@ -1126,6 +1140,16 @@ fn test_annotated_llm_types_and_builtin_codecs_cover_mutators_and_codecs() { let updated_messages = json_to_py(py, &json!([{"role": "user", "content": "updated"}])).unwrap(); annotated.set_messages(updated_messages.bind(py)).unwrap(); + let updated_instructions = + json_to_py(py, &json!([{"type": "text", "text": "Updated policy"}])).unwrap(); + annotated + .set_instructions(updated_instructions.bind(py)) + .unwrap(); + let updated_api_specific = + json_to_py(py, &json!({"api": "openai_responses", "background": true})).unwrap(); + annotated + .set_api_specific(updated_api_specific.bind(py)) + .unwrap(); annotated.set_model(Some("updated-model".into())); let updated_params = json_to_py(py, &json!({"temperature": 0.7})).unwrap(); annotated.set_params(updated_params.bind(py)).unwrap(); @@ -1163,6 +1187,14 @@ fn test_annotated_llm_types_and_builtin_codecs_cover_mutators_and_codecs() { annotated.set_extra(updated_extra.bind(py)).unwrap(); assert_eq!(annotated.model(), Some("updated-model".into())); assert_eq!(annotated.last_user_message(), Some("updated".into())); + assert_eq!( + py_to_json(annotated.instructions(py).unwrap().bind(py)).unwrap(), + json!([{"type": "text", "text": "Updated policy"}]) + ); + assert_eq!( + py_to_json(annotated.api_specific(py).unwrap().bind(py)).unwrap(), + json!({"api": "openai_responses", "background": true}) + ); assert_eq!(annotated.store(), Some(true)); assert_eq!(annotated.previous_response_id(), Some("resp_1".into())); assert_eq!( @@ -1194,6 +1226,8 @@ fn test_annotated_llm_types_and_builtin_codecs_cover_mutators_and_codecs() { ); annotated.set_params(py.None().bind(py)).unwrap(); + annotated.set_instructions(py.None().bind(py)).unwrap(); + annotated.set_api_specific(py.None().bind(py)).unwrap(); annotated.set_tools(py.None().bind(py)).unwrap(); annotated.set_tool_choice(py.None().bind(py)).unwrap(); annotated.set_truncation(py.None().bind(py)).unwrap(); @@ -1201,6 +1235,8 @@ fn test_annotated_llm_types_and_builtin_codecs_cover_mutators_and_codecs() { annotated.set_include(py.None().bind(py)).unwrap(); annotated.set_metadata(py.None().bind(py)).unwrap(); assert!(annotated.params(py).unwrap().bind(py).is_none()); + assert!(annotated.instructions(py).unwrap().bind(py).is_none()); + assert!(annotated.api_specific(py).unwrap().bind(py).is_none()); assert!(annotated.tools(py).unwrap().bind(py).is_none()); assert!(annotated.tool_choice(py).unwrap().bind(py).is_none()); assert!(annotated.truncation(py).unwrap().bind(py).is_none()); @@ -1209,9 +1245,18 @@ fn test_annotated_llm_types_and_builtin_codecs_cover_mutators_and_codecs() { assert!(annotated.metadata(py).unwrap().bind(py).is_none()); let bad_messages = json_to_py(py, &json!([{"content": "missing role"}])).unwrap(); - let err = PyAnnotatedLLMRequest::new(bad_messages.bind(py), None, None, None, None, None) - .err() - .unwrap(); + let err = PyAnnotatedLLMRequest::new( + bad_messages.bind(py), + None, + None, + None, + None, + None, + None, + None, + ) + .err() + .unwrap(); assert!(err.to_string().contains("invalid messages")); let bad_params = json_to_py(py, &json!({"temperature": "hot"})).unwrap(); assert!(annotated.set_params(bad_params.bind(py)).is_err()); @@ -1219,6 +1264,18 @@ fn test_annotated_llm_types_and_builtin_codecs_cover_mutators_and_codecs() { assert!(annotated.set_tools(bad_tools.bind(py)).is_err()); let bad_choice = json_to_py(py, &json!({"bad": true})).unwrap(); assert!(annotated.set_tool_choice(bad_choice.bind(py)).is_err()); + let bad_instructions = json_to_py(py, &json!(true)).unwrap(); + assert!( + annotated + .set_instructions(bad_instructions.bind(py)) + .is_err() + ); + let bad_api_specific = json_to_py(py, &json!({"api": "unknown"})).unwrap(); + assert!( + annotated + .set_api_specific(bad_api_specific.bind(py)) + .is_err() + ); let bad_extra = PyList::empty(py); assert!(annotated.set_extra(&bad_extra.into_any()).is_err()); @@ -1477,6 +1534,10 @@ fn test_forced_serialization_error_hooks_cover_unreachable_wrappers() { let annotated = PyAnnotatedLLMRequest { inner: AnnotatedLLMRequest { + instructions: Some(MessageContent::Text("system instructions".into())), + api_specific: Some( + serde_json::from_value(json!({"api": "openai_chat", "seed": 7})).unwrap(), + ), messages: vec![ Message::System { content: MessageContent::Text("system".into()), @@ -1493,13 +1554,15 @@ fn test_forced_serialization_error_hooks_cover_unreachable_wrappers() { max_tokens: Some(8), ..Default::default() }), - tools: Some(vec![nemo_relay::codec::request::ToolDefinition { - tool_type: "function".into(), + tools: Some(vec![nemo_relay::codec::request::ToolDefinition::Function { function: nemo_relay::codec::request::FunctionDefinition { name: "lookup".into(), description: None, parameters: Some(json!({"type": "object"})), + strict: None, + extra: serde_json::Map::new(), }, + extra: serde_json::Map::new(), }]), tool_choice: Some(nemo_relay::codec::request::ToolChoice::Auto), store: None, @@ -1581,6 +1644,11 @@ fn test_forced_serialization_error_hooks_cover_unreachable_wrappers() { "forced serialization failure", |py, _, annotated, _| annotated.messages(py).map(|_| ()), ), + ( + FORCE_ANNOTATED_REQUEST_INSTRUCTIONS_SERIALIZATION_ERROR, + "forced serialization failure", + |py, _, annotated, _| annotated.instructions(py).map(|_| ()), + ), ( FORCE_ANNOTATED_REQUEST_PARAMS_SERIALIZATION_ERROR, "forced serialization failure", @@ -1596,6 +1664,11 @@ fn test_forced_serialization_error_hooks_cover_unreachable_wrappers() { "forced serialization failure", |py, _, annotated, _| annotated.tool_choice(py).map(|_| ()), ), + ( + FORCE_ANNOTATED_REQUEST_API_SPECIFIC_SERIALIZATION_ERROR, + "forced serialization failure", + |py, _, annotated, _| annotated.api_specific(py).map(|_| ()), + ), ( FORCE_ANNOTATED_RESPONSE_MESSAGE_SERIALIZATION_ERROR, "forced serialization failure", diff --git a/crates/types/src/api/llm.rs b/crates/types/src/api/llm.rs index 359244bee..c7cb2f4fd 100644 --- a/crates/types/src/api/llm.rs +++ b/crates/types/src/api/llm.rs @@ -11,6 +11,11 @@ use crate::api::event::PendingMarkSpec; use crate::codec::optimization::LlmOptimizationContribution; use crate::codec::request::AnnotatedLlmRequest; +/// Versioned JSON envelope schema for [`LlmRequestInterceptOutcome`]. +/// +/// Version 2 carries the version 2 annotated-request schema. +pub const LLM_REQUEST_INTERCEPT_OUTCOME_SCHEMA: &str = "nemo.relay.LlmRequestInterceptOutcome@2"; + bitflags! { /// Bitflags that modify LLM-call behavior and observability. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] diff --git a/crates/types/src/codec/request.rs b/crates/types/src/codec/request.rs index f54c0d267..1167b92e6 100644 --- a/crates/types/src/codec/request.rs +++ b/crates/types/src/codec/request.rs @@ -12,6 +12,12 @@ use serde::{Deserialize, Serialize}; use crate::Json; +/// Versioned JSON envelope schema for [`AnnotatedLlmRequest`]. +/// +/// Version 2 adds tagged request components that older exhaustive Rust +/// consumers cannot deserialize safely. +pub const ANNOTATED_LLM_REQUEST_SCHEMA: &str = "nemo.relay.AnnotatedLlmRequest@2"; + // --------------------------------------------------------------------------- // AnnotatedLlmRequest type hierarchy // --------------------------------------------------------------------------- @@ -19,12 +25,18 @@ use crate::Json; /// Structured view of an LLM request, produced by a Codec from opaque /// [`LlmRequest`](crate::api::llm::LlmRequest) content. /// -/// The `extra` field captures any provider-specific keys not modeled by the -/// known fields, ensuring lossless round-trip through `decode`/`encode`. -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +/// The `extra` field captures unknown future top-level keys. Modeled +/// provider-specific controls belong in [`AnnotatedLlmRequest::api_specific`]. +#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)] pub struct AnnotatedLlmRequest { /// Parsed conversation messages. + #[serde(default)] pub messages: Vec, + /// Provider-level instructions that are not part of the conversation array. + /// + /// Anthropic encodes this as `system`; OpenAI Responses uses `instructions`. + #[serde(skip_serializing_if = "Option::is_none")] + pub instructions: Option, /// Model identifier (e.g., `"gpt-4"`, `"claude-sonnet-4-20250514"`). #[serde(skip_serializing_if = "Option::is_none")] pub model: Option, @@ -76,8 +88,13 @@ pub struct AnnotatedLlmRequest { /// OpenAI streaming toggle. #[serde(skip_serializing_if = "Option::is_none")] pub stream: Option, - /// Extensible key-value pairs for unmodeled provider-specific fields. - /// Merged back into the request body during encode via `serde(flatten)`. + /// API-specific request data that does not have portable semantics. + #[serde(skip_serializing_if = "Option::is_none")] + pub api_specific: Option, + /// Unknown future top-level fields. + /// + /// Baseline-aware codecs remove deleted keys and overlay changed or added + /// keys without rebuilding untouched provider JSON. #[serde(flatten)] pub extra: serde_json::Map, } @@ -102,6 +119,14 @@ pub enum Message { #[serde(skip_serializing_if = "Option::is_none")] name: Option, }, + /// A developer instruction message used by OpenAI APIs. + Developer { + /// The message content. + content: MessageContent, + /// Optional sender name. + #[serde(skip_serializing_if = "Option::is_none")] + name: Option, + }, /// An assistant response, optionally containing tool calls. Assistant { /// The message content (optional — may be absent when tool calls are present). @@ -121,6 +146,53 @@ pub enum Message { /// The ID of the tool call this result corresponds to. tool_call_id: String, }, + /// A legacy OpenAI function-result message. + Function { + /// The function result. OpenAI permits an explicit null value. + content: Option, + /// Function name. + name: String, + }, + /// A portable top-level tool-call item, primarily used by OpenAI Responses. + #[serde(rename = "tool_call")] + ToolCallItem { + /// Optional provider item ID. + #[serde(skip_serializing_if = "Option::is_none")] + id: Option, + /// Provider call ID used to correlate the result. + call_id: String, + /// Tool/function name. + name: String, + /// Parsed tool arguments. + arguments: Json, + /// Provider fields without portable semantics. + #[serde(default, flatten)] + extra: serde_json::Map, + }, + /// A portable top-level tool-result item, primarily used by OpenAI Responses. + #[serde(rename = "tool_result")] + ToolResultItem { + /// Optional provider item ID. + #[serde(skip_serializing_if = "Option::is_none")] + id: Option, + /// Provider call ID used to correlate the result. + call_id: String, + /// Provider result payload. + output: Json, + /// Provider fields without portable semantics. + #[serde(default, flatten)] + extra: serde_json::Map, + }, + /// Lossless provider-native request item. + #[serde(rename = "provider_native")] + ProviderNative { + /// Provider surface that owns the value. + provider: String, + /// Native discriminator or a descriptive fallback. + kind: String, + /// Exact provider JSON value. + value: Json, + }, } /// Message content: either a plain string or multimodal parts array. @@ -143,11 +215,83 @@ pub enum ContentPart { Text { /// The text content. text: String, + /// Provider fields without portable semantics. + #[serde(default, flatten)] + extra: serde_json::Map, }, /// An image URL content part. ImageUrl { /// Image URL payload. image_url: OpenAiImageUrl, + /// Provider fields without portable semantics. + #[serde(default, flatten)] + extra: serde_json::Map, + }, + /// Portable image input payload. + Image { + /// Provider-neutral image data object. + image: Json, + /// Provider fields without portable semantics. + #[serde(default, flatten)] + extra: serde_json::Map, + }, + /// Portable audio input payload. + Audio { + /// Provider-neutral audio data object. + audio: Json, + /// Provider fields without portable semantics. + #[serde(default, flatten)] + extra: serde_json::Map, + }, + /// Portable file or document input payload. + File { + /// Provider-neutral file data object. + file: Json, + /// Provider fields without portable semantics. + #[serde(default, flatten)] + extra: serde_json::Map, + }, + /// Assistant refusal content. + Refusal { + /// Refusal text. + refusal: String, + /// Provider fields without portable semantics. + #[serde(default, flatten)] + extra: serde_json::Map, + }, + /// Tool call embedded in a provider content-block array. + ToolUse { + /// Tool call identifier. + id: String, + /// Tool name. + name: String, + /// Parsed arguments. + input: Json, + /// Provider fields without portable semantics. + #[serde(default, flatten)] + extra: serde_json::Map, + }, + /// Tool result embedded in a provider content-block array. + ToolResult { + /// Tool call identifier. + tool_use_id: String, + /// Tool result payload. + content: Json, + /// Whether the tool failed. + #[serde(skip_serializing_if = "Option::is_none")] + is_error: Option, + /// Provider fields without portable semantics. + #[serde(default, flatten)] + extra: serde_json::Map, + }, + /// Lossless provider-native content block. + ProviderNative { + /// Provider surface that owns the value. + provider: String, + /// Native block discriminator or a descriptive fallback. + kind: String, + /// Exact provider JSON value. + value: Json, }, } @@ -182,14 +326,27 @@ pub struct FunctionCall { pub arguments: String, } -/// A tool definition (function schema) available to the model. +/// A tool definition available to the model. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] -pub struct ToolDefinition { - /// The type of tool (typically `"function"`). - #[serde(rename = "type")] - pub tool_type: String, - /// The function definition. - pub function: FunctionDefinition, +#[serde(tag = "type", rename_all = "snake_case")] +pub enum ToolDefinition { + /// Portable function tool. + Function { + /// The function definition. + function: FunctionDefinition, + /// Provider fields on the function-tool wrapper. + #[serde(default, flatten)] + extra: serde_json::Map, + }, + /// Lossless provider-native tool definition. + ProviderNative { + /// Provider surface that owns the value. + provider: String, + /// Native tool discriminator or a descriptive fallback. + kind: String, + /// Exact provider JSON value. + value: Json, + }, } /// A function definition within a tool definition. @@ -203,6 +360,12 @@ pub struct FunctionDefinition { /// The JSON Schema for the function parameters. #[serde(skip_serializing_if = "Option::is_none")] pub parameters: Option, + /// Whether the provider should enforce the parameter schema strictly. + #[serde(skip_serializing_if = "Option::is_none")] + pub strict: Option, + /// Provider fields without portable semantics. + #[serde(default, flatten)] + pub extra: serde_json::Map, } /// Tool choice control: how the model should use available tools. @@ -218,6 +381,167 @@ pub enum ToolChoice { /// Force a specific function by name. #[serde(untagged)] Specific(ToolChoiceFunction), + /// Lossless provider-native tool choice. + #[serde(untagged)] + ProviderNative(ProviderNativeComponent), +} + +/// Lossless provider-native component embedded in the annotated request. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct ProviderNativeComponent { + /// Provider surface that owns the value. + pub provider: String, + /// Native discriminator or a descriptive fallback. + pub kind: String, + /// Exact provider JSON value. + pub value: Json, +} + +/// API-specific request fields that do not have portable semantics. +#[allow( + clippy::large_enum_variant, + reason = "provider wire-schema fields stay directly mutable on each public variant" +)] +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(tag = "api")] +pub enum ApiSpecificRequest { + /// Anthropic Messages-specific request fields. + #[serde(rename = "anthropic_messages")] + AnthropicMessages { + /// Top-level prompt cache control. + #[serde(skip_serializing_if = "Option::is_none")] + cache_control: Option, + /// Reusable container identifier. + #[serde(skip_serializing_if = "Option::is_none")] + container: Option, + /// Requested inference geography. + #[serde(skip_serializing_if = "Option::is_none")] + inference_geo: Option, + /// Provider output configuration. + #[serde(skip_serializing_if = "Option::is_none")] + output_config: Option, + /// Extended-thinking configuration. + #[serde(skip_serializing_if = "Option::is_none")] + thinking: Option, + /// Top-k sampling limit. + #[serde(skip_serializing_if = "Option::is_none")] + top_k: Option, + /// User profile attribution identifier. + #[serde(skip_serializing_if = "Option::is_none")] + user_profile_id: Option, + }, + /// OpenAI Chat Completions-specific request fields. + #[serde(rename = "openai_chat")] + OpenAIChat { + /// Audio output configuration. + #[serde(skip_serializing_if = "Option::is_none")] + audio: Option, + /// Frequency penalty. + #[serde(skip_serializing_if = "Option::is_none")] + frequency_penalty: Option, + /// Deprecated function-call control. + #[serde(skip_serializing_if = "Option::is_none")] + function_call: Option, + /// Deprecated function definitions. + #[serde(skip_serializing_if = "Option::is_none")] + functions: Option>, + /// Token logit bias map. + #[serde(skip_serializing_if = "Option::is_none")] + logit_bias: Option, + /// Whether token log probabilities are requested. + #[serde(skip_serializing_if = "Option::is_none")] + logprobs: Option, + /// Requested output modalities. + #[serde(skip_serializing_if = "Option::is_none")] + modalities: Option>, + /// Request moderation configuration. + #[serde(skip_serializing_if = "Option::is_none")] + moderation: Option, + /// Number of completion choices. + #[serde(skip_serializing_if = "Option::is_none")] + n: Option, + /// Predicted output content. + #[serde(skip_serializing_if = "Option::is_none")] + prediction: Option, + /// Presence penalty. + #[serde(skip_serializing_if = "Option::is_none")] + presence_penalty: Option, + /// Prompt cache routing key. + #[serde(skip_serializing_if = "Option::is_none")] + prompt_cache_key: Option, + /// Prompt cache configuration. + #[serde(skip_serializing_if = "Option::is_none")] + prompt_cache_options: Option, + /// Deprecated prompt cache retention policy. + #[serde(skip_serializing_if = "Option::is_none")] + prompt_cache_retention: Option, + /// Requested reasoning effort. + #[serde(skip_serializing_if = "Option::is_none")] + reasoning_effort: Option, + /// Structured response format configuration. + #[serde(skip_serializing_if = "Option::is_none")] + response_format: Option, + /// Stable safety identifier. + #[serde(skip_serializing_if = "Option::is_none")] + safety_identifier: Option, + /// Best-effort deterministic sampling seed. + #[serde(skip_serializing_if = "Option::is_none")] + seed: Option, + /// Streaming response configuration. + #[serde(skip_serializing_if = "Option::is_none")] + stream_options: Option, + /// Requested response verbosity. + #[serde(skip_serializing_if = "Option::is_none")] + verbosity: Option, + /// Web-search configuration. + #[serde(skip_serializing_if = "Option::is_none")] + web_search_options: Option, + }, + /// OpenAI Responses-specific request fields. + #[serde(rename = "openai_responses")] + OpenAIResponses { + /// Whether the response should run in the background. + #[serde(skip_serializing_if = "Option::is_none")] + background: Option, + /// Context-management entries. + #[serde(skip_serializing_if = "Option::is_none")] + context_management: Option, + /// Conversation identifier or object. + #[serde(skip_serializing_if = "Option::is_none")] + conversation: Option, + /// Request moderation configuration. + #[serde(skip_serializing_if = "Option::is_none")] + moderation: Option, + /// Reusable prompt template reference. + #[serde(skip_serializing_if = "Option::is_none")] + prompt: Option, + /// Prompt cache routing key. + #[serde(skip_serializing_if = "Option::is_none")] + prompt_cache_key: Option, + /// Prompt cache configuration. + #[serde(skip_serializing_if = "Option::is_none")] + prompt_cache_options: Option, + /// Deprecated prompt cache retention policy. + #[serde(skip_serializing_if = "Option::is_none")] + prompt_cache_retention: Option, + /// Stable safety identifier. + #[serde(skip_serializing_if = "Option::is_none")] + safety_identifier: Option, + /// Streaming response configuration. + #[serde(skip_serializing_if = "Option::is_none")] + stream_options: Option, + /// Text output configuration. + #[serde(skip_serializing_if = "Option::is_none")] + text: Option, + }, + /// Custom provider request fields. + #[serde(rename = "custom")] + Custom { + /// Custom API identifier. + api_name: String, + /// Opaque custom API data. + data: Json, + }, } /// A specific tool choice that forces a named function. @@ -265,14 +589,13 @@ impl AnnotatedLlmRequest { /// For [`MessageContent::Parts`], returns the text of the first /// [`ContentPart::Text`] part. pub fn system_prompt(&self) -> Option<&str> { + if let Some(text) = self.instructions.as_ref().and_then(first_content_text) { + return Some(text); + } self.messages.iter().find_map(|m| match m { - Message::System { content, .. } => match content { - MessageContent::Text(s) => Some(s.as_str()), - MessageContent::Parts(parts) => parts.iter().find_map(|p| match p { - ContentPart::Text { text } => Some(text.as_str()), - ContentPart::ImageUrl { .. } => None, - }), - }, + Message::System { content, .. } | Message::Developer { content, .. } => { + first_content_text(content) + } _ => None, }) } @@ -284,13 +607,14 @@ impl AnnotatedLlmRequest { /// the first [`ContentPart::Text`] part. pub fn last_user_message(&self) -> Option<&str> { self.messages.iter().rev().find_map(|m| match m { - Message::User { content, .. } => match content { - MessageContent::Text(s) => Some(s.as_str()), - MessageContent::Parts(parts) => parts.iter().find_map(|p| match p { - ContentPart::Text { text } => Some(text.as_str()), - ContentPart::ImageUrl { .. } => None, - }), - }, + Message::User { content, .. } => first_content_text(content), + Message::ProviderNative { + provider, value, .. + } if provider == "openai_responses" + && value.get("role").and_then(Json::as_str) == Some("user") => + { + native_message_text(value) + } _ => None, }) } @@ -300,11 +624,68 @@ impl AnnotatedLlmRequest { /// Returns `true` if at least one [`Message::Assistant`] variant has a /// non-empty `tool_calls` field. pub fn has_tool_calls(&self) -> bool { - self.messages.iter().any(|m| { - matches!( - m, - Message::Assistant { tool_calls: Some(calls), .. } if !calls.is_empty() - ) + self.messages.iter().any(|m| match m { + Message::Assistant { + tool_calls: Some(calls), + content, + .. + } => !calls.is_empty() || content.as_ref().is_some_and(content_has_tool_use), + Message::Assistant { + content: Some(content), + .. + } => content_has_tool_use(content), + Message::ToolCallItem { .. } => true, + Message::ProviderNative { value, .. } => matches!( + value.get("type").and_then(Json::as_str), + Some("function_call" | "custom_tool_call" | "tool_use") + ), + _ => false, }) } } + +fn first_content_text(content: &MessageContent) -> Option<&str> { + match content { + MessageContent::Text(text) => Some(text.as_str()), + MessageContent::Parts(parts) => parts.iter().find_map(|part| match part { + ContentPart::Text { text, .. } => Some(text.as_str()), + ContentPart::ProviderNative { value, .. } => value + .get("text") + .and_then(Json::as_str) + .or_else(|| value.get("refusal").and_then(Json::as_str)), + ContentPart::ImageUrl { .. } + | ContentPart::Image { .. } + | ContentPart::Audio { .. } + | ContentPart::File { .. } + | ContentPart::Refusal { .. } + | ContentPart::ToolUse { .. } + | ContentPart::ToolResult { .. } => None, + }), + } +} + +fn content_has_tool_use(content: &MessageContent) -> bool { + match content { + MessageContent::Text(_) => false, + MessageContent::Parts(parts) => parts.iter().any(|part| match part { + ContentPart::ToolUse { .. } => true, + ContentPart::ProviderNative { value, .. } => matches!( + value.get("type").and_then(Json::as_str), + Some("tool_use" | "mcp_tool_use" | "server_tool_use") + ), + _ => false, + }), + } +} + +fn native_message_text(value: &Json) -> Option<&str> { + match value.get("content")? { + Json::String(text) => Some(text.as_str()), + Json::Array(parts) => parts.iter().find_map(|part| { + part.get("text") + .and_then(Json::as_str) + .or_else(|| part.get("refusal").and_then(Json::as_str)) + }), + _ => None, + } +} diff --git a/crates/types/src/codec/response.rs b/crates/types/src/codec/response.rs index 87000ee0e..731617966 100644 --- a/crates/types/src/codec/response.rs +++ b/crates/types/src/codec/response.rs @@ -352,8 +352,12 @@ impl AnnotatedLlmResponse { match self.message.as_ref()? { MessageContent::Text(s) => Some(s.as_str()), MessageContent::Parts(parts) => parts.iter().find_map(|p| match p { - crate::codec::request::ContentPart::Text { text } => Some(text.as_str()), - crate::codec::request::ContentPart::ImageUrl { .. } => None, + crate::codec::request::ContentPart::Text { text, .. } => Some(text.as_str()), + crate::codec::request::ContentPart::ProviderNative { value, .. } => value + .get("text") + .and_then(crate::Json::as_str) + .or_else(|| value.get("refusal").and_then(crate::Json::as_str)), + _ => None, }), } } diff --git a/crates/types/tests/serialization_tests.rs b/crates/types/tests/serialization_tests.rs index ddf6feedf..ee0819947 100644 --- a/crates/types/tests/serialization_tests.rs +++ b/crates/types/tests/serialization_tests.rs @@ -11,13 +11,15 @@ use nemo_relay_types::api::event::{ }; use nemo_relay_types::api::llm::{LlmAttributes, LlmRequest, LlmRequestInterceptOutcome}; use nemo_relay_types::api::tool::ToolExecutionInterceptOutcome; -use nemo_relay_types::codec::request::{AnnotatedLlmRequest, Message, MessageContent}; +use nemo_relay_types::codec::request::{AnnotatedLlmRequest, ContentPart, Message, MessageContent}; use nemo_relay_types::codec::response::AnnotatedLlmResponse; use serde_json::{Map, json}; #[test] fn event_round_trips_with_annotated_llm_profiles() { let request = AnnotatedLlmRequest { + instructions: None, + api_specific: None, messages: vec![Message::User { content: MessageContent::Text("hello".into()), name: None, @@ -223,3 +225,99 @@ fn tool_execution_intercept_outcome_converts_from_json() { let outcome: ToolExecutionInterceptOutcome = result.clone().into(); assert_eq!(outcome, ToolExecutionInterceptOutcome::new(result)); } + +#[test] +fn annotated_request_helpers_cover_portable_and_native_components() { + let mut request = AnnotatedLlmRequest { + instructions: Some(MessageContent::Parts(vec![ContentPart::ProviderNative { + provider: "openai_responses".into(), + kind: "refusal".into(), + value: json!({"refusal": "instruction refusal"}), + }])), + ..AnnotatedLlmRequest::default() + }; + assert_eq!(request.system_prompt(), Some("instruction refusal")); + + request.instructions = None; + request.messages = vec![Message::Developer { + content: MessageContent::Parts(vec![ContentPart::Text { + text: "developer instruction".into(), + extra: Map::new(), + }]), + name: None, + }]; + assert_eq!(request.system_prompt(), Some("developer instruction")); + + for (content, expected) in [ + (json!("native string"), Some("native string")), + ( + json!([{"type": "input_text", "text": "native array"}]), + Some("native array"), + ), + (json!({"not": "content"}), None), + ] { + request.messages = vec![Message::ProviderNative { + provider: "openai_responses".into(), + kind: "message".into(), + value: json!({"role": "user", "content": content}), + }]; + assert_eq!(request.last_user_message(), expected); + } + + request.messages = vec![Message::Assistant { + content: Some(MessageContent::Parts(vec![ContentPart::ToolUse { + id: "call_1".into(), + name: "lookup".into(), + input: json!({}), + extra: Map::new(), + }])), + tool_calls: Some(vec![]), + name: None, + }]; + assert!(request.has_tool_calls()); + + request.messages = vec![Message::Assistant { + content: Some(MessageContent::Parts(vec![ContentPart::ProviderNative { + provider: "anthropic_messages".into(), + kind: "server_tool_use".into(), + value: json!({"type": "server_tool_use"}), + }])), + tool_calls: None, + name: None, + }]; + assert!(request.has_tool_calls()); + + request.messages = vec![Message::ToolCallItem { + id: None, + call_id: "call_2".into(), + name: "lookup".into(), + arguments: json!({}), + extra: Map::new(), + }]; + assert!(request.has_tool_calls()); + + request.messages = vec![Message::ProviderNative { + provider: "openai_responses".into(), + kind: "function_call".into(), + value: json!({"type": "function_call"}), + }]; + assert!(request.has_tool_calls()); +} + +#[test] +fn annotated_response_text_reads_provider_native_text_and_refusal() { + for (value, expected) in [ + (json!({"text": "native text"}), "native text"), + (json!({"refusal": "native refusal"}), "native refusal"), + ] { + let response = AnnotatedLlmResponse { + message: Some(MessageContent::Parts(vec![ContentPart::ProviderNative { + provider: "openai_responses".into(), + kind: "native".into(), + value, + }])), + ..AnnotatedLlmResponse::default() + }; + assert_eq!(response.response_text(), Some(expected)); + } +} diff --git a/crates/worker/src/lib.rs b/crates/worker/src/lib.rs index 86a91a49b..982501b64 100644 --- a/crates/worker/src/lib.rs +++ b/crates/worker/src/lib.rs @@ -44,7 +44,7 @@ pub use nemo_relay_types::codec::optimization::{ LlmOptimizationSummary, LlmOptimizationSummaryStatus, LlmOptimizationTokenImpact, LlmOptimizationTokens, }; -pub use nemo_relay_types::codec::request::AnnotatedLlmRequest; +pub use nemo_relay_types::codec::request::{ANNOTATED_LLM_REQUEST_SCHEMA, AnnotatedLlmRequest}; pub use nemo_relay_types::codec::response::AnnotatedLlmResponse; pub use nemo_relay_types::plugin::{ConfigDiagnostic, DiagnosticLevel}; use nemo_relay_worker_proto::v1::plugin_worker_server::{PluginWorker, PluginWorkerServer}; @@ -1518,7 +1518,13 @@ impl WorkerService { let request_value = required_json::(payload.request, "llm request")?; let annotated = payload .annotated_request - .map(|value| decode_json_envelope::(&value)) + .map(|value| { + decode_expected_json_envelope::( + &value, + "annotated llm request", + ANNOTATED_LLM_REQUEST_SCHEMA, + ) + }) .transpose()?; let handler = self.llm_request(&request.registration_name)?; let outcome = with_thread_scope(&scope, || { @@ -1765,6 +1771,20 @@ fn required_json( Ok(decode_json_envelope::(&value)?) } +fn decode_expected_json_envelope( + value: &JsonEnvelope, + field: &str, + expected_schema: &str, +) -> Result { + if value.schema != expected_schema { + return Err(WorkerSdkError::InvalidInput(format!( + "{field} has schema {:?}; expected {expected_schema:?}", + value.schema + ))); + } + Ok(decode_json_envelope(value)?) +} + fn empty_response() -> InvokeResponse { InvokeResponse { result: Some(nemo_relay_worker_proto::v1::invoke_response::Result::Empty( @@ -1800,7 +1820,7 @@ fn llm_request_response(outcome: LlmRequestInterceptOutcome) -> Result InvokeRequest { request.payload.as_mut() { payload.annotated_request = Some(JsonEnvelope { - schema: "nemo.relay.AnnotatedLlmRequest@1".into(), + schema: ANNOTATED_LLM_REQUEST_SCHEMA.into(), json: b"{".to_vec(), }); } request } +fn llm_invoke_with_legacy_annotation(registration_name: &str) -> InvokeRequest { + let mut request = llm_invoke( + registration_name, + RegistrationSurface::LlmRequestIntercept, + llm_request(), + None, + None, + ); + if let Some(nemo_relay_worker_proto::v1::invoke_request::Payload::Llm(payload)) = + request.payload.as_mut() + { + payload.annotated_request = Some(JsonEnvelope { + schema: "nemo.relay.AnnotatedLlmRequest@1".into(), + json: br#"{"messages":[]}"#.to_vec(), + }); + } + request +} + fn scope_context() -> ScopeContext { ScopeContext { scope_stack_id: "stack-1".into(), diff --git a/docs/getting-started/agent-runtime-primer.mdx b/docs/about-nemo-relay/agent-runtime-primer.mdx similarity index 67% rename from docs/getting-started/agent-runtime-primer.mdx rename to docs/about-nemo-relay/agent-runtime-primer.mdx index 36dacbb70..b70cef34a 100644 --- a/docs/getting-started/agent-runtime-primer.mdx +++ b/docs/about-nemo-relay/agent-runtime-primer.mdx @@ -1,7 +1,7 @@ --- title: "Agent Runtime Primer" description: "Learn the NeMo Relay runtime model for scopes, middleware, events, and integration boundaries." -position: 1 +position: 2 --- import { MermaidStyles } from "@/components/MermaidStyles"; @@ -11,7 +11,8 @@ SPDX-License-Identifier: Apache-2.0 */} NeMo Relay is a portable runtime layer for agent systems that already have an application, framework, or model provider. Use this primer when you need to -understand what NeMo Relay adds before running [Quick Start](/getting-started/quick-start). +understand what NeMo Relay adds before choosing an installation, quick-start, or +integration path. Agent applications usually cross several boundaries in one request: an entry point starts work, the agent calls a model, the model asks for tools, tools call @@ -91,29 +92,4 @@ It does not replace: Instead, it gives those systems a shared runtime contract for call boundaries, policy hooks, event emission, and export. -## Find the Right Starting Point - -Use this section as a router, not a setup checklist. Start from the destination -or outcome that matches your task; the linked page points you to the setup, -configuration, and validation steps for that path. - -- **Direct Python, Node.js, or Rust application APIs:** Application code owns - callback placement, provider authentication, and plugin initialization. Start - with [Quick Start](/getting-started/quick-start), then use - [Instrument Applications](/instrument-applications/about) when you own the - tool or LLM call site. -- **LangChain, LangGraph, Deep Agents, or OpenClaw:** Start with - [Supported Integrations](/supported-integrations/about) to choose the - maintained path and support level. Local wiring still belongs to the - application, framework, or OpenClaw plugin configuration. -- **New framework, orchestration, SDK, or provider integration:** The framework - or adapter owns scheduling, retries, callbacks, and provider payloads. Start - with [Integrate into Frameworks](/integrate-into-frameworks/about). -- **Local Claude Code, Codex, or Hermes runs:** The coding-agent harness - owns invocation while Relay observes hooks, gateway-routed model traffic, and - exporter output. Start with [NeMo Relay CLI](/nemo-relay-cli/about). -- **Reusable runtime behavior across services or teams:** Runtime plugin - configuration owns reusable middleware, subscribers, exporters, model pricing, - policy, or adaptive behavior. Start with [Build Plugins](/build-plugins/about), - use [Observability](/configure-plugins/observability/about) for export setup, and use - [Adaptive](/configure-plugins/adaptive/about) after baseline instrumentation is working. +For setup routing, start with [Getting Started](/getting-started/about). diff --git a/docs/about-nemo-relay/architecture.mdx b/docs/about-nemo-relay/architecture.mdx index 155c3aeef..8af4a69d2 100644 --- a/docs/about-nemo-relay/architecture.mdx +++ b/docs/about-nemo-relay/architecture.mdx @@ -1,7 +1,7 @@ --- title: "Architecture" description: "" -position: 2 +position: 3 --- import { MermaidStyles } from "@/components/MermaidStyles"; diff --git a/docs/about-nemo-relay/concepts/codecs.mdx b/docs/about-nemo-relay/concepts/codecs.mdx index 311380dec..312e9d11c 100644 --- a/docs/about-nemo-relay/concepts/codecs.mdx +++ b/docs/about-nemo-relay/concepts/codecs.mdx @@ -67,6 +67,36 @@ Response decoding improves observability and downstream consistency. It does not automatically change the value returned to the application unless a separate typed value boundary also does so. +The built-in request codecs are lossless patch codecs. For an unchanged +annotation, the following identity holds at the JSON-value level: + +```text +encode(decode(original), original) == original +``` + +Encoding compares the edited annotation with a freshly decoded baseline and +patches only fields that changed. This preserves caller-selected forms such as +string content versus content blocks, explicit `null` values, legacy field +names, ordered input items, metadata, and unknown fields. Removing a key from +top-level `extra` removes that unknown key from the provider request; adding or +changing a key overlays it. + +`AnnotatedLlmRequest` separates portable and provider-specific data: + +- `instructions` represents Anthropic `system` and OpenAI Responses + `instructions`. +- `messages`, content parts, function calls and results, tools, and tool choice + expose portable components when the shapes have shared semantics. +- `api_specific` is a tagged, mutable surface for modeled Anthropic Messages, + OpenAI Chat Completions, or OpenAI Responses controls. +- Provider-only messages, content blocks, input items, tools, and tool choices + use `{ provider, kind, value }`, where `value` is the exact native JSON. +- Top-level `extra` is reserved for unknown future fields. + +A provider-native component can only be encoded by the provider surface named +in its `provider` field. Provider mismatches and portable edits that the target +API cannot represent fail before the provider callback. + ## Normalized Data Consumption Normalized codec output is applicable to several runtime layers: @@ -101,11 +131,12 @@ disambiguate an otherwise identical request shape, such as an Anthropic Messages request without a top-level `system` field, but no provider annotation is invented without either a matching provider surface or a recognized hint. -Provider extraction covers model names, messages, generation parameters, tool -definitions, tool calls, finish reasons, usage, cost, provider-specific fields, -and replayable request or response JSON when the source payload contains enough -information. Provider codecs should preserve unknown fields and treat request -encoding as a merge over the original provider payload. +Provider extraction covers model names, instructions, messages, generation +parameters, tool definitions, tool calls, finish reasons, usage, cost, +provider-specific fields, and replayable request or response JSON when the +source payload contains enough information. Provider codecs preserve unknown +fields and treat request encoding as a baseline-aware patch over the original +provider payload. The response-extraction interface is the existing response codec contract: `LlmResponseCodec::decode_response` returns `AnnotatedLlmResponse`. Built-in @@ -133,6 +164,13 @@ hints let Anthropic Messages requests without a top-level `system` field decode through the Anthropic provider surface instead of being treated as shape-only OpenAI Chat payloads. +The `nemo-relay` gateway always enables matching request codecs for +`/v1/messages`, `/v1/chat/completions`, and `/v1/responses`, for both buffered +and streaming execution. Count-token, model, probe, and non-LLM passthrough +routes do not enable request codecs. Gateway request intercepts must therefore +edit generation bodies through `annotated_request`; raw `request.content` +remains writable on routes without a request codec. + ### Agent Payload Extraction Agent payload extraction is separate from provider codecs. Coding agents, diff --git a/docs/about-nemo-relay/concepts/events.mdx b/docs/about-nemo-relay/concepts/events.mdx index 93741a5d5..f3ae6a864 100644 --- a/docs/about-nemo-relay/concepts/events.mdx +++ b/docs/about-nemo-relay/concepts/events.mdx @@ -149,6 +149,36 @@ read. The following payload excerpt shows the stable skill-load fields. Standard fields from the [shared envelope](#shared-envelope), such as `atof_version`, `uuid`, `parent_uuid`, and `timestamp`, are omitted. + +Automatic skill-load extraction is best effort. Relay emits `skill.load` only +through one of the following detection paths: + +- An integration supplies a precomputed detection with a nonempty skill name + and supported source before removing the original tool arguments. +- Relay sees a recognized first-class skill tool with a nonempty `skill`, + `skill_name`, or `name` argument. +- Relay sees the original arguments for a supported reader request that meets + all of the reader criteria below. + +For a reader request, all of the following conditions must be true: + +- The tool name matches a supported reader, either directly or through one + supported shell wrapper. +- A recognized argument field or parsed command contains a path or URI. Its + final path component equals `SKILL.md` under an ASCII case-insensitive + comparison. +- The path's immediate parent directory has a nonempty name, which Relay uses + as the skill name. +- The request reads the entire file. Relay ignores requests that use partial or + range controls, unsupported transformations, metadata-only operations, + pipelines, redirections, command substitutions, or compound commands. + Relay also ignores output-transforming and unreviewed reader options. + +Relay checks the request before tool execution. The mark records an attempted +load even if execution later fails. If Relay doesn't emit a mark, that doesn't +prove that the skill wasn't loaded. + + ```json { "kind": "mark", @@ -167,23 +197,67 @@ Relay uses these `skill_load_source` values: | --- | --- | | `skill_tool` | A first-class skill tool, such as `Skill` or `skill_view` | | `structured_read` | A structured request for a whole-file read | -| `shell_read` | A standalone, complete `cat`, `bat`, `batcat`, or PowerShell `Get-Content` command | - -For reader detections, the path must end in `SKILL.md`; the skill name is its -immediate parent directory. A nonzero structured offset, any structured limit -or range control, shell pipelines, redirections, compound commands, and -range-limited shell readers do not count. An explicit structured `offset` of -`0` still represents a complete read. Relay omits the full path and shell -command from the mark. +| `shell_read` | A complete `cat`, `bat`, `batcat`, or PowerShell `Get-Content` command, either direct or within one supported wrapper | + +For structured reads, Relay recognizes only the following reviewed tool-name +suffixes and argument fields: + +- Tool suffixes: `read`, `read_file`, `read_text_file`, + `read_multiple_files`, `file_read`, `read_resource`, `get_file_contents`, + and `read_file_content`. +- Argument fields: `path`, `file_path`, `filepath`, `filename`, `file`, + `paths`, `uri`, and `absolute_path`. + +For shell reads, Relay recognizes the following direct shell tool names after +normalizing case and nonalphanumeric separators: + +- Tool names: `sh`, `bash`, `zsh`, `fish`, `shell`, `shell_command`, `exec`, + `exec_command`, `execute`, `terminal`, `run_command`, `run_shell_command`, + `shell_exec`, `powershell`, and `pwsh`. +- Argument fields: `command` and `cmd`. + +Relay also recognizes a single wrapper layer inside a direct shell tool. Each +wrapper must use one of the following forms: + +- POSIX shells: `sh -c`, `bash -c`, `bash -lc`, `zsh -c`, or `zsh -lc`. +- Fish: `fish -c` or `fish --command=...`. +- PowerShell: `powershell -Command` or `pwsh -Command`. + +The `sh`, `bash`, `zsh`, and `fish` wrappers support `cat`, `bat`, and `batcat`. +PowerShell wrappers support `Get-Content`. Relay treats unknown tool names, +argument fields, wrapper forms, Windows `type`, and PowerShell aliases as +unsupported until an integration explicitly adds them to the registry. + +The shell reader option registry is deliberately narrow: + +- `cat` supports no options or the exact content-preserving `-u` option. Other + options, including bundled short options, are unsupported. +- `bat` and `batcat` require one exact plain-output option: `-p` or `--plain`. + Relay rejects every additional option. +- `Get-Content` supports positional paths and the exact `-Path`, + `-LiteralPath`, `-Raw`, `-Encoding`, `-ReadCount`, `-AsByteStream`, and + `-Force` options. PowerShell abbreviations and other options are unsupported. + +For a reader request, the final path component must equal `SKILL.md` under an +ASCII case-insensitive comparison. Relay uses the immediate parent directory as +the skill name. Relay doesn't emit a mark when a structured request has a +`offset` other than the integer `0`, a `limit`, or a range control. It also +doesn't emit a mark for shell commands that use pipelines, redirections, +compound commands, or range-limited readers, or help, version, listing, or +diagnostic operations. It also doesn't emit a mark for output-transforming or +unreviewed reader options. An explicit structured `offset` of `0` still +represents a complete read. Relay omits the full path and shell command from +the mark. The mark has these lifecycle semantics: - Relay emits it immediately after the tool-start event and before execution, with the tool span as its parent. -- It records an eager load attempt and remains present if the tool later fails. -- A blocked call that never starts does not emit the mark. -- Relay emits each skill once per tool call. A later call that attempts the - same load emits another mark. +- The mark records an eager load attempt and remains in the event stream if the + tool later fails. +- Relay doesn't emit the mark for a blocked call that never starts. +- Relay emits at most one mark for each skill in a tool call. A later tool call + that attempts the same load emits another mark. Claude Code `Skill`, Codex complete-reader shell calls, Hermes `skill_view`, and tool calls made through the Rust, Python, Node.js, Go/FFI, LangChain, LangGraph, diff --git a/docs/about-nemo-relay/concepts/index.mdx b/docs/about-nemo-relay/concepts/index.mdx index 5a6d93bf8..0ba4db236 100644 --- a/docs/about-nemo-relay/concepts/index.mdx +++ b/docs/about-nemo-relay/concepts/index.mdx @@ -1,7 +1,7 @@ --- title: "Concepts" description: "" -position: 4 +position: 5 --- {/* SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 */} diff --git a/docs/about-nemo-relay/concepts/subscribers.mdx b/docs/about-nemo-relay/concepts/subscribers.mdx index b794a7966..6fdc12bff 100644 --- a/docs/about-nemo-relay/concepts/subscribers.mdx +++ b/docs/about-nemo-relay/concepts/subscribers.mdx @@ -24,6 +24,11 @@ waiting for subscriber callbacks or exporter work. Events describe what happened. Subscribers are the components that watch those events. +In this documentation, an event is **emitted** when the runtime submits it for +subscriber dispatch. An event is **delivered** to a subscriber when its callback +runs. An event is **exported** when an exporter completes the downstream work +defined by its API. On native targets, these are separate milestones. + That separation matters: - The runtime can emit one canonical event stream @@ -127,8 +132,21 @@ order and subscriber snapshot order. ## Waiting for Delivery -Use the flush API when application shutdown, tests, or examples must observe -side effects from subscriber callbacks that have already been queued: +An event-producing call returning is not a delivery barrier. These guarantees +apply when you call a barrier outside a native subscriber callback. Use the +following milestones to choose the barrier that matches the output you need to +observe: + +| Milestone | What It Establishes | What It Does Not Establish | +|---|---|---| +| Event-producing call returns | The runtime call completed. On native targets, the runtime normally queues subscriber work for background dispatch. | Subscriber callbacks ran, exporter work completed, or output is visible or durable. | +| Subscriber flush returns | All native subscriber callbacks queued before the flush call have completed. | Exporter-owned workers drained or a remote backend accepted the data. | +| Exporter barrier returns | The runtime drained the shared subscriber queue, and the exporter completed its documented `force_flush()`, `export()`, or equivalent operation. For a worker-backed exporter, that operation can stop after its configured timeout. | Exactly-once delivery, machine-crash durability, confirmed completion by a worker that timed out, or success beyond the exporter's documented error and timeout behavior. | +| Plugin clear or exporter shutdown returns | NeMo Relay completed the component's documented graceful teardown steps, or reported an error. | Retention if the process terminates before teardown completes or completion by a worker that timed out during teardown. | +| Process terminates before a barrier | You can observe only callback side effects and exporter work that completed before termination. | Retention of queued subscriber callbacks or pending exporter work. | + +Use the subscriber flush API when application shutdown, tests, or examples must +observe side effects from callbacks that were already queued before the barrier: - Rust: `nemo_relay::api::subscriber::flush_subscribers()?` - Python: `nemo_relay.subscribers.flush()` @@ -136,6 +154,24 @@ side effects from subscriber callbacks that have already been queued: callback side effects - FFI: `nemo_relay_flush_subscribers()` +Events emitted concurrently after the flush barrier are not covered by that +call. Exporters with their own workers or batch processors require the +exporter-specific flush, export, or shutdown operation described in +[Observability](/configure-plugins/observability/about). + + +Do not invoke a subscriber flush, an exporter barrier, or plugin clear from a +native subscriber callback. To avoid blocking its worker, the native dispatcher +returns a re-entrant subscriber flush without waiting. Callbacks later in the +active dispatch snapshot can still run after the barrier returns. Run those +operations after the callback returns. + +If the process terminates before subscriber delivery and exporter teardown +complete, queued telemetry can be lost. Returning from the event-producing API +does not make that telemetry crash-safe. + + + ### Forwarding and Export Some subscribers translate the event stream into external formats or transport @@ -190,8 +226,10 @@ plugin component. Use these practices when applying the concept in application or integration code. - Use a plain subscriber when you want in-process custom behavior. -- Flush subscribers before asserting on printed output, captured lists, or files - written by subscriber callbacks. +- Flush subscribers before inspecting printed output, captured lists, or other + custom callback side effects. +- Use the exporter's documented barrier before inspecting exporter output. +- Clear plugin-managed exporters during graceful shutdown. - Use `event.to_dict()` or `event.to_json()` when a host runtime or exporter needs the canonical event JSON shape in-process. - Use a scope-local subscriber when the observation should disappear with the diff --git a/docs/about-nemo-relay/ecosystem.mdx b/docs/about-nemo-relay/ecosystem.mdx index 61a16ce18..3db5c47fa 100644 --- a/docs/about-nemo-relay/ecosystem.mdx +++ b/docs/about-nemo-relay/ecosystem.mdx @@ -1,7 +1,7 @@ --- title: "Ecosystem" description: "Understand how NeMo Relay fits with agent frameworks, providers, and the NVIDIA NeMo ecosystem." -position: 3 +position: 4 --- import { MermaidStyles } from "@/components/MermaidStyles"; @@ -42,6 +42,22 @@ agent products. A framework asks, "What should the agent do next?" NeMo Relay asks, "When the agent does work, which scope owns it, which middleware applies, what events are emitted, and which subscribers can consume the result?" +## How NeMo Relay Relates to Other Tooling + +NeMo Relay is an execution-time runtime contract. It captures work at the +scope, tool, and LLM boundaries, and can apply middleware before or around that +work. Telemetry conventions and observability or evaluation products serve +different roles and can be used alongside Relay. + +| Tooling | Primary Role | Relationship to NeMo Relay | +|---|---|---| +| OpenTelemetry GenAI conventions | Define a common representation for AI telemetry. | Relay owns runtime capture and control, then can export its lifecycle events as generic OTLP spans for an existing OpenTelemetry tracing pipeline. | +| Langfuse, LangSmith, and Arize Phoenix | Store, explore, and evaluate traces and agent runs. | These products are outside Relay's execution control loop. Use Relay to establish consistent execution boundaries and event data, then send data through the configured subscriber or export path that fits your backend. | + +Relay is not a replacement for a telemetry standard or an observability backend. +Its role is to make the real execution path observable and controllable before +the resulting lifecycle data is stored, visualized, or evaluated elsewhere. + ```mermaid diff --git a/docs/about-nemo-relay/overview.mdx b/docs/about-nemo-relay/overview.mdx index 8262b319b..5f6a23b9c 100644 --- a/docs/about-nemo-relay/overview.mdx +++ b/docs/about-nemo-relay/overview.mdx @@ -31,6 +31,9 @@ observability backend, or guardrail authoring system. It gives those systems a common runtime boundary to meet at. +For how Relay complements OpenTelemetry GenAI conventions and observability or +evaluation products, see [How NeMo Relay Relates to Other Tooling](/about-nemo-relay/ecosystem#how-nemo-relay-relates-to-other-tooling). + The first design question is simple: where can Relay observe or control the real work? The answer determines whether you should use a CLI sidecar, direct SDK instrumentation, a maintained integration, a framework wrapper, or a plugin. @@ -115,7 +118,7 @@ Use the tasks below to build your understanding and set up Relay: | Task | Start With | |---|---| | Install packages | [Installation](/getting-started/installation) | -| Understand the mental model | [Agent Runtime Primer](/getting-started/agent-runtime-primer) | +| Understand the mental model | [Agent Runtime Primer](/about-nemo-relay/agent-runtime-primer) | | Configure plugin files | [Plugin Configuration Files](/configure-plugins/plugin-configuration-files) | | Export traces or trajectories | [Observability](/configure-plugins/observability/about) | | Tune performance with adaptive behavior | [Adaptive](/configure-plugins/adaptive/about) | diff --git a/docs/about-nemo-relay/release-notes/highlights.mdx b/docs/about-nemo-relay/release-notes/highlights.mdx index 822e0d705..d92cefd4d 100644 --- a/docs/about-nemo-relay/release-notes/highlights.mdx +++ b/docs/about-nemo-relay/release-notes/highlights.mdx @@ -45,12 +45,26 @@ dynamic plugins. - The built-in PII redaction plugin now cleans mark data, category profiles, metadata, and generic scope inputs and outputs before subscribers or exporters receive them. +- The PII redaction plugin now supports ordered, composable profiles. A single + component can apply multiple privacy policies to marks, LLM and tool + inputs/outputs, and scope metadata. Relay runs profiles by priority and uses + array order to break ties. +- The opt-in `trajectory_context` preset removes request and response content + from buffered and streaming events for OpenAI Chat Completions, OpenAI + Responses, and Anthropic Messages. It preserves conversation topology, + tool-call relationships, agent hierarchy, routing, usage, cost, and + optimization telemetry. You can preserve custom mark payloads or redact them + recursively. - Relay now tracks LLM event-history freshness per agent. A new agent and the first LLM start after a `compaction` mark retain the complete sanitized history. Later starts keep the system instructions, latest user turn, and following messages. The provider request itself does not change. - A central provider-codec factory now resolves canonical provider names and - the built-in request, response, and streaming codecs. + the built-in request, response, and streaming codecs. The coding-agent + gateway uses lossless request codecs for Anthropic Messages, OpenAI Chat + Completions, and OpenAI Responses generation routes, so middleware can edit + portable or provider-native request annotations without rebuilding untouched + wire data. ### Observability and Optimization @@ -117,8 +131,11 @@ dynamic plugins. typed attribute paths at the same time. - Check ATIF and event consumers for assumptions about marks, requested model names, or complete request histories. -- Update Rust exhaustive matches, direct struct literals, and editor metadata - consumers for the new public variants and fields. +- Update Rust exhaustive matches and direct `AnnotatedLlmRequest`, message, + content, and tool literals for the new public variants and fields. Rebuild + Rust native plugins and `grpc-v1` workers that consume request annotations + against 0.6, and narrow their manifest requirement to + `compat.relay = ">=0.6,<1.0"`. - Migrate Rust code that directly constructs `LlmJsonStream`, and update Go call sites to handle the error returned by `LlmStream.Close`. Call an explicit close method when a consumer stops reading an LLM stream early. diff --git a/docs/about-nemo-relay/release-notes/index.mdx b/docs/about-nemo-relay/release-notes/index.mdx index 953985ffb..28bac2c95 100644 --- a/docs/about-nemo-relay/release-notes/index.mdx +++ b/docs/about-nemo-relay/release-notes/index.mdx @@ -1,7 +1,7 @@ --- title: "Release Notes" description: "Read NeMo Relay release notes and find the official release history." -position: 5 +position: 6 --- {/* SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 */} @@ -31,8 +31,11 @@ If you're upgrading from 0.5, plan for these changes: LLM event contains the full request history, ATIF contains marks, or `step.model_name` always matches the requested model. - Update Rust code that exhaustively matches public enums or directly - constructs mark types, ATOF sink types, or editor metadata types. The - provided builders and constructors are the safest migration path. + constructs `AnnotatedLlmRequest`, message/content/tool variants, mark types, + ATOF sink types, or editor metadata types. Rebuild Rust native plugins and + `grpc-v1` workers that consume request annotations against 0.6, and declare + `compat.relay = ">=0.6,<1.0"` for those artifacts. The provided builders and + constructors are the safest migration path. - Update managed LLM stream consumers that construct `LlmJsonStream` directly or ignore Go `LlmStream.Close` results. Use the stream constructors and explicit close methods described in the compatibility notes when exiting a @@ -56,14 +59,20 @@ Here's what changed across the main product surfaces: and Hermes can share that gateway, coordinate a single recovery attempt, and let it shut down after it becomes idle. - **Events and security:** Mark, scope-start, and scope-end sanitizers now work - at global, scope-local, and plugin levels. The PII plugin sanitizes those - surfaces, later LLM events can omit stale history, and Relay records - skill-load attempts with `skill.load` marks. + at global, scope-local, and plugin levels. The PII plugin sanitizes these + surfaces and supports ordered, composable profiles plus an opt-in + `trajectory_context` preset. Later LLM events can omit stale history, and + Relay records skill-load attempts with `skill.load` marks. - **Observability and optimization:** ATOF can send every event to multiple file and stream sinks. OTLP exporters use typed attributes, trace exporters can show marks as tool spans, session identifiers correlate across formats, ATIF reports the effective response model, and plugins can contribute optimization evidence. +- **Provider request codecs:** The coding-agent gateway now enables lossless + request annotations for Anthropic Messages, OpenAI Chat Completions, and + OpenAI Responses generation routes. Interceptors edit `annotated_request`, + while no-op round trips retain provider-native fields, explicit nulls, and + wire representations. - **Plugins and configuration:** Rust, Python, and Node.js embedding hosts can own native and worker dynamic plugin lifecycles. Experimental C and Go entry points expose the same lifecycle, and the CLI can edit nested lists, maps, and @@ -99,6 +108,7 @@ Use these guides to explore the main 0.6 capabilities: - [Tool Execution Intercept Outcomes](/reference/tool-execution-intercept-outcomes) - [Instrument an LLM Call](/instrument-applications/instrument-llm-call) - [Provider Response Codecs and Model Pricing](/integrate-into-frameworks/provider-response-codecs) +- [Codecs](/about-nemo-relay/concepts/codecs) - [Observability Configuration](/configure-plugins/observability/configuration) - [ATOF Export](/configure-plugins/observability/atof) - [OpenTelemetry Export](/configure-plugins/observability/opentelemetry) diff --git a/docs/about-nemo-relay/release-notes/known-issues.mdx b/docs/about-nemo-relay/release-notes/known-issues.mdx index d17fbb98b..ad638d561 100644 --- a/docs/about-nemo-relay/release-notes/known-issues.mdx +++ b/docs/about-nemo-relay/release-notes/known-issues.mdx @@ -81,6 +81,86 @@ The following limitations and migration steps apply to NVIDIA NeMo Relay 0.6. ### Compatibility and Migration Notes +#### Migrate Composable PII Redaction Configuration + + +If your Rust code constructs `PiiRedactionConfig` with a struct literal, +initialize the new `profiles` field or use `..Default::default()`. Serialization +now omits the default-valued legacy fields (`mode`, `input`, `output`, `mark`, +`tool_input`, `tool_output`, and `priority`). Update any code that depends on +the previous exact serialized shape. + +Existing TOML and JSON files still work. You can continue to use the legacy +single-policy form, but don't combine its fields with `profiles`. Relay runs +profiles in ascending priority and uses array order to break ties. Profiles can +apply to marks, LLM input/output, tool input/output, and scope metadata. +Existing PII redaction configurations now also sanitize LLM and tool scope +metadata, so downstream telemetry can change even when its schema does not. + +For profile examples and migration guidance, refer to the +[PII Redaction Configuration](/configure-plugins/pii-redaction/configuration) +guide. PR +[#512](https://github.com/NVIDIA/NeMo-Relay/pull/512) introduced the composable +configuration API. + + +#### Configure Structure-Preserving Trajectory Redaction + + +If your Rust code constructs `BuiltinBackendConfig` with a struct literal, +initialize the new `preset` and `custom_mark_payload_policy` fields or use +`..Default::default()`. Serialization now omits the default `action = "remove"` +value. Update any code that depends on the previous exact serialized shape. + +The `trajectory_context` preset is opt-in. If you leave `preset` unset, existing +runtime behavior stays the same. With the preset enabled, Relay removes request +and response content from buffered and streaming events for OpenAI Chat +Completions, OpenAI Responses, and Anthropic Messages. It preserves conversation +topology, tool-call relationships, agent hierarchy, routing, usage, cost, and +optimization fields. Relay preserves opaque custom marks by default; set +`custom_mark_payload_policy = "redact_all_leaves"` to redact their scalar +values recursively. If retained metadata or custom marks can contain email +addresses, apply an email-redaction profile later in the profile list. + +The policy affects observability data only. It doesn't change the provider +request or the client-visible provider response. For the full configuration +options, refer to the +[PII Redaction Configuration](/configure-plugins/pii-redaction/configuration) +guide. PR +[#513](https://github.com/NVIDIA/NeMo-Relay/pull/513) added trajectory-context +redaction on top of [PR #512](https://github.com/NVIDIA/NeMo-Relay/pull/512). + + +#### Update Plugins and Workers That Consume Request Annotations + + +`AnnotatedLlmRequest` now includes shared instructions, a tagged +`api_specific` surface, and additional message, content, tool, and tool-choice +variants. This is a Rust source-compatibility break for exhaustive matches and +direct struct literals. + +Rebuild Rust native plugins and Rust `grpc-v1` workers that accept annotated +LLM requests against NeMo Relay 0.6. Update Python workers to the 0.6 worker +SDK. Set their manifest requirement to +`compat.relay = ">=0.6,<1.0"` instead of the broader `>=0.5,<1.0` guidance. A +0.5 consumer cannot decode the added fields and enum variants. Relay 0.6 labels +request annotations as `nemo.relay.AnnotatedLlmRequest@2` and request-intercept +results as `nemo.relay.LlmRequestInterceptOutcome@2`. The host rejects a +dynamic plugin that registers an LLM request intercept while its +`compat.relay` range still admits Relay 0.5. + +Python workers continue to receive annotations as `dict[str, Any]`, but the +worker SDK validates the versioned envelope. Node.js, Go, and raw C FFI +request-intercept callbacks continue to carry the annotation as JSON, so they +do not require new binding DTOs. Consumers in those languages should still +avoid closed-world matching on role or component discriminator strings. + +The [Codecs](/about-nemo-relay/concepts/codecs) guide documents the expanded +annotation schema and mutation contract. PR +[#509](https://github.com/NVIDIA/NeMo-Relay/pull/509) introduced the request +codec integration and Rust source break. + + #### Update Managed LLM Stream Consumers @@ -277,6 +357,10 @@ catalog lists the current paths and task routing. PR ### Fixed in NVIDIA NeMo Relay 0.6 +- Coding-agent gateway generation routes now decode request annotations for + Anthropic Messages, OpenAI Chat Completions, and OpenAI Responses. Request + interceptors can edit portable and provider-native fields while unchanged + nested fields, explicit nulls, and provider representations survive encoding. - Installed coding-agent sessions now acquire the shared gateway through MCP before hooks or routed provider traffic. This prevents cold-start loss and coordinates concurrent startup and recovery. diff --git a/docs/build-plugins/dynamic-plugins/about.mdx b/docs/build-plugins/dynamic-plugins/about.mdx index bd3d0bcee..68f2331ba 100644 --- a/docs/build-plugins/dynamic-plugins/about.mdx +++ b/docs/build-plugins/dynamic-plugins/about.mdx @@ -64,6 +64,13 @@ operators must enable a registered dynamic plugin explicitly. Declare only the capabilities the plugin needs. Add `config_schema.path` only with the `config_schema` capability. + +Dynamic plugins that register an LLM request intercept consume the versioned +request annotation contract. For Relay 0.6 request intercepts, set +`compat.relay = ">=0.6,<1.0"`, or another range that excludes Relay 0.5. Relay +rejects the registration when the manifest still admits Relay 0.5. + + The following requirements vary by execution lane: | Manifest area | Native dynamic plugin | Worker plugin | diff --git a/docs/build-plugins/dynamic-plugins/grpc-worker/about.mdx b/docs/build-plugins/dynamic-plugins/grpc-worker/about.mdx index abb068f39..07c27d5b6 100644 --- a/docs/build-plugins/dynamic-plugins/grpc-worker/about.mdx +++ b/docs/build-plugins/dynamic-plugins/grpc-worker/about.mdx @@ -64,6 +64,12 @@ Set `compat.worker_protocol` to `grpc-v1` and include `plugin_worker` in the capability list. Refer to [Choose a Worker Runtime](#choose-a-worker-runtime) for the runtime-specific `load` and `source` requirements. +If the worker registers an LLM request intercept, set +`compat.relay = ">=0.6,<1.0"`, or another range that excludes Relay 0.5. The +Relay 0.6 host sends `nemo.relay.AnnotatedLlmRequest@2` and expects +`nemo.relay.LlmRequestInterceptOutcome@2`; it rejects request-intercept +registrations whose manifest still admits Relay 0.5. + The worker receives its activation ID, plugin ID, worker and host endpoints, and a local activation token through environment variables. Do not start the worker directly during normal operation. The host supplies these values. diff --git a/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx b/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx index 602fcc3b6..db11f2acc 100644 --- a/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx +++ b/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx @@ -55,6 +55,11 @@ Set `plugin.kind` to `rust_dynamic`, `compat.native_api` to `"1"`, and from the manifest. The exported symbol must return a descriptor whose `plugin_kind` matches `plugin.id` exactly. +If the plugin registers an LLM request intercept, set +`compat.relay = ">=0.6,<1.0"`, or another range that excludes Relay 0.5. Relay +rejects the registration when the manifest admits a host that predates the +version 2 request annotation envelope. + ## Create a Native Plugin Create a Rust library project with the following `Cargo.toml` file: diff --git a/docs/build-plugins/dynamic-plugins/native-dynamic/rust-native-plugin-example.mdx b/docs/build-plugins/dynamic-plugins/native-dynamic/rust-native-plugin-example.mdx index 47495d106..89bf0ca82 100644 --- a/docs/build-plugins/dynamic-plugins/native-dynamic/rust-native-plugin-example.mdx +++ b/docs/build-plugins/dynamic-plugins/native-dynamic/rust-native-plugin-example.mdx @@ -29,7 +29,7 @@ edition = "2024" crate-type = ["cdylib"] [dependencies] -nemo-relay-plugin = "0.5.0" +nemo-relay-plugin = "0.6.0" serde_json = "1" ``` @@ -70,7 +70,7 @@ id = "acme.native_policy" kind = "rust_dynamic" [compat] -relay = ">=0.5,<1.0" +relay = ">=0.6,<1.0" native_api = "1" [defaults] diff --git a/docs/configure-plugins/observability/about.mdx b/docs/configure-plugins/observability/about.mdx index 85ef99e54..bb2bef277 100644 --- a/docs/configure-plugins/observability/about.mdx +++ b/docs/configure-plugins/observability/about.mdx @@ -19,7 +19,8 @@ manual subscribers. A common first path is: 1. Instrument one scope, tool call, or LLM call in your application. 2. Enable the Observability plugin with only one exporter, such as ATOF JSONL. 3. Run one request through the instrumented path. -4. Confirm the emitted output before adding more exporters or sanitization rules. +4. Clear the plugin, or call the manual exporter's documented barrier. +5. Confirm the output before adding more exporters or sanitization rules. Observability in NeMo Relay starts with events. Scopes, marks, managed tool calls, managed LLM calls, middleware, and manual lifecycle APIs emit the @@ -69,6 +70,28 @@ Both paths consume the same runtime event contract. Choose plugin-managed export for reusable process configuration, and choose manual APIs only when the caller needs direct control over registration or lifetime. +## Delivery and Shutdown + +All built-in exporters begin with the same asynchronous native subscriber +dispatcher. An event-producing API can return before the exporter subscriber +runs. The subscriber flush API waits for callbacks queued before that barrier, +but some exporters then hand work to their own stream worker, trajectory +builder, batch processor, or remote backend. + +Use the barrier that matches how you manage the exporter lifecycle: + +- For a custom subscriber, call the subscriber flush API before inspecting + callback side effects. +- For a manual exporter, call its documented `force_flush()`, `export()`, or + `shutdown()` operation. +- For plugin-managed export, call `plugin.clear()` or + `clear_plugin_configuration()` during graceful shutdown. + +If the process terminates before those operations complete, the output can omit +queued telemetry even when the instrumented application work completed +successfully. Refer to [Subscribers](/about-nemo-relay/concepts/subscribers#waiting-for-delivery) +for the complete delivery-barrier contract. + ## Use Observability When Start here when you need to: @@ -101,10 +124,14 @@ guardrails before exporters receive sensitive payloads. Your first exporter path works correctly when: - The instrumented request still completes successfully. -- The chosen exporter receives scope data for that request, plus tool or LLM - lifecycle data when the request exercises those paths. +- After the appropriate exporter or plugin barrier returns, the output contains + scope data for that request, plus tool or LLM lifecycle data when the request + exercises those paths. - The exported output shows the same NeMo Relay scope hierarchy you expect from that request. +Application success alone does not prove telemetry delivery because exporter +work is downstream from runtime execution. + If that basic check fails, use the [Trace Incident Runbook](/resources/troubleshooting/trace-incident-runbook) before adding more exporters or extra config layers. diff --git a/docs/configure-plugins/observability/atif.mdx b/docs/configure-plugins/observability/atif.mdx index 096282b50..3acb8b227 100644 --- a/docs/configure-plugins/observability/atif.mdx +++ b/docs/configure-plugins/observability/atif.mdx @@ -217,7 +217,10 @@ single-file ATIF v1.7 artifact. The plugin writes each trajectory when its top-level Agent scope or supported coding-agent turn scope closes. If the plugin is cleared while that root scope -is still open, teardown flushes the partial trajectory. +is still open, teardown first drains queued subscriber callbacks and then +writes the partial trajectory. Terminating the process before root-scope +closure or plugin teardown completes can leave the trajectory absent or +incomplete. To correlate ATIF with OpenTelemetry or OpenInference traces from the same run, join on NeMo Relay UUIDs. The plugin-managed ATIF `session_id` is the @@ -375,7 +378,12 @@ Ok(()) ## Manual API Use the manual `AtifExporter` API when you need explicit collection boundaries -or one exporter object per run. +or one exporter object per run. `export()` and `export_json()` first wait for +subscriber callbacks queued before the call, then take a snapshot of the +collected event history. A separate subscriber flush immediately before export +is unnecessary. The snapshot can also include events delivered after the +internal flush returns but before the snapshot is taken. Events delivered after +the snapshot are not included. @@ -416,7 +424,7 @@ try { ```rust -use nemo_relay::api::subscriber::{deregister_subscriber, flush_subscribers, register_subscriber}; +use nemo_relay::api::subscriber::{deregister_subscriber, register_subscriber}; use nemo_relay::observability::atif::{AtifAgentInfo, AtifExporter}; fn main() -> Result<(), Box> { @@ -434,7 +442,6 @@ fn main() -> Result<(), Box> { // Run instrumented application work here. - flush_subscribers()?; let trajectory = exporter.export()?; let trajectory_json = serde_json::to_string_pretty(&trajectory)?; println!("{trajectory_json}"); @@ -456,6 +463,10 @@ fn main() -> Result<(), Box> { - Tool definitions or `extra` metadata are not JSON-compatible. - The application never opens a top-level Agent scope or a supported coding-agent turn scope, so no trajectory file is created. +- The application expects manual trajectory output without calling `export()` + or `export_json()`. +- The process terminates before manual `export()`, supported root-scope closure, + or plugin teardown establishes the export boundary. - `storage[i].type` is unknown or `storage[i].bucket` is empty for some entry. - `storage` is non-empty in a build that was compiled without the `atif-storage` feature. diff --git a/docs/configure-plugins/observability/atof.mdx b/docs/configure-plugins/observability/atof.mdx index f3013e665..b37447fca 100644 --- a/docs/configure-plugins/observability/atof.mdx +++ b/docs/configure-plugins/observability/atof.mdx @@ -66,7 +66,7 @@ The following table describes the top-level ATOF settings: | Field | Default | Notes | |---|---|---| | `enabled` | `false` | Must be `true` to export events. | -| `sinks` | `[]` | File and stream destinations. An enabled ATOF section requires at least one sink, and every sink receives each event. | +| `sinks` | `[]` | File and stream destinations. An enabled ATOF section requires at least one sink. Each event delivered to the ATOF subscriber is sent to every configured sink. | ## File Sinks @@ -80,6 +80,20 @@ table describes each file sink: | `filename` | Timestamped `nemo-relay-events-*.jsonl` | Output filename. | | `mode` | `append` | `append` or `overwrite`. | + +Native event-producing APIs return before the ATOF subscriber necessarily +runs. After it receives an event, a file sink writes one complete JSONL record +and flushes the process writer. A subscriber flush waits for the file callback +to complete its write attempt for every event queued before the barrier. A +manual exporter reports any stored file errors during `force_flush()` or +`shutdown()`. Plugin teardown also reports stored file errors. + +Flushing the process writer does not call `fsync` or guarantee durability after +a machine crash. If the process terminates before the subscriber receives an +event, the corresponding record can still be lost. + + + ## Stream Sinks Each stream sink receives the same canonical ATOF event that file sinks write. @@ -98,7 +112,7 @@ The following table describes each stream sink: | `transport` | `http_post` | `http_post`, `websocket`, or `ndjson`. | | `headers` | `{}` | String-to-string headers for requests or handshakes. | | `header_env` | `{}` | Maps each header name to an environment variable that contains the full header value. Each variable must be set and nonblank, and the same header cannot also appear in `headers`. | -| `timeout_millis` | `3000` | Per-sink timeout. Must be greater than `0`. | +| `timeout_millis` | `3000` | Per-sink timeout for stream I/O and flush or close acknowledgement. Must be greater than `0`. | | `field_name_policy` | `preserve` | Field-name handling before Relay sends stream events. Accepted values are `preserve` and `replace_dots`. | Use `header_env` for secrets so the configuration stores the environment @@ -120,9 +134,23 @@ The following transports are supported: - `ndjson` opens one long-lived HTTP upload and writes each event as one newline-delimited JSON record. -`force_flush()` flushes file output and drains queued stream events while -keeping stream connections open. `shutdown()` is terminal: it flushes pending -work, closes the connections, and ignores events that arrive later. +For a file sink, the shared subscriber flush completes the synchronous file +callback. For a stream sink, it only confirms that the ATOF callback queued the +stream work. A manual exporter's `force_flush()` first waits for queued +subscriber callbacks, then flushes file output or asks each stream worker to +drain its queued events while keeping the exporter open. It waits for each +stream worker up to that sink's `timeout_millis` value. If a worker does not +acknowledge in time, NeMo Relay logs a warning and `force_flush()` can still +return successfully; the worker can still have queued work. + +`shutdown()` follows the same delivery barrier, flushes file output, and asks +each stream worker to drain and close within the configured timeout. It ignores +events delivered after the exporter closes. A timeout is logged and does not by +itself cause `shutdown()` to return an error. + +During graceful shutdown, `plugin.clear()` or `clear_plugin_configuration()` +uses that same terminal stream-sink operation. It does not guarantee that a +worker which timed out drained or closed. ## Migrate From Version 1 @@ -138,15 +166,16 @@ Make the following changes when you migrate a plugin configuration: - Keep static request headers in `headers`, or use `header_env` to read complete header values from environment variables at activation time. -The plugin sends every event to every configured sink. The manual -`AtofExporter` API is different: each exporter owns one file or stream sink. +The plugin fans out each event delivered to its ATOF subscriber to every +configured sink. The manual `AtofExporter` API is different: each exporter owns +one file or stream sink. ## Expected Output -Each file sink writes an emitted scope, tool, LLM, middleware, or mark event as -one ATOF JSON object per line. Stream sinks encode the event as described in -[Stream Sinks](#stream-sinks). For event field semantics, refer to -[Events](/about-nemo-relay/concepts/events). +Each file sink writes a scope, tool, LLM, middleware, or mark event delivered to +the exporter as one ATOF JSON object per line. Stream sinks encode the delivered +event as described in [Stream Sinks](#stream-sinks). For event field semantics, +refer to [Events](/about-nemo-relay/concepts/events). Coding-agent sessions emit a `session.start` mark with the logical `metadata.session_id`, a Relay-owned UUIDv7 `metadata.session_instance_id`, and @@ -167,6 +196,8 @@ annotation. The request used for provider execution remains unchanged. Refer to Register the plugin before instrumented work starts. Clear it during shutdown so every sink flushes pending events and stream connections close cleanly. +Terminating the process before subscriber delivery or plugin teardown completes +can leave the final records missing. ## Plugin Configuration @@ -426,3 +457,5 @@ Common configuration and runtime issues include: configured stream sink. - A file sink is configured in a target that cannot access the native filesystem. +- The process terminates before a subscriber or exporter barrier, so the file + or stream destination is missing queued tail records. diff --git a/docs/configure-plugins/observability/configuration.mdx b/docs/configure-plugins/observability/configuration.mdx index 27f5e280f..4eca73d5b 100644 --- a/docs/configure-plugins/observability/configuration.mdx +++ b/docs/configure-plugins/observability/configuration.mdx @@ -161,6 +161,7 @@ The following table describes how each failure affects application work: | ATIF dispatcher serialization or subscriber-management failure | The ATIF dispatcher records a fatal exporter error and stops observing later ATIF events. Other observability sections continue to run. | | OpenTelemetry or OpenInference construction failure | Plugin initialization fails before the subscriber is registered. | | OpenTelemetry or OpenInference export failure after registration | Application work continues. The OTLP exporter reports failures through its runtime logging and flush or shutdown path. | +| Process termination before plugin teardown completes | The output can omit events that were still in the subscriber queue or an exporter-owned worker. NeMo Relay does not backfill those events after restart, and the process can end before an exporter reports an error. | Missing or delayed telemetry is represented as absence of exporter output, not as synthetic success or failure events. NeMo Relay does not backfill events for @@ -438,9 +439,20 @@ filename templates, unknown fields according to policy, and enabled exporters that are unavailable in the current build or target. Call `plugin.clear()` or `clear_plugin_configuration()` during teardown. -Clearing the plugin configuration deregisters inferred subscribers, flushes file -exporters, drains and closes ATOF stream sinks, and shuts down owned OTLP -subscribers. +Clearing the plugin configuration waits for previously queued subscriber +callbacks, runs exporter-specific shutdown work, and removes plugin-owned +registrations. That work flushes file exporters, asks ATOF stream sink workers +to drain and close within their configured timeout, writes any eligible partial +ATIF trajectories, and shuts down owned OTLP subscribers. If an ATOF stream +worker times out, NeMo Relay logs a warning and continues teardown; clearing +the plugin can still return successfully without confirming that worker drained +or closed. + +This is a graceful teardown barrier, not a crash-durability guarantee. If the +process terminates before clearing begins or finishes, the output can omit +telemetry for application work that completed successfully. Telemetry delivery +failures remain fail-open and do not retroactively change the application +result. Use manual subscriber/exporter APIs instead of the plugin when you need custom subscriber names, explicit per-run exporter objects, or direct control over the diff --git a/docs/configure-plugins/observability/openinference.mdx b/docs/configure-plugins/observability/openinference.mdx index e4108396b..0a5ff9688 100644 --- a/docs/configure-plugins/observability/openinference.mdx +++ b/docs/configure-plugins/observability/openinference.mdx @@ -158,6 +158,11 @@ the alias. Redact sensitive event payloads with sanitize guardrails before production export. +Register the plugin before the first instrumented request and clear it during +graceful shutdown. Calling subscriber flush delivers queued events to the +tracing subscriber, but it does not necessarily drain the tracer provider's +export processor. + ## Plugin Configuration Use plugin configuration when the application should let NeMo Relay own the @@ -309,7 +314,11 @@ Ok(()) ## Manual API Use the manual subscriber API when you need an explicit subscriber name or -direct `force_flush` control. +direct control over when spans are exported. `force_flush()` first waits for +shared subscriber delivery and then asks the tracer provider to export finished +spans. `shutdown()` waits for shared subscriber work and then shuts down the +provider. During plugin teardown, `plugin.clear()` or +`clear_plugin_configuration()` invokes the owned subscriber shutdown path. @@ -394,3 +403,6 @@ fn main() -> Result<(), Box> { - The OpenInference feature is unavailable in the current build or target. - Tool and LLM calls do not use managed helpers, so spans contain only scope lifecycle data. +- Abrupt termination can discard events that are still in the shared dispatcher + or spans that are still in the provider queue. Backend visibility also + depends on transport success, timeouts, and backend processing. diff --git a/docs/configure-plugins/observability/opentelemetry.mdx b/docs/configure-plugins/observability/opentelemetry.mdx index 6c434c9cd..48ed44383 100644 --- a/docs/configure-plugins/observability/opentelemetry.mdx +++ b/docs/configure-plugins/observability/opentelemetry.mdx @@ -149,8 +149,10 @@ discrete attributes. Refer to for the full mapping. Register the plugin before the first instrumented request, use stable service -identity fields, keep credentials outside source code, and flush during -graceful shutdown. +identity fields, keep credentials outside source code, and clear the plugin +during graceful shutdown. Calling subscriber flush delivers queued events to +the tracing subscriber, but it does not necessarily drain the tracer provider's +export processor. ## Plugin Configuration @@ -303,7 +305,11 @@ Ok(()) ## Manual API Use the manual subscriber API when you need an explicit subscriber name or -direct `force_flush` control. +direct control over when spans are exported. `force_flush()` first waits for +shared subscriber delivery and then asks the tracer provider to export finished +spans. `shutdown()` waits for shared subscriber work and then shuts down the +provider. During plugin teardown, `plugin.clear()` or +`clear_plugin_configuration()` invokes the owned subscriber shutdown path. @@ -384,3 +390,6 @@ fn main() -> Result<(), Box> { - Headers or resource attributes are not string-to-string maps. - The exporter feature is unavailable in the current build or target. - The endpoint is unreachable at runtime. +- Abrupt termination can discard events that are still in the shared dispatcher + or spans that are still in the provider queue. Backend visibility also + depends on transport success, timeouts, and backend processing. diff --git a/docs/contribute/runtime-contract-docs.mdx b/docs/contribute/runtime-contract-docs.mdx index ee1a0022c..748751adf 100644 --- a/docs/contribute/runtime-contract-docs.mdx +++ b/docs/contribute/runtime-contract-docs.mdx @@ -27,7 +27,8 @@ These pages define or route the shared runtime model: - `docs/about-nemo-relay/concepts/events.mdx` - `docs/about-nemo-relay/concepts/framework-integrations.mdx` - `docs/about-nemo-relay/concepts/codecs.mdx` -- `docs/getting-started/agent-runtime-primer.mdx` +- `docs/about-nemo-relay/agent-runtime-primer.mdx` +- `docs/getting-started/about.mdx` - `docs/getting-started/installation.mdx` - `docs/getting-started/quick-start/index.mdx` - `docs/reference/api/index.mdx` diff --git a/docs/getting-started/about.mdx b/docs/getting-started/about.mdx new file mode 100644 index 000000000..1e73b0ccc --- /dev/null +++ b/docs/getting-started/about.mdx @@ -0,0 +1,100 @@ +--- +title: "About" +description: "Choose the right NeMo Relay getting-started path for your agent system." +position: 1 +--- +{/* SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +SPDX-License-Identifier: Apache-2.0 */} + +Use Getting Started to choose the smallest NeMo Relay path that matches the +boundary you want to observe or control. Relay can run beside a local +coding-agent harness, wrap application-owned tool or LLM calls, connect through +a maintained framework integration, or activate reusable plugin behavior from +configuration. + + + + +Install [NVIDIA agent skills](https://docs.nvidia.com/skills) to give +compatible coding agents product-specific Relay workflows and references. After +the NeMo Relay skills are published in `nvidia/skills`, install the full Relay +skill set without cloning the catalog: + +```bash +npx skills add nvidia/skills \ + --skill nemo-relay-install \ + --skill nemo-relay-get-started \ + --skill nemo-relay-instrument-calls \ + --skill nemo-relay-instrument-context-isolation \ + --skill nemo-relay-instrument-typed-wrappers \ + --skill nemo-relay-plugin-observability \ + --skill nemo-relay-plugin-adaptive-tuning \ + --skill nemo-relay-plugin-build \ + --skill nemo-relay-migrate-from-flow \ + --skill nemo-relay-debug-runtime-integration \ + --agent codex +``` + +Use `--agent claude-code`, `--agent cursor`, or another supported agent target +when that is where you want the skill installed. To inspect the current NVIDIA +skill catalog first, run: + +```bash +npx skills add nvidia/skills --list +``` + +Agent skills help your coding assistant apply Relay documentation and workflows. +They do not replace installing NeMo Relay packages, configuring runtime +plugins, or validating emitted events. + + + +## Find the Right Starting Point + +Start with the fastest path that can prove Relay is working, then move toward +more involved integration and development workflows only when you need them. + +- **Try now with the CLI:** Use [NeMo Relay CLI](/nemo-relay-cli/about) when a + local Claude Code, Codex, or Hermes coding-agent harness owns invocation. + Relay can observe hooks, gateway-routed model traffic, and exporter output + without changing application code. +- **Use a supported integration:** Start with + [Supported Integrations](/supported-integrations/about) when LangChain, + LangGraph, Deep Agents, or OpenClaw already owns scheduling, callbacks, tool + execution, provider payloads, or agent lifecycle hooks. +- **Instrument application-owned calls:** Start with + [Quick Start](/getting-started/quick-start), then use + [Instrument Applications](/instrument-applications/about) when your Python, + Node.js, or Rust application owns the tool or LLM call site. +- **Build a custom framework or provider integration:** Use + [Integrate into Frameworks](/integrate-into-frameworks/about) when a new + framework, orchestration layer, SDK, or provider adapter owns scheduling, + retries, callbacks, and provider payloads. +- **Package reusable runtime behavior:** Use + [Build Plugins](/build-plugins/about) for shared middleware, subscribers, + exporters, model pricing, policy, or adaptive behavior. Use + [Observability](/configure-plugins/observability/about) for export setup and + [Adaptive](/configure-plugins/adaptive/about) after baseline instrumentation + is working. + +## Getting Started Process + +Use the pages in this section in order when you are setting up Relay for the +first time: + +1. Review [Prerequisites](/getting-started/prerequisites) for language, + package, and repository tooling. +2. Use [Installation](/getting-started/installation) to install the CLI, + language binding, or integration package that matches your boundary. +3. Use [Configuration](/getting-started/configuration) when you need runtime, + plugin, observability, or adaptive setup guidance. +4. Use [Quick Start](/getting-started/quick-start) to run the smallest + working path and confirm Relay emits the expected output. + +If you need the conceptual model before setup, read +[Agent Runtime Primer](/about-nemo-relay/agent-runtime-primer) and +[Concepts](/about-nemo-relay/concepts). diff --git a/docs/getting-started/installation.mdx b/docs/getting-started/installation.mdx index 2b3086de6..4f11aa039 100644 --- a/docs/getting-started/installation.mdx +++ b/docs/getting-started/installation.mdx @@ -194,7 +194,7 @@ Install the Node.js package when your application uses NeMo Relay through the JavaScript API and owns the tool or LLM call boundary. ```bash -npm install nemo-relay-node +npm install nemo-relay-node@0.6.0 ``` ## Rust diff --git a/docs/getting-started/quick-start/index.mdx b/docs/getting-started/quick-start/index.mdx index 29e61b57c..972debb96 100644 --- a/docs/getting-started/quick-start/index.mdx +++ b/docs/getting-started/quick-start/index.mdx @@ -15,9 +15,42 @@ or control: boundary. If you still need the shared runtime vocabulary for scopes, middleware, events, -or plugins, refer to the [Agent Runtime Primer](/getting-started/agent-runtime-primer) +or plugins, refer to the [Agent Runtime Primer](/about-nemo-relay/agent-runtime-primer) or [Concepts](/about-nemo-relay/concepts) first. + + + +Install [NVIDIA agent skills](https://docs.nvidia.com/skills) to give +compatible coding agents product-specific Relay workflows and references. After +the NeMo Relay skills are published in `nvidia/skills`, install the full Relay +skill set without cloning the catalog: + +```bash +npx skills add nvidia/skills \ + --skill nemo-relay-install \ + --skill nemo-relay-get-started \ + --skill nemo-relay-instrument-calls \ + --skill nemo-relay-instrument-context-isolation \ + --skill nemo-relay-instrument-typed-wrappers \ + --skill nemo-relay-plugin-observability \ + --skill nemo-relay-plugin-adaptive-tuning \ + --skill nemo-relay-plugin-build \ + --skill nemo-relay-migrate-from-flow \ + --skill nemo-relay-debug-runtime-integration \ + --agent codex +``` + +Use `--agent claude-code`, `--agent cursor`, or another supported agent target +when that is where you want the skill installed. To inspect the current NVIDIA +skill catalog first, run `npx skills add nvidia/skills --list`. + + + ## Examples by Integration Layer Use this map when you know what Relay boundary you want to verify, but you do diff --git a/docs/getting-started/quick-start/nodejs.mdx b/docs/getting-started/quick-start/nodejs.mdx index dacd01a7d..8f47996be 100644 --- a/docs/getting-started/quick-start/nodejs.mdx +++ b/docs/getting-started/quick-start/nodejs.mdx @@ -19,7 +19,7 @@ local checkout. Use this path when you want the published package for application development. ```bash -npm install nemo-relay-node +npm install nemo-relay-node@0.6.0 ``` ### Install from the Repository diff --git a/docs/integrate-into-frameworks/provider-codecs.mdx b/docs/integrate-into-frameworks/provider-codecs.mdx index 90df6c7f6..4012984a2 100644 --- a/docs/integrate-into-frameworks/provider-codecs.mdx +++ b/docs/integrate-into-frameworks/provider-codecs.mdx @@ -13,7 +13,7 @@ Use this guide when a framework integration needs NeMo Relay middleware, interce You will attach request and response codecs to a managed LLM wrapper so that: -- Request intercepts can work with normalized messages, model names, tools, generation parameters, and provider-specific extras +- Request intercepts can work with normalized instructions, messages, model names, tools, generation parameters, and tagged provider-specific fields - The provider callback still receives the provider payload that the framework expects - Response subscribers can receive normalized response annotations without changing the caller-visible provider response @@ -52,9 +52,37 @@ annotation, Relay rejects the outcome before creating the LLM lifecycle. When no request codec is active, the raw request remains fully writable and is the provider-visible source of truth. +Gateway interceptors must preserve the request's effective `stream` mode. The +gateway chooses buffered or streaming response handling before request +intercepts run and rejects changes between those modes before contacting the +provider. + When a managed LLM call has a response codec, NeMo Relay decodes the raw provider response for observability and attaches the result to the emitted LLM end event. The response codec does not rewrite the value returned to the application. Use [Provider Response Codecs](/integrate-into-frameworks/provider-response-codecs) for response-only behavior and custom response codec examples. -Codec implementations should preserve fields they do not understand. Treat `encode` as a merge operation over the original provider payload, not as a full replacement. +The built-in codecs guarantee +`encode(decode(original), original) == original` at the JSON-value level. Their +encoders compare an edited annotation with a decoded baseline and patch only +changed fields. As a result, unchanged string-versus-block content, explicit +nulls, legacy field names, ordered input items, metadata, and unknown fields +remain untouched. Removing an unknown top-level key from `extra` removes it +from the provider payload; adding or changing a key overlays it. + +Use the annotated request surfaces according to their ownership: + +- `instructions` for Anthropic `system` or OpenAI Responses `instructions`. +- Portable `messages`, content parts, function calls and results, `tools`, and + `tool_choice` when the component has shared semantics. +- Tagged `api_specific` fields for modeled controls that belong to Anthropic + Messages, OpenAI Chat Completions, or OpenAI Responses. +- `{ provider, kind, value }` native components for provider-only input items, + blocks, tools, tool choices, and future union members. `value` is the exact + provider JSON. +- Top-level `extra` only for unknown future fields. + +Provider-native components are intentionally surface-bound. If an intercept +places an Anthropic native block in an OpenAI request, or makes a portable edit +that the target API cannot encode, Relay returns an explicit error before the +provider callback. ## Built-in Provider Codecs @@ -81,6 +109,11 @@ The built-in provider codecs expose the same core methods: Choose the provider codec that matches the payload shape the framework already sends to the provider. Do not translate to a different provider shape only to make the codec fit. +The `nemo-relay` gateway selects these request codecs automatically for +`/v1/messages`, `/v1/chat/completions`, and `/v1/responses`, for both buffered +and streaming calls. Count-token, model, probe, and non-LLM passthrough routes +do not use request codecs. + ## Example: Add a System Message with a Provider Codec This example uses a request intercept to edit the normalized request. The codec writes the edited messages back into the provider payload before the provider callback runs. @@ -252,6 +285,44 @@ let response = llm_call_execute( +## Example: Edit Provider-Specific and Native Data + +The tagged `api_specific` object is mutable. Python getters return JSON values, +so retrieve, edit, and reassign the value. The same reassignment pattern works +for a provider-native component inside `messages`. + +```python +def tune_responses_request(_name, request, annotated): + if annotated is None: + return nemo_relay.LLMRequestInterceptOutcome(request) + + api_specific = dict(annotated.api_specific or {}) + if api_specific.get("api") == "openai_responses": + api_specific["background"] = True + api_specific["stream_options"] = {"include_obfuscation": False} + annotated.api_specific = api_specific + + messages = list(annotated.messages) + for index, item in enumerate(messages): + if ( + item.get("role") == "provider_native" + and item.get("provider") == "openai_responses" + and item.get("kind") == "reasoning" + ): + native_value = dict(item["value"]) + native_value["encrypted_content"] = "updated-ciphertext" + messages[index] = {**item, "value": native_value} + annotated.messages = messages + + # request.headers remains mutable; request.content is read-only here. + request.headers["x-request-policy"] = "codec-aware" + return nemo_relay.LLMRequestInterceptOutcome(request, annotated) +``` + +Native values are not normalized across providers. Keep the `provider` tag +unchanged unless you also replace the entire component with a representation +that the target codec owns. + ## Example: Write a Custom Framework Codec Use a custom codec when a framework uses a payload shape that does not directly match a built-in provider format. The codec decodes the framework shape into `AnnotatedLLMRequest`, and encodes edits back into the framework shape. @@ -369,11 +440,33 @@ Use this checklist to confirm the implementation preserves the expected runtime contract. - Intercepts receive `annotated` only when the managed call supplies a request codec. -- `encode` preserves provider fields that the annotated model does not represent. +- An unchanged annotation encodes to the exact original JSON value. +- Portable edits preserve adjacent provider-native components and unknown fields. +- Provider-specific edits use `api_specific`; `extra` is limited to unknown future top-level fields. +- Provider-mismatched native components fail before the provider callback. - Response codecs are used only for event annotations, not caller-visible response rewriting. - Codec implementations are pure data transforms and do not perform provider I/O. - Framework-owned clients, sockets, streams, callbacks, and file handles stay outside codec results. +## Migrate Raw Gateway Body Mutation + +The 0.6 request annotation expansion uses +`nemo.relay.AnnotatedLlmRequest@2` and +`nemo.relay.LlmRequestInterceptOutcome@2` at dynamic plugin boundaries. Rebuild +Rust native plugins and Rust `grpc-v1` workers against NeMo Relay 0.6, and +update Python workers to the 0.6 worker SDK. A dynamic plugin that registers an +LLM request intercept must set `compat.relay = ">=0.6,<1.0"`, or another range +that excludes Relay 0.5. Relay enforces this floor when the intercept is +registered. + +Gateway interceptors that mutate provider request bodies must now edit the +annotation on the three generation routes. For example, replace +`request.content["messages"] = messages` with +`annotated.messages = messages`, and return the edited annotation with the raw +request. Keep transport-only changes in `request.headers`. Relay rejects a raw +body edit while the route codec is active, so the provider never receives an +ambiguous mix of raw and annotated changes. + ## Next Steps Use these links to continue from this workflow into the next related task. diff --git a/docs/nemo-relay-cli/codex.mdx b/docs/nemo-relay-cli/codex.mdx index 65b86222e..f0725dec3 100644 --- a/docs/nemo-relay-cli/codex.mdx +++ b/docs/nemo-relay-cli/codex.mdx @@ -343,9 +343,19 @@ discover exactly one enabled, trusted handler for every generated event. `PostToolUseFailure`, `Notification`, and `SessionEnd` are not in the Codex 0.143 plugin hook schema, so Relay does not generate undiscoverable handlers for them. Relay maps delivered events to agent, turn, subagent, tool, and mark lifecycle -events; prompt and stop payloads also provide LLM correlation context. Complete -standalone `cat`, `bat`, `batcat`, or PowerShell `Get-Content` reads of -`SKILL.md` emit observed `skill.load` marks. +events; prompt and stop payloads also provide LLM correlation context. + +Relay emits an observed `skill.load` mark when it sees a complete file read +through `cat`, `bat`, `batcat`, or PowerShell `Get-Content`. The final path +component must equal `SKILL.md` under an ASCII case-insensitive comparison. +Relay also recognizes one supported wrapper around that reader: `sh -c`, +`bash -c`, `bash -lc`, `zsh -c`, `zsh -lc`, `fish -c`, +`fish --command=...`, `powershell -Command`, or `pwsh -Command`. Skill-load +detection is best effort. Relay must be able to observe a supported request +shape. If Relay doesn't emit a mark, that doesn't prove that the skill wasn't +loaded. Refer to +[Automatic Skill-Load Marks](/about-nemo-relay/concepts/events#automatic-skill-load-marks) +for the complete criteria. The transparent wrapper passes hook entries as Codex CLI config overrides and sets `features.hooks=true` for that launched process. It also disables the known diff --git a/docs/reference/api/index.mdx b/docs/reference/api/index.mdx index 51e06ab4f..33455f244 100644 --- a/docs/reference/api/index.mdx +++ b/docs/reference/api/index.mdx @@ -19,7 +19,7 @@ semantics such as: - Support boundaries For those topics, start with [Concepts](/about-nemo-relay/concepts), the -[Agent Runtime Primer](/getting-started/agent-runtime-primer), or the +[Agent Runtime Primer](/about-nemo-relay/agent-runtime-primer), or the workflow-specific documentation path that matches the boundary you own. Primary binding references: diff --git a/docs/reference/llm-request-intercept-outcomes.mdx b/docs/reference/llm-request-intercept-outcomes.mdx index f1012553f..5e968a4b8 100644 --- a/docs/reference/llm-request-intercept-outcomes.mdx +++ b/docs/reference/llm-request-intercept-outcomes.mdx @@ -43,9 +43,23 @@ separate response-side path used to attach normalized data to lifecycle events. With an active codec, `request.content` is read-only context. Every intercept must return an annotation and make provider-body changes through that -annotation, including its flattened `extra` fields for provider-specific data. -Relay rejects a changed raw body or missing annotation at the offending -intercept before invoking later middleware or creating an LLM lifecycle. +annotation. Use portable fields such as `messages` and `instructions`, the +tagged `api_specific` surface for modeled provider controls, provider-native +components for provider-only unions, and flattened `extra` only for unknown +future top-level fields. Relay rejects a changed raw body or missing annotation +at the offending intercept before invoking later middleware or creating an LLM +lifecycle. + +The `nemo-relay` gateway enables request codecs for the generation routes +`/v1/messages`, `/v1/chat/completions`, and `/v1/responses`, including streaming +requests. Interceptors written for those routes must migrate provider-body +mutation from `request.content` to `annotated_request`. Count-token, model, +probe, and non-LLM passthrough routes retain raw-body authority. + +For gateway generation routes, preserve the effective `stream` mode in the +annotation. Relay rejects buffered-to-streaming and streaming-to-buffered edits +before the provider callback because response handling is selected from the +original request. The following diagram shows how Relay resolves an intercept outcome before managed execution. @@ -84,7 +98,7 @@ or object shape: - Public C callbacks return one owned canonical outcome JSON string, and native ABI v1 callbacks return one host-owned outcome JSON string. - Rust and Python `grpc-v1` worker SDKs return their canonical outcome in a - `JsonEnvelope` with schema `nemo.relay.LlmRequestInterceptOutcome@1`. + `JsonEnvelope` with schema `nemo.relay.LlmRequestInterceptOutcome@2`. The standalone request-intercept helper returns the complete outcome but does not emit its pending marks because it does not own an LLM lifecycle. @@ -114,6 +128,21 @@ outputs, metadata envelopes, and parallel mark-aware registrations with the canonical outcome and the existing `register_llm_request_intercept` registration name. +The 0.6 annotation expansion changes the request and outcome envelope schemas +to `nemo.relay.AnnotatedLlmRequest@2` and +`nemo.relay.LlmRequestInterceptOutcome@2`. Rebuild Rust native plugins and Rust +`grpc-v1` workers against NeMo Relay 0.6, and update Python workers to the 0.6 +worker SDK. Dynamic plugins that register an LLM request intercept must declare +`compat.relay = ">=0.6,<1.0"`, or another range that excludes Relay 0.5. The +host enforces this compatibility floor during registration. + +For gateway generation requests, replace code such as +`request.content["messages"] = ...` with an edit to +`annotated_request.messages`, then return both the unchanged raw request and the +edited annotation. Header edits continue to use `request.headers`. An attempt +to change the raw body while its request codec is active returns an explicit +error before the provider callback. + ## Related Topics - [Tool Execution Intercept Outcomes](/reference/tool-execution-intercept-outcomes) diff --git a/docs/resources/troubleshooting/trace-incident-runbook.mdx b/docs/resources/troubleshooting/trace-incident-runbook.mdx index dd571d4e6..9a89f4366 100644 --- a/docs/resources/troubleshooting/trace-incident-runbook.mdx +++ b/docs/resources/troubleshooting/trace-incident-runbook.mdx @@ -14,7 +14,7 @@ baseline scope and call instrumentation path. For first-time setup problems, start with the [Troubleshooting Guide](/resources/troubleshooting). For conceptual grounding, -refer to [Agent Runtime Primer](/getting-started/agent-runtime-primer), +refer to [Agent Runtime Primer](/about-nemo-relay/agent-runtime-primer), [Scopes](/about-nemo-relay/concepts/scopes), [Events](/about-nemo-relay/concepts/events), and [Subscribers](/about-nemo-relay/concepts/subscribers). diff --git a/examples/rust-native-plugin/relay-plugin.toml b/examples/rust-native-plugin/relay-plugin.toml index 84ed4fcf3..1ef8916e1 100644 --- a/examples/rust-native-plugin/relay-plugin.toml +++ b/examples/rust-native-plugin/relay-plugin.toml @@ -8,7 +8,7 @@ id = "examples.rust_native_policy" kind = "rust_dynamic" [compat] -relay = ">=0.5,<1.0" +relay = ">=0.6,<1.0" native_api = "1" [defaults] diff --git a/fern/docs.yml b/fern/docs.yml index c59775cfe..6c7d106b5 100644 --- a/fern/docs.yml +++ b/fern/docs.yml @@ -80,6 +80,10 @@ redirects: - source: /nemo/relay/nemo-relay-cli/cursor destination: /nemo/relay/nemo-relay-cli/about +# Getting started IA +- source: /nemo/relay/getting-started/agent-runtime-primer + destination: /nemo/relay/about-nemo-relay/agent-runtime-primer + experimental: mdx-components: - ./components diff --git a/go/nemo_relay/nemo_relay.go b/go/nemo_relay/nemo_relay.go index b8b5280dd..727f8380c 100644 --- a/go/nemo_relay/nemo_relay.go +++ b/go/nemo_relay/nemo_relay.go @@ -1549,7 +1549,10 @@ func DeregisterSubscriber(name string) error { // FlushSubscribers waits for subscriber callbacks queued before this call to // finish. Native event-producing APIs enqueue subscriber work and return -// without waiting for observer callbacks. +// without waiting for observer callbacks. Call this function outside native +// subscriber callbacks. A re-entrant call returns without waiting to avoid +// blocking the dispatcher, so callbacks later in the same dispatch snapshot +// can still run. func FlushSubscribers() error { return checkStatus(C.nemo_relay_flush_subscribers()) } @@ -1827,13 +1830,19 @@ func (e *AtofExporter) Deregister(name string) error { return checkStatus(status) } -// ForceFlush flushes the output file. +// ForceFlush, outside a native subscriber callback, waits for queued subscriber delivery and then +// flushes the configured file sink or asks the configured stream sink to drain up to its timeout. +// A re-entrant call does not establish the delivery barrier. A stream timeout is logged and does +// not by itself return an error. func (e *AtofExporter) ForceFlush() error { status := C.nemo_relay_atof_exporter_force_flush(e.ptr) return checkStatus(status) } -// Shutdown flushes the output file. +// Shutdown, outside a native subscriber callback, waits for queued subscriber delivery and then +// flushes the configured file sink or asks the configured stream sink to drain and close up to its +// timeout. A re-entrant call does not establish the delivery barrier. A stream timeout is logged +// and does not by itself return an error. func (e *AtofExporter) Shutdown() error { status := C.nemo_relay_atof_exporter_shutdown(e.ptr) return checkStatus(status) diff --git a/go/nemo_relay/subscribers/subscribers.go b/go/nemo_relay/subscribers/subscribers.go index e7225cbfc..b6e0c7b38 100644 --- a/go/nemo_relay/subscribers/subscribers.go +++ b/go/nemo_relay/subscribers/subscribers.go @@ -42,8 +42,8 @@ func Deregister(name string) error { return nemo_relay.DeregisterSubscriber(name) } -// Flush waits for subscriber callbacks queued before this call to finish. This -// is a shorthand for [nemo_relay.FlushSubscribers]. +// Flush has the same completion semantics and re-entrant limitation as +// [nemo_relay.FlushSubscribers]. func Flush() error { return nemo_relay.FlushSubscribers() } diff --git a/python/nemo_relay/_native.pyi b/python/nemo_relay/_native.pyi index fac3ed23d..aaec3027a 100644 --- a/python/nemo_relay/_native.pyi +++ b/python/nemo_relay/_native.pyi @@ -30,6 +30,7 @@ _JsonPrimitive: TypeAlias = str | int | float | bool | None _JsonValue: TypeAlias = _JsonPrimitive | list["_JsonValue"] | dict[str, "_JsonValue"] _JsonObject: TypeAlias = dict[str, _JsonValue] _Json: TypeAlias = _JsonValue +_MessageContent: TypeAlias = str | Sequence[Mapping[str, _JsonValue]] class _EventSanitizeFields(TypedDict): data: _Json | None @@ -463,21 +464,25 @@ class AnnotatedLLMRequest: self, messages: Sequence[Mapping[str, _JsonValue]], *, + instructions: Optional[_MessageContent] = None, model: Optional[str] = None, params: Optional[Mapping[str, _JsonValue]] = None, tools: Optional[Sequence[Mapping[str, _JsonValue]]] = None, tool_choice: Optional[str | Mapping[str, _JsonValue]] = None, + api_specific: Optional[Mapping[str, _JsonValue]] = None, extra: Optional[Mapping[str, _JsonValue]] = None, ) -> None: """Create a normalized LLM request view. Args: messages: Provider-normalized message objects. + instructions: Optional provider-level instructions. model: Optional model name. params: Optional provider parameters. tools: Optional tool declarations. tool_choice: Optional tool-selection directive. - extra: Optional provider-specific fields. + api_specific: Optional tagged provider-specific request fields. + extra: Optional unknown future top-level fields. Returns: ``None``. @@ -496,6 +501,14 @@ class AnnotatedLLMRequest: """Replace normalized message objects.""" ... @property + def instructions(self) -> Optional[_MessageContent]: + """Return provider-level instructions, if present.""" + ... + @instructions.setter + def instructions(self, value: Optional[_MessageContent]) -> None: + """Set or clear provider-level instructions.""" + ... + @property def model(self) -> Optional[str]: """Return the normalized model name, if present.""" ... @@ -528,12 +541,20 @@ class AnnotatedLLMRequest: """Set or clear the normalized tool-choice directive.""" ... @property + def api_specific(self) -> Optional[_JsonObject]: + """Return tagged provider-specific request fields, if present.""" + ... + @api_specific.setter + def api_specific(self, value: Optional[Mapping[str, _JsonValue]]) -> None: + """Set or clear tagged provider-specific request fields.""" + ... + @property def extra(self) -> _JsonObject: - """Return provider-specific request fields.""" + """Return unknown future top-level request fields.""" ... @extra.setter def extra(self, value: Mapping[str, _JsonValue]) -> None: - """Replace provider-specific request fields.""" + """Replace unknown future top-level request fields.""" ... def system_prompt(self) -> Optional[str]: """Return the first normalized system prompt, if one is present.""" @@ -891,10 +912,24 @@ class AtofExporter: """Deregister ``name`` and return whether it existed.""" ... def force_flush(self) -> None: - """Flush the output file.""" + """Flush the exporter. + + Outside a native subscriber callback, wait for queued subscriber + delivery, then flush the file sink or ask the stream sink to drain up + to its timeout. A re-entrant call does not establish the delivery + barrier. A stream timeout is logged and does not by itself return an + error. + """ ... def shutdown(self) -> None: - """Flush the output file before shutdown.""" + """Flush the exporter and shut it down. + + Outside a native subscriber callback, wait for queued subscriber + delivery, then flush the file sink or ask the stream sink to drain and + close up to its timeout. A re-entrant call does not establish the + delivery barrier. A stream timeout is logged and does not by itself + return an error. + """ ... class ScopeStack: @@ -1947,7 +1982,11 @@ def deregister_subscriber(name: str) -> bool: ... def flush_subscribers() -> None: - """Wait for subscriber callbacks already queued by native event emission.""" + """Wait for subscriber callbacks queued by native event emission. + + Call this function outside subscriber callbacks. A re-entrant call returns + without waiting, so callbacks later in the same dispatch snapshot can run. + """ ... def scope_register_tool_sanitize_request_guardrail( diff --git a/python/nemo_relay/subscribers.py b/python/nemo_relay/subscribers.py index 24476504b..b54b42e2e 100644 --- a/python/nemo_relay/subscribers.py +++ b/python/nemo_relay/subscribers.py @@ -93,6 +93,10 @@ def flush() -> None: Native NeMo Relay event APIs enqueue subscriber callbacks and return without waiting for observer work. Use this barrier in tests and shutdown paths when captured subscriber output must be complete before continuing. + + Call this function outside subscriber callbacks. A re-entrant call returns + without waiting to avoid blocking the dispatcher, so callbacks later in the + same dispatch snapshot can still run. """ return _native_flush() diff --git a/python/plugin/src/nemo_relay_plugin/_api.py b/python/plugin/src/nemo_relay_plugin/_api.py index b688addc8..b8facbd76 100644 --- a/python/plugin/src/nemo_relay_plugin/_api.py +++ b/python/plugin/src/nemo_relay_plugin/_api.py @@ -106,8 +106,8 @@ class EventSanitizeFields(TypedDict): JSON_SCHEMA = "nemo.relay.Json@1" EVENT_SCHEMA = "nemo.relay.Event@1" LLM_REQUEST_SCHEMA = "nemo.relay.LlmRequest@1" -ANNOTATED_LLM_REQUEST_SCHEMA = "nemo.relay.AnnotatedLlmRequest@1" -LLM_REQUEST_INTERCEPT_OUTCOME_SCHEMA = "nemo.relay.LlmRequestInterceptOutcome@1" +ANNOTATED_LLM_REQUEST_SCHEMA = "nemo.relay.AnnotatedLlmRequest@2" +LLM_REQUEST_INTERCEPT_OUTCOME_SCHEMA = "nemo.relay.LlmRequestInterceptOutcome@2" TOOL_EXECUTION_INTERCEPT_OUTCOME_SCHEMA = "nemo.relay.ToolExecutionInterceptOutcome@1" PLUGIN_DIAGNOSTICS_SCHEMA = "nemo.relay.PluginDiagnostics@1" _OBJECT_SCHEMAS = frozenset( diff --git a/python/tests/plugin/test_worker_sdk.py b/python/tests/plugin/test_worker_sdk.py index 7ecafbeea..5a521fb1e 100644 --- a/python/tests/plugin/test_worker_sdk.py +++ b/python/tests/plugin/test_worker_sdk.py @@ -1296,7 +1296,7 @@ def legacy_result(name: str, value: Json, next_call: ToolNext) -> Any: pb.LLM_REQUEST_INTERCEPT, llm=_llm_payload(request={"content": {}}, annotated={"messages": []}), ), - "expected 'nemo.relay.AnnotatedLlmRequest@1'", + "expected 'nemo.relay.AnnotatedLlmRequest@2'", ), ], ) diff --git a/python/tests/test_builtin_codecs.py b/python/tests/test_builtin_codecs.py index 96fc0b2ab..99beb9054 100644 --- a/python/tests/test_builtin_codecs.py +++ b/python/tests/test_builtin_codecs.py @@ -115,6 +115,72 @@ def test_openai_chat_encode(self): assert cast(float, encoded_content["temperature"]) == 0.7 assert len(cast(list[JsonObject], encoded_content["messages"])) == 2 + def test_anthropic_issue_501_roundtrip_and_annotated_edit(self): + """Anthropic cache blocks survive unchanged and edited Python annotations.""" + codec = AnthropicMessagesCodec() + original = LLMRequest( + {}, + { + "model": "claude-sonnet-4-20250514", + "max_tokens": 128, + "system": [ + { + "type": "text", + "text": "Keep this cached.", + "cache_control": {"type": "ephemeral"}, + } + ], + "messages": [ + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "toolu_1", + "name": "lookup", + "input": {"q": "x"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_1", + "content": "done", + } + ], + }, + ], + "future_field": None, + }, + ) + + annotated = codec.decode(original) + assert annotated.instructions == [ + { + "type": "text", + "text": "Keep this cached.", + "cache_control": {"type": "ephemeral"}, + } + ] + assert codec.encode(annotated, original).content == original.content + + instructions = cast(list[JsonObject], annotated.instructions) + instructions[0]["text"] = "Edited safely." + annotated.instructions = instructions + encoded = codec.encode(annotated, original) + expected = dict(cast(JsonObject, original.content)) + expected["system"] = [ + { + "type": "text", + "text": "Edited safely.", + "cache_control": {"type": "ephemeral"}, + } + ] + assert encoded.content == expected + # --------------------------------------------------------------------------- # 3. Built-in codec decode_response diff --git a/python/tests/test_codecs.py b/python/tests/test_codecs.py index d5ef6810b..5fd18bfd5 100644 --- a/python/tests/test_codecs.py +++ b/python/tests/test_codecs.py @@ -168,20 +168,31 @@ def test_annotated_llm_request_all_fields(self): annotated = AnnotatedLLMRequest( messages, + instructions=[{"type": "text", "text": "Follow policy"}], model="gpt-4-turbo", params=params, tools=tools_list, tool_choice="auto", + api_specific={"api": "custom", "api_name": "test", "data": {"flag": True}}, extra=extra, ) assert annotated.messages == messages + assert annotated.instructions == [{"type": "text", "text": "Follow policy"}] assert annotated.model == "gpt-4-turbo" assert annotated.params == params assert annotated.tools == tools_list assert annotated.tool_choice == "auto" + assert annotated.api_specific == {"api": "custom", "api_name": "test", "data": {"flag": True}} assert annotated.extra == extra + annotated.instructions = "Updated policy" + annotated.api_specific = None + assert annotated.instructions == "Updated policy" + assert annotated.api_specific is None + annotated.instructions = None + assert annotated.instructions is None + # --------------------------------------------------------------------------- # 2. LlmCodec subclass decode/encode round-trip diff --git a/python/tests/test_tools.py b/python/tests/test_tools.py index 0fa667afa..3619902f3 100644 --- a/python/tests/test_tools.py +++ b/python/tests/test_tools.py @@ -3,6 +3,8 @@ """Tests for NeMo Relay tool lifecycle, guardrails, and intercepts.""" +from collections import UserDict, UserList +from dataclasses import dataclass from typing import cast import pytest @@ -91,6 +93,61 @@ def my_func(args): result = await tools.execute("double", {"x": 5}, my_func) assert result == {"result": 10} + async def test_execute_rejects_cyclic_results_and_remains_usable(self): + def cyclic_result(_args): + result = {} + result["self"] = result + return result + + with pytest.raises(RuntimeError, match="circular reference detected"): + await tools.execute("cyclic_result", {}, cyclic_result) + + async def async_cyclic_result(_args): + return cyclic_result(_args) + + with pytest.raises(RuntimeError, match="circular reference detected"): + await tools.execute("async_cyclic_result", {}, async_cyclic_result) + + cyclic_mapping = UserDict() + cyclic_mapping["self"] = cyclic_mapping + with pytest.raises(RuntimeError, match="circular reference detected"): + await tools.execute("cyclic_mapping", {}, lambda _args: cyclic_mapping) + + cyclic_sequence = UserList() + cyclic_sequence.append(cyclic_sequence) + with pytest.raises(RuntimeError, match="circular reference detected"): + await tools.execute("cyclic_sequence", {}, lambda _args: cyclic_sequence) + + @dataclass + class CyclicDataclass: + child: object | None = None + + cyclic_dataclass = CyclicDataclass() + cyclic_dataclass.child = cyclic_dataclass + with pytest.raises(RuntimeError, match="circular reference detected"): + await tools.execute("cyclic_dataclass", {}, lambda _args: cyclic_dataclass) + + result = await tools.execute( + "post_cycle_result", + {}, + lambda _args: {"status": "ok"}, + ) + assert result == {"status": "ok"} + + async def test_execute_allows_shared_non_cyclic_results(self): + shared = {"value": True} + + result = await tools.execute( + "shared_result", + {}, + lambda _args: {"first": shared, "second": shared}, + ) + + assert result == { + "first": {"value": True}, + "second": {"value": True}, + } + async def test_execute_returns_string(self): def func(args): return "hello" diff --git a/skills/nemo-relay-instrument-typed-wrappers/SKILL.md b/skills/nemo-relay-instrument-typed-wrappers/SKILL.md index 2805cb88a..5b7218c9e 100644 --- a/skills/nemo-relay-instrument-typed-wrappers/SKILL.md +++ b/skills/nemo-relay-instrument-typed-wrappers/SKILL.md @@ -42,6 +42,17 @@ Keep typed boundaries explicit so middleware still sees predictable JSON. - Request codecs run before LLM request intercepts. Intercepts receive both the raw `LLMRequest` and optional annotated request; `encode` merges annotated edits back before execution intercepts and the provider callback run. +- Built-in request codecs guarantee JSON-value identity for an unchanged + annotation. They compare edits with a decoded baseline and patch only changed + fields, preserving native representation details and unknown fields. +- Use `instructions`, portable messages and components, and the tagged + `api_specific` request surface for normalized edits. Provider-only union + members use explicit `{ provider, kind, value }` native components. Reserve + top-level `extra` for unknown future fields. +- The `nemo-relay` gateway always supplies matching request codecs on Anthropic + Messages, OpenAI Chat Completions, and OpenAI Responses generation routes. + On those routes, treat raw `request.content` as read-only and return body + edits through the annotated request. Header edits still use the raw request. ## Key Rules @@ -69,6 +80,8 @@ Keep typed boundaries explicit so middleware still sees predictable JSON. - [ ] Response codec failures do not break the underlying LLM call - [ ] Request codec `encode` preserves original provider fields unless an intercept intentionally changes them +- [ ] Provider-native components belong to the codec's provider surface +- [ ] Gateway generation intercepts do not mutate raw `request.content` ## Related Skills