diff --git a/crates/mesh-llm-guardrails/src/collaboration.rs b/crates/mesh-llm-guardrails/src/collaboration.rs new file mode 100644 index 0000000000..768e1bafeb --- /dev/null +++ b/crates/mesh-llm-guardrails/src/collaboration.rs @@ -0,0 +1,332 @@ +use serde::Deserialize; +use serde_json::{Map, Value}; + +pub const MESH_COLLABORATION_FIELD: &str = "mesh_collaboration"; + +#[derive(Debug, Clone, PartialEq, Eq, Deserialize)] +#[serde(tag = "mode", rename_all = "snake_case", deny_unknown_fields)] +pub enum CollaborationContract { + ReportRequired(ReportRequired), +} + +#[derive(Debug, Clone, PartialEq, Eq, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ReportRequired { + pub version: u8, + pub tool: String, + pub body_argument: String, + pub locked_arguments: Map, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum CollaborationContractError { + Invalid(String), + UnknownTool(String), + BodyArgumentLocked, + ArgumentsDoNotMatchToolSchema(String), +} + +impl std::fmt::Display for CollaborationContractError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Invalid(reason) => write!( + f, + "mesh_collaboration must be a valid report_required contract: {reason}" + ), + Self::UnknownTool(tool) => write!( + f, + "mesh_collaboration tool `{tool}` is not declared in tools" + ), + Self::BodyArgumentLocked => { + f.write_str("mesh_collaboration body_argument must not also be locked") + } + Self::ArgumentsDoNotMatchToolSchema(reason) => write!( + f, + "mesh_collaboration arguments do not match the declared tool schema: {reason}" + ), + } + } +} + +impl std::error::Error for CollaborationContractError {} + +impl CollaborationContract { + pub fn parse(value: &Value, tools: Option<&Value>) -> Result { + let contract: Self = serde_json::from_value(value.clone()) + .map_err(|error| CollaborationContractError::Invalid(error.to_string()))?; + let report = contract.report_required(); + if report.version != 1 { + return Err(CollaborationContractError::Invalid( + "version must be 1".into(), + )); + } + if report.tool.trim().is_empty() || report.body_argument.trim().is_empty() { + return Err(CollaborationContractError::Invalid( + "tool and body_argument must be non-empty".into(), + )); + } + if report.locked_arguments.contains_key(&report.body_argument) { + return Err(CollaborationContractError::BodyArgumentLocked); + } + if !tool_is_declared(tools, &report.tool) { + return Err(CollaborationContractError::UnknownTool(report.tool.clone())); + } + let tool = declared_tools(tools, &report.tool) + .next() + .expect("unique tool checked"); + if tool + .pointer(&format!( + "/function/parameters/properties/{}/type", + report.body_argument + )) + .and_then(Value::as_str) + != Some("string") + { + return Err(CollaborationContractError::ArgumentsDoNotMatchToolSchema( + "body_argument must name a top-level string property".into(), + )); + } + let mut candidate = report.locked_arguments.clone(); + candidate.insert(report.body_argument.clone(), Value::String("report".into())); + let sanitized = crate::tools::sanitize_tool_arguments_for_tool( + &report.tool, + &Value::Object(candidate.clone()), + tools, + ) + .map_err(|error| { + CollaborationContractError::ArgumentsDoNotMatchToolSchema(error.to_string()) + })?; + if sanitized != Value::Object(candidate) { + return Err(CollaborationContractError::ArgumentsDoNotMatchToolSchema( + "an argument was rejected by the schema".into(), + )); + } + Ok(contract) + } + + pub fn report_required(&self) -> &ReportRequired { + match self { + Self::ReportRequired(report) => report, + } + } + + pub fn final_arguments( + &self, + content: Option<&str>, + tool_calls: Option<&Value>, + ) -> Result, CollaborationOutputError> { + let report = self.report_required(); + let native_arguments = match tool_calls { + Some(calls) => arguments_from_single_call(calls, &report.tool)?, + None => Map::new(), + }; + let body = native_arguments + .get(&report.body_argument) + .and_then(Value::as_str) + .map(str::to_owned) + .or_else(|| { + content + .map(str::trim) + .filter(|text| !text.is_empty()) + .map(str::to_owned) + }) + .ok_or(CollaborationOutputError::MissingReportBody)?; + let mut arguments = report.locked_arguments.clone(); + arguments.insert(report.body_argument.clone(), Value::String(body)); + Ok(arguments) + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum CollaborationOutputError { + InvalidChoiceCount, + InvalidToolCallCount, + WrongTool { expected: String, actual: String }, + InvalidArguments, + MissingReportBody, +} + +impl std::fmt::Display for CollaborationOutputError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::InvalidChoiceCount => { + f.write_str("report_required output must contain exactly one choice") + } + Self::InvalidToolCallCount => { + f.write_str("report_required output must contain exactly one tool call") + } + Self::WrongTool { expected, actual } => write!( + f, + "report_required output called `{actual}` instead of `{expected}`" + ), + Self::InvalidArguments => { + f.write_str("report_required tool arguments must be a JSON object") + } + Self::MissingReportBody => { + f.write_str("report_required output did not contain a report body") + } + } + } +} + +impl std::error::Error for CollaborationOutputError {} + +fn tool_is_declared(tools: Option<&Value>, expected: &str) -> bool { + declared_tools(tools, expected).count() == 1 +} + +fn declared_tools<'a>( + tools: Option<&'a Value>, + expected: &'a str, +) -> impl Iterator { + tools + .and_then(Value::as_array) + .into_iter() + .flatten() + .filter(move |tool| { + tool.pointer("/function/name").and_then(Value::as_str) == Some(expected) + }) +} + +pub fn collaboration_calls_output_tool(calls: &Value, contract: &CollaborationContract) -> bool { + calls.as_array().is_some_and(|calls| { + calls.iter().any(|call| { + call.pointer("/function/name").and_then(Value::as_str) + == Some(contract.report_required().tool.as_str()) + }) + }) +} + +fn arguments_from_single_call( + calls: &Value, + expected_tool: &str, +) -> Result, CollaborationOutputError> { + let calls = calls + .as_array() + .filter(|calls| calls.len() == 1) + .ok_or(CollaborationOutputError::InvalidToolCallCount)?; + let call = &calls[0]; + let actual = call + .pointer("/function/name") + .and_then(Value::as_str) + .ok_or(CollaborationOutputError::InvalidArguments)?; + if actual != expected_tool { + return Err(CollaborationOutputError::WrongTool { + expected: expected_tool.to_owned(), + actual: actual.to_owned(), + }); + } + let arguments = call + .pointer("/function/arguments") + .ok_or(CollaborationOutputError::InvalidArguments)?; + match arguments { + Value::Object(arguments) => Ok(arguments.clone()), + Value::String(arguments) => serde_json::from_str::>(arguments) + .map_err(|_| CollaborationOutputError::InvalidArguments), + _ => Err(CollaborationOutputError::InvalidArguments), + } +} + +pub fn finalize_openai_response_value( + response: &mut Value, + contract: &CollaborationContract, +) -> Result { + let choices = response + .get_mut("choices") + .and_then(Value::as_array_mut) + .filter(|choices| choices.len() == 1) + .ok_or(CollaborationOutputError::InvalidChoiceCount)?; + let choice = choices + .first_mut() + .and_then(Value::as_object_mut) + .ok_or(CollaborationOutputError::InvalidChoiceCount)?; + let message = choice + .get_mut("message") + .and_then(Value::as_object_mut) + .ok_or(CollaborationOutputError::MissingReportBody)?; + let tool_calls = message.get("tool_calls").cloned(); + if let Some(calls) = tool_calls.as_ref() + && !collaboration_calls_output_tool(calls, contract) + { + return Ok(false); + } + let content = message.get("content").and_then(Value::as_str); + let arguments = contract.final_arguments(content, tool_calls.as_ref())?; + message.insert("content".into(), Value::Null); + message.insert( + "tool_calls".into(), + serde_json::json!([{ + "id":"call_mesh_collaboration", + "type":"function", + "function":{ + "name":contract.report_required().tool, + "arguments":Value::Object(arguments).to_string() + } + }]), + ); + choice.insert("finish_reason".into(), Value::String("tool_calls".into())); + Ok(true) +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + fn tools() -> Value { + json!([{"type":"function","function":{"name":"buzz_send","parameters":{"type":"object","properties":{"content":{"type":"string"},"channel":{"type":"string"},"reply_to":{"type":"string"}},"required":["content"]}}}]) + } + + #[test] + fn prose_becomes_report_arguments_with_locked_values() { + let contract = CollaborationContract::parse( + &json!({ + "mode":"report_required", "version":1, "tool":"buzz_send", "body_argument":"content", + "locked_arguments":{"channel":"locked", "reply_to":"root"} + }), + Some(&tools()), + ) + .expect("valid contract"); + let arguments = contract + .final_arguments(Some("Done."), None) + .expect("report"); + assert_eq!( + arguments, + json!({"content":"Done.","channel":"locked","reply_to":"root"}) + .as_object() + .unwrap() + .clone() + ); + } + + #[test] + fn native_call_body_survives_but_locked_values_win() { + let contract = CollaborationContract::parse( + &json!({ + "mode":"report_required", "version":1, "tool":"buzz_send", "body_argument":"content", + "locked_arguments":{"channel":"locked"} + }), + Some(&tools()), + ) + .expect("valid contract"); + let calls = json!([{"function":{"name":"buzz_send","arguments":"{\"content\":\"Done.\",\"channel\":\"wrong\"}"}}]); + let arguments = contract + .final_arguments(None, Some(&calls)) + .expect("report"); + assert_eq!(arguments["content"], "Done."); + assert_eq!(arguments["channel"], "locked"); + } + + #[test] + fn contract_rejects_undeclared_tool() { + let error = CollaborationContract::parse( + &json!({ + "mode":"report_required", "version":1, "tool":"missing", "body_argument":"content", + "locked_arguments":{} + }), + Some(&tools()), + ) + .expect_err("unknown tool"); + assert!(matches!(error, CollaborationContractError::UnknownTool(_))); + } +} diff --git a/crates/mesh-llm-guardrails/src/lib.rs b/crates/mesh-llm-guardrails/src/lib.rs index d4e05f6356..b47815babd 100644 --- a/crates/mesh-llm-guardrails/src/lib.rs +++ b/crates/mesh-llm-guardrails/src/lib.rs @@ -1,3 +1,4 @@ +pub mod collaboration; pub mod compact; pub mod content; pub mod policy; @@ -5,6 +6,11 @@ pub mod request_contract; pub mod structured; pub mod tools; +pub use collaboration::{ + CollaborationContract, CollaborationContractError, CollaborationOutputError, + MESH_COLLABORATION_FIELD, ReportRequired, collaboration_calls_output_tool, + finalize_openai_response_value, +}; pub use compact::{ CompactionConfig, CompactionDecision, CompactionOverride, CompactionReport, CompactionRequest, MESH_COMPACT_FIELD, compact_messages, estimate_message_tokens, diff --git a/crates/mesh-mixture-of-agents/src/gateway.rs b/crates/mesh-mixture-of-agents/src/gateway.rs index a9e9fe2212..a0a139c2a3 100644 --- a/crates/mesh-mixture-of-agents/src/gateway.rs +++ b/crates/mesh-mixture-of-agents/src/gateway.rs @@ -52,7 +52,7 @@ pub async fn handle_turn(config: &GatewayConfig, body: &Value) -> TurnResult { let allowed_tools = session.tool_names(); - match turn_type { + let mut result = match turn_type { session::TurnType::ToolResult => { handle_tool_result(config, &session, has_tools, &allowed_tools, start).await } @@ -67,7 +67,28 @@ pub async fn handle_turn(config: &GatewayConfig, body: &Value) -> TurnResult { ) .await } + }; + if let Some(raw_contract) = body.get(mesh_llm_guardrails::MESH_COLLABORATION_FIELD) { + match mesh_llm_guardrails::CollaborationContract::parse(raw_contract, tools.as_ref()) + .and_then(|contract| { + mesh_llm_guardrails::finalize_openai_response_value( + &mut result.response_body, + &contract, + ) + .map(|_| ()) + .map_err(|error| { + mesh_llm_guardrails::CollaborationContractError::Invalid(error.to_string()) + }) + }) { + Ok(()) => {} + Err(error) => { + result.response_body = + error_response(&error.to_string(), "mesh_collaboration_invalid"); + result.turn_kind = TurnKind::Failed; + } + } } + result } // ─── Query handling ────────────────────────────────────────────────── diff --git a/crates/openai-frontend/src/guardrails/mod.rs b/crates/openai-frontend/src/guardrails/mod.rs index 09a4924fc4..636c176569 100644 --- a/crates/openai-frontend/src/guardrails/mod.rs +++ b/crates/openai-frontend/src/guardrails/mod.rs @@ -1,6 +1,7 @@ use std::sync::Arc; use async_trait::async_trait; +use serde_json::{Value, json}; use crate::{ backend::{ @@ -72,6 +73,31 @@ impl GuardedOpenAiBackend { self } + async fn collaborative_chat_completion( + &self, + mut request: ChatCompletionRequest, + context: OpenAiRequestContext, + ) -> OpenAiResult { + let Some(raw_contract) = request + .extra + .remove(mesh_llm_guardrails::MESH_COLLABORATION_FIELD) + else { + return self.guarded_chat_completion(request, context).await; + }; + if request.stream { + return Err(crate::errors::OpenAiError::unsupported( + "mesh_collaboration does not support streaming", + )); + } + let contract = mesh_llm_guardrails::CollaborationContract::parse( + &raw_contract, + request.tools.as_ref(), + ) + .map_err(|error| crate::errors::OpenAiError::invalid_request(error.to_string()))?; + let response = self.guarded_chat_completion(request, context).await?; + finalize_collaboration_response(response, &contract) + } + async fn guarded_chat_completion( &self, request: ChatCompletionRequest, @@ -300,6 +326,41 @@ fn telemetry_attempt_bucket(attempts: u8) -> GuardrailTelemetryAttemptBucket { } } +fn finalize_collaboration_response( + mut response: ChatCompletionResponse, + contract: &mesh_llm_guardrails::CollaborationContract, +) -> OpenAiResult { + if response.choices.len() != 1 { + return Err(crate::errors::OpenAiError::backend( + "report_required output must contain exactly one choice", + )); + } + let choice = response.choices.first_mut().expect("choice count checked"); + if let Some(tool_calls) = choice.message.tool_calls.as_ref() + && !mesh_llm_guardrails::collaboration_calls_output_tool(tool_calls, contract) + { + return Ok(response); + } + let arguments = contract + .final_arguments( + choice.message.content.as_deref(), + choice.message.tool_calls.as_ref(), + ) + .map_err(|error| crate::errors::OpenAiError::backend(error.to_string()))?; + let tool = &contract.report_required().tool; + choice.message.content = None; + choice.message.tool_calls = Some(json!([{ + "id": "call_mesh_collaboration", + "type": "function", + "function": { + "name": tool, + "arguments": Value::Object(arguments).to_string(), + } + }])); + choice.finish_reason = Some(crate::common::FinishReason::ToolCalls); + Ok(response) +} + #[async_trait] impl OpenAiBackend for GuardedOpenAiBackend { async fn models(&self) -> OpenAiResult> { @@ -310,7 +371,7 @@ impl OpenAiBackend for GuardedOpenAiBackend { &self, request: ChatCompletionRequest, ) -> OpenAiResult { - self.guarded_chat_completion(request, OpenAiRequestContext::new()) + self.collaborative_chat_completion(request, OpenAiRequestContext::new()) .await } @@ -319,7 +380,7 @@ impl OpenAiBackend for GuardedOpenAiBackend { request: ChatCompletionRequest, context: OpenAiRequestContext, ) -> OpenAiResult { - self.guarded_chat_completion(request, context).await + self.collaborative_chat_completion(request, context).await } async fn chat_completion_stream( @@ -327,6 +388,14 @@ impl OpenAiBackend for GuardedOpenAiBackend { request: ChatCompletionRequest, context: OpenAiRequestContext, ) -> OpenAiResult { + if request + .extra + .contains_key(mesh_llm_guardrails::MESH_COLLABORATION_FIELD) + { + return Err(crate::errors::OpenAiError::unsupported( + "mesh_collaboration does not support streaming", + )); + } self.backend.chat_completion_stream(request, context).await } diff --git a/crates/openai-frontend/src/guardrails/tests.rs b/crates/openai-frontend/src/guardrails/tests.rs index e4982e0eb1..d8269fcb13 100644 --- a/crates/openai-frontend/src/guardrails/tests.rs +++ b/crates/openai-frontend/src/guardrails/tests.rs @@ -1422,3 +1422,131 @@ fn supported_json_schema_response_format() -> serde_json::Value { } }) } + +#[tokio::test] +async fn report_required_wraps_prose_and_locks_routing() { + let backend = Arc::new(SequencedBackend::new(vec![Ok(response_with_content( + "small-model", + "Tests pass.", + ))])); + let guarded = GuardedOpenAiBackend::new(backend, GuardrailPolicy::default()); + let request: ChatCompletionRequest = serde_json::from_value(json!({ + "model": "small-model", + "messages": [{"role":"user","content":"Report the result"}], + "tools": [{"type":"function","function":{ + "name":"buzz_send", + "parameters":{ + "type":"object", + "properties":{ + "content":{"type":"string"}, + "channel":{"type":"string"}, + "reply_to":{"type":"string"} + }, + "required":["content","channel","reply_to"], + "additionalProperties":false + } + }}], + "mesh_collaboration": { + "mode":"report_required", + "version":1, + "tool":"buzz_send", + "body_argument":"content", + "locked_arguments":{"channel":"channel-1","reply_to":"event-1"} + } + })) + .expect("request"); + + let response = guarded.chat_completion(request).await.expect("response"); + let choice = &response.choices[0]; + assert_eq!(choice.message.content, None); + assert_eq!( + choice.finish_reason, + Some(crate::common::FinishReason::ToolCalls) + ); + let call = &choice.message.tool_calls.as_ref().expect("call")[0]; + assert_eq!(call["function"]["name"], "buzz_send"); + let arguments: serde_json::Value = serde_json::from_str( + call["function"]["arguments"] + .as_str() + .expect("arguments string"), + ) + .expect("arguments json"); + assert_eq!( + arguments, + json!({ + "content":"Tests pass.", "channel":"channel-1", "reply_to":"event-1" + }) + ); +} + +#[tokio::test] +async fn report_required_rejects_multiple_choices() { + let mut backend_response = response_with_content("small-model", "first"); + let mut second = backend_response.choices[0].clone(); + second.index = 1; + second.message.content = Some("second".to_string()); + backend_response.choices.push(second); + let backend = Arc::new(SequencedBackend::new(vec![Ok(backend_response)])); + let guarded = GuardedOpenAiBackend::new(backend, GuardrailPolicy::default()); + let request: ChatCompletionRequest = serde_json::from_value(json!({ + "model":"small-model", + "messages":[{"role":"user","content":"Report the result"}], + "tools":[{"type":"function","function":{"name":"buzz_send","parameters":{"type":"object","properties":{"content":{"type":"string"}},"required":["content"],"additionalProperties":false}}}], + "mesh_collaboration":{"mode":"report_required","version":1,"tool":"buzz_send","body_argument":"content","locked_arguments":{}} + })) + .expect("request"); + + let error = guarded + .chat_completion(request) + .await + .expect_err("multiple choices must fail closed"); + assert!( + error + .to_string() + .contains("report_required output must contain exactly one choice") + ); +} + +#[tokio::test] +async fn absent_collaboration_contract_preserves_ordinary_response() { + let backend = Arc::new(SequencedBackend::new(vec![Ok(response_with_content( + "small-model", + "ordinary", + ))])); + let guarded = GuardedOpenAiBackend::new(backend, enforce_policy()); + let request: ChatCompletionRequest = serde_json::from_value(json!({ + "model":"small-model", + "messages":[{"role":"user","content":"hello"}] + })) + .expect("request"); + + let response = guarded.chat_completion(request).await.expect("response"); + assert_eq!( + response.choices[0].message.content.as_deref(), + Some("ordinary") + ); + assert!(response.choices[0].message.tool_calls.is_none()); +} + +#[tokio::test] +async fn report_required_passes_intermediate_tool_call_through() { + let intermediate = response_with_tool_calls( + "small-model", + json!([{"type":"function","function":{"name":"read_file","arguments":"{\"path\":\"README.md\"}"}}]), + None, + ); + let backend = Arc::new(SequencedBackend::new(vec![Ok(intermediate.clone())])); + let guarded = GuardedOpenAiBackend::new(backend, GuardrailPolicy::default()); + let request: ChatCompletionRequest = serde_json::from_value(json!({ + "model":"small-model", + "messages":[{"role":"user","content":"Inspect then report"}], + "tools":[ + {"type":"function","function":{"name":"read_file","parameters":{"type":"object","properties":{"path":{"type":"string"}},"required":["path"]}}}, + {"type":"function","function":{"name":"buzz_send","parameters":{"type":"object","properties":{"content":{"type":"string"},"channel":{"type":"string"}},"required":["content","channel"],"additionalProperties":false}}} + ], + "mesh_collaboration":{"mode":"report_required","version":1,"tool":"buzz_send","body_argument":"content","locked_arguments":{"channel":"locked"}} + })).expect("request"); + + let response = guarded.chat_completion(request).await.expect("response"); + assert_eq!(response, intermediate); +} diff --git a/docs/specs/mesh-collaboration-report.md b/docs/specs/mesh-collaboration-report.md new file mode 100644 index 0000000000..dff856fd26 --- /dev/null +++ b/docs/specs/mesh-collaboration-report.md @@ -0,0 +1,36 @@ +# Mesh collaboration report contract + +`mesh_collaboration` is an opt-in extension to non-streaming +`/v1/chat/completions`. Its initial mode guarantees that a final prose answer is +returned as one caller-declared report tool call while leaving intermediate tool +calls unchanged. + +```json +{ + "mesh_collaboration": { + "mode": "report_required", + "version": 1, + "tool": "buzz_send", + "body_argument": "content", + "locked_arguments": { + "channel": "channel-id", + "reply_to": "event-id" + } + } +} +``` + +The named tool must be declared exactly once in `tools`. `body_argument` must be +a top-level string property and cannot also be locked. Locked values must satisfy +the tool schema. Mesh constructs final arguments from only the generated report +body and `locked_arguments`; model-supplied extra arguments are discarded. + +If the model returns another declared tool call, Mesh passes it through so the +agent can continue investigating. A prose response is wrapped as the report +call. A native call to the report tool is normalized to the same safe output. +Requests without `mesh_collaboration` are unchanged. Streaming is rejected in +this first version. + +This contract deliberately requires a distinct structured report tool. It does +not interpret shell commands or try to distinguish investigative and publishing +uses of a generic shell tool.