From d70a13f07d5c5bfd1a82dd8520c6085ef549f874 Mon Sep 17 00:00:00 2001 From: Will Killian Date: Fri, 7 Aug 2026 16:18:50 -0400 Subject: [PATCH 1/4] chore: address Sonar code quality findings Signed-off-by: Will Killian --- crates/adaptive/src/response_cache/key.rs | 80 +- crates/cli/src/diagnostics/mod.rs | 250 ++-- .../core/src/codec/gemini_generate_content.rs | 1183 +++++++++-------- .../src/observability/plugin_component.rs | 111 +- crates/core/src/plugin.rs | 257 ++-- .../src/plugins/nemo_guardrails/python.rs | 117 +- crates/node/src/callable.rs | 106 +- crates/pii-redaction/src/overlay.rs | 113 +- go/nemo_relay/logging_test.go | 114 +- go/nemo_relay/plugin_activation_test.go | 45 +- .../latency_benchmark/src/report/report.js | 30 +- .../src/report/template.html | 3 - 12 files changed, 1302 insertions(+), 1107 deletions(-) diff --git a/crates/adaptive/src/response_cache/key.rs b/crates/adaptive/src/response_cache/key.rs index 29f1b6555..dbdb1726b 100644 --- a/crates/adaptive/src/response_cache/key.rs +++ b/crates/adaptive/src/response_cache/key.rs @@ -453,45 +453,55 @@ fn lossy_gemini_part(part: &Json, plain_text_parts: &mut usize) -> bool { return true; }; match data_key { - "text" => { - if !object.get("text").is_some_and(Json::is_string) { - return true; - } - if object.len() == 1 { - *plain_text_parts += 1; + "text" => lossy_gemini_text_part(object, plain_text_parts), + "functionCall" => lossy_gemini_function_call_part(object), + "functionResponse" => lossy_gemini_function_response_part(object), + _ => false, + } +} + +fn lossy_gemini_text_part( + object: &serde_json::Map, + plain_text_parts: &mut usize, +) -> bool { + if !object.get("text").is_some_and(Json::is_string) { + return true; + } + if object.len() == 1 { + *plain_text_parts += 1; + } + false +} + +fn lossy_gemini_function_call_part(object: &serde_json::Map) -> bool { + object.keys().any(|key| key != "functionCall") + || match object.get("functionCall").and_then(Json::as_object) { + Some(call) => { + call.keys() + .any(|key| !matches!(key.as_str(), "name" | "id" | "args")) + || call.get("args").is_some_and(|args| !args.is_object()) } - false + None => true, } - "functionCall" => { - object.keys().any(|key| key != "functionCall") - || match object.get("functionCall").and_then(Json::as_object) { - Some(fc) => { - fc.keys() - .any(|key| !matches!(key.as_str(), "name" | "id" | "args")) - || fc.get("args").is_some_and(|args| !args.is_object()) - } - None => true, - } - } - "functionResponse" => { - object.keys().any(|key| key != "functionResponse") - || match object.get("functionResponse").and_then(Json::as_object) { - Some(fr) => { - fr.keys().any(|key| { - !matches!(key.as_str(), "id" | "name" | "response" | "parts") - }) || match ( - fr.get("id").and_then(Json::as_str), - fr.get("name").and_then(Json::as_str), - ) { - (Some(id), Some(name)) => id != name, - _ => false, - } +} + +fn lossy_gemini_function_response_part(object: &serde_json::Map) -> bool { + object.keys().any(|key| key != "functionResponse") + || match object.get("functionResponse").and_then(Json::as_object) { + Some(response) => { + response + .keys() + .any(|key| !matches!(key.as_str(), "id" | "name" | "response" | "parts")) + || match ( + response.get("id").and_then(Json::as_str), + response.get("name").and_then(Json::as_str), + ) { + (Some(id), Some(name)) => id != name, + _ => false, } - None => true, - } + } + None => true, } - _ => false, - } } fn is_gemini_part_data_key(key: &str) -> bool { diff --git a/crates/cli/src/diagnostics/mod.rs b/crates/cli/src/diagnostics/mod.rs index c033d7d54..e8de8e53b 100644 --- a/crates/cli/src/diagnostics/mod.rs +++ b/crates/cli/src/diagnostics/mod.rs @@ -888,84 +888,82 @@ async fn observability_http_exporter_checks( return Vec::new(); }; join_all( - endpoints - .iter() - .enumerate() - .map(|(index, endpoint)| async move { - let endpoint_type = endpoint - .get("type") - .and_then(Value::as_str) - .unwrap_or("unknown"); - let label = "OpenTelemetry endpoint"; - let transport = endpoint - .get("transport") - .and_then(Value::as_str) - .unwrap_or("http_binary"); - match endpoint.get("endpoint").and_then(Value::as_str) { - Some(url) => { - let mut check = if transport == "grpc" { - if probe_mode.is_offline() { - match validate_grpc_endpoint(url) { - Ok(_) => Check { - name: label, - status: Status::Info, - details: format!( - "endpoints[{index}] ({endpoint_type}): live network probe skipped (--offline)" - ), - }, - Err(details) => Check { - name: label, - status: Status::Fail, - details: format!( - "endpoints[{index}] ({endpoint_type}): {details}" - ), - }, - } - } else { - probe_tcp_named(label, url).await - } - } else { - let effective_url = resolve_http_trace_endpoint(url); - if probe_mode.is_offline() { - match validate_otlp_http_endpoint(effective_url.as_ref()) { - Ok(()) => Check { - name: label, - status: Status::Info, - details: format!( - "endpoints[{index}] ({endpoint_type}): live network probe skipped (--offline)" - ), - }, - Err(details) => Check { - name: label, - status: Status::Fail, - details: format!( - "endpoints[{index}] ({endpoint_type}): {details}" - ), - }, - } - } else { - probe_otlp_http_named(label, effective_url.as_ref()).await - } - }; - if !probe_mode.is_offline() { - check.details = - format!("endpoints[{index}] ({endpoint_type}): {}", check.details); - } - check - } - None => Check { - name: label, - status: Status::Fail, - details: format!( - "endpoints[{index}] ({endpoint_type}): endpoint is required" - ), - }, - } - }), + endpoints.iter().enumerate().map(|(index, endpoint)| { + observability_http_exporter_check(index, endpoint, probe_mode) + }), ) .await } +async fn observability_http_exporter_check( + index: usize, + endpoint: &Value, + probe_mode: DoctorProbeMode, +) -> Check { + let endpoint_type = endpoint + .get("type") + .and_then(Value::as_str) + .unwrap_or("unknown"); + let label = "OpenTelemetry endpoint"; + let Some(url) = endpoint.get("endpoint").and_then(Value::as_str) else { + return Check { + name: label, + status: Status::Fail, + details: format!("endpoints[{index}] ({endpoint_type}): endpoint is required"), + }; + }; + let transport = endpoint + .get("transport") + .and_then(Value::as_str) + .unwrap_or("http_binary"); + let mut check = if transport == "grpc" { + probe_grpc_endpoint(label, url, probe_mode).await + } else { + probe_http_endpoint(label, url, probe_mode).await + }; + if !probe_mode.is_offline() { + check.details = format!("endpoints[{index}] ({endpoint_type}): {}", check.details); + } + check +} + +async fn probe_grpc_endpoint(label: &'static str, url: &str, mode: DoctorProbeMode) -> Check { + if mode.is_offline() { + return match validate_grpc_endpoint(url) { + Ok(_) => Check { + name: label, + status: Status::Info, + details: "live network probe skipped (--offline)".into(), + }, + Err(details) => Check { + name: label, + status: Status::Fail, + details, + }, + }; + } + probe_tcp_named(label, url).await +} + +async fn probe_http_endpoint(label: &'static str, url: &str, mode: DoctorProbeMode) -> Check { + let effective_url = resolve_http_trace_endpoint(url); + if mode.is_offline() { + return match validate_otlp_http_endpoint(effective_url.as_ref()) { + Ok(()) => Check { + name: label, + status: Status::Info, + details: "live network probe skipped (--offline)".into(), + }, + Err(details) => Check { + name: label, + status: Status::Fail, + details, + }, + }; + } + probe_otlp_http_named(label, effective_url.as_ref()).await +} + fn observability_component_config(plugin_value: &Value) -> Option<&Value> { plugin_value .get("components") @@ -1229,54 +1227,74 @@ fn validate_atof_stream_probe_target( fn endpoint_headers(endpoint: &Value) -> Result, String> { let mut out = Vec::new(); let mut names = std::collections::HashSet::new(); - if let Some(headers) = endpoint.get("headers") { - let Some(object) = headers.as_object() else { - return Err("headers must be an object of string values".into()); + append_configured_headers(endpoint.get("headers"), &mut names, &mut out)?; + append_environment_headers(endpoint.get("header_env"), &mut names, &mut out)?; + Ok(out) +} + +fn append_configured_headers( + headers: Option<&Value>, + names: &mut std::collections::HashSet, + out: &mut Vec<(String, String)>, +) -> Result<(), String> { + let Some(headers) = headers else { + return Ok(()); + }; + let Some(object) = headers.as_object() else { + return Err("headers must be an object of string values".into()); + }; + for (key, value) in object { + let name = reqwest::header::HeaderName::from_bytes(key.as_bytes()) + .map_err(|error| error.to_string())?; + let Some(value) = value.as_str() else { + return Err(format!("headers.{key} must be a string")); }; - for (key, value) in object { - let name = reqwest::header::HeaderName::from_bytes(key.as_bytes()) - .map_err(|error| error.to_string())?; - let Some(value) = value.as_str() else { - return Err(format!("headers.{key} must be a string")); - }; - if value.trim().is_empty() { - return Err(format!("headers.{key} must not be blank")); - } - reqwest::header::HeaderValue::from_bytes(value.as_bytes()) - .map_err(|error| format!("headers.{key} invalid: {error}"))?; - if !names.insert(name) { - return Err(format!("header {key:?} appears more than once")); - } - out.push((key.clone(), value.to_string())); + if value.trim().is_empty() { + return Err(format!("headers.{key} must not be blank")); } + reqwest::header::HeaderValue::from_bytes(value.as_bytes()) + .map_err(|error| format!("headers.{key} invalid: {error}"))?; + if !names.insert(name) { + return Err(format!("header {key:?} appears more than once")); + } + out.push((key.clone(), value.to_string())); } - if let Some(header_env) = endpoint.get("header_env") { - let Some(object) = header_env.as_object() else { - return Err("header_env must be an object of string values".into()); + Ok(()) +} + +fn append_environment_headers( + header_env: Option<&Value>, + names: &mut std::collections::HashSet, + out: &mut Vec<(String, String)>, +) -> Result<(), String> { + let Some(header_env) = header_env else { + return Ok(()); + }; + let Some(object) = header_env.as_object() else { + return Err("header_env must be an object of string values".into()); + }; + for (key, variable) in object { + let name = reqwest::header::HeaderName::from_bytes(key.as_bytes()) + .map_err(|error| error.to_string())?; + if names.contains(&name) { + return Err(format!( + "header {key:?} cannot appear in both headers and header_env" + )); + } + let Some(variable) = variable.as_str() else { + return Err(format!("header_env.{key} must be a string")); }; - for (key, variable) in object { - let name = reqwest::header::HeaderName::from_bytes(key.as_bytes()) - .map_err(|error| error.to_string())?; - if names.contains(&name) { - return Err(format!( - "header {key:?} cannot appear in both headers and header_env" - )); - } - let Some(variable) = variable.as_str() else { - return Err(format!("header_env.{key} must be a string")); - }; - let value = std::env::var(variable) - .map_err(|_| format!("environment variable {variable:?} is not set"))?; - if value.trim().is_empty() { - return Err(format!("environment variable {variable:?} is blank")); - } - reqwest::header::HeaderValue::from_bytes(value.as_bytes()) - .map_err(|error| format!("header_env.{key} invalid: {error}"))?; - names.insert(name); - out.push((key.clone(), value)); + let value = std::env::var(variable) + .map_err(|_| format!("environment variable {variable:?} is not set"))?; + if value.trim().is_empty() { + return Err(format!("environment variable {variable:?} is blank")); } + reqwest::header::HeaderValue::from_bytes(value.as_bytes()) + .map_err(|error| format!("header_env.{key} invalid: {error}"))?; + names.insert(name); + out.push((key.clone(), value)); } - Ok(out) + Ok(()) } fn doctor_atof_probe_payload() -> Result { diff --git a/crates/core/src/codec/gemini_generate_content.rs b/crates/core/src/codec/gemini_generate_content.rs index c4137f23d..dd3b06a47 100644 --- a/crates/core/src/codec/gemini_generate_content.rs +++ b/crates/core/src/codec/gemini_generate_content.rs @@ -576,198 +576,215 @@ fn system_instruction_text(val: &Json) -> Option { /// - `functionCall` parts (model role) → `Message::Assistant { tool_calls: Some([…]) }` /// - text parts → plain message content fn gemini_content_to_messages(content: &Json) -> Result> { - // Each contents item must be a JSON object. let obj = content.as_object().ok_or_else(|| { FlowError::InvalidArgument("Gemini contents item must be an object".into()) })?; + let role = gemini_content_role(obj)?; + let parts = obj.get("parts").and_then(Json::as_array).ok_or_else(|| { + FlowError::InvalidArgument("Gemini contents item must have an array 'parts' field".into()) + })?; + let (fr_parts, fn_call_parts) = validate_gemini_content_parts(parts)?; + validate_gemini_content_roles(role, &fr_parts, &fn_call_parts)?; + if !fr_parts.is_empty() { + return gemini_function_response_messages(parts, &fr_parts); + } + let content = gemini_parts_to_message_content(parts, "request")?; + if !fn_call_parts.is_empty() { + return gemini_function_call_messages(content, &fn_call_parts); + } + Ok(vec![gemini_plain_message(role, content)]) +} - // Role is optional; absent role defaults to "user" per Google's REST spec. - // When present it must be a string; only "user" and "model" are accepted. - let role = match obj.get("role") { - None => "user", - Some(Json::String(s)) if s == "user" => "user", - Some(Json::String(s)) if s == "model" => "model", - Some(Json::String(other)) => { - return Err(FlowError::InvalidArgument(format!( - "Gemini contents item has unsupported role '{other}'; expected 'user' or 'model'" - ))); - } - Some(_) => { - return Err(FlowError::InvalidArgument( - "Gemini contents item 'role' must be a string".into(), - )); - } - }; - - // `parts` is required and must be an array. - let parts = obj - .get("parts") - .ok_or_else(|| { - FlowError::InvalidArgument("Gemini contents item is missing 'parts'".into()) - })? - .as_array() - .ok_or_else(|| { - FlowError::InvalidArgument("Gemini contents item 'parts' must be an array".into()) - })?; - - // Validate each part; collect functionResponse and functionCall parts with - // strict name checks so invalid items surface as errors rather than silent drops. - let mut fr_parts: Vec<&Json> = Vec::new(); - let mut fn_call_parts: Vec<&Json> = Vec::new(); +fn gemini_content_role(obj: &serde_json::Map) -> Result<&str> { + match obj.get("role") { + None => Ok("user"), + Some(Json::String(role)) if role == "user" || role == "model" => Ok(role), + Some(Json::String(other)) => Err(FlowError::InvalidArgument(format!( + "Gemini contents item has unsupported role '{other}'; expected 'user' or 'model'" + ))), + Some(_) => Err(FlowError::InvalidArgument( + "Gemini contents item 'role' must be a string".into(), + )), + } +} +fn validate_gemini_content_parts(parts: &[Json]) -> Result<(Vec<&Json>, Vec<&Json>)> { + let mut responses = Vec::new(); + let mut calls = Vec::new(); for part in parts { let part_obj = part.as_object().ok_or_else(|| { FlowError::InvalidArgument("Gemini parts item must be an object".into()) })?; - let data_key = validate_single_gemini_part_data_field(part_obj, "request")?; - if data_key == Some("functionResponse") { - let fr = part.get("functionResponse").unwrap(); - let fr_obj = fr.as_object().ok_or_else(|| { - FlowError::InvalidArgument("Gemini functionResponse must be an object".into()) - })?; - let name = fr_obj.get("name").and_then(Json::as_str).unwrap_or(""); - if name.is_empty() { - return Err(FlowError::InvalidArgument( - "Gemini functionResponse is missing a non-empty 'name'".into(), - )); + match validate_single_gemini_part_data_field(part_obj, "request")? { + Some("functionResponse") => { + validate_gemini_function_response_part(part)?; + responses.push(part); } - // `response` is required and must be an object (Gemini spec). - match fr_obj.get("response") { - None => { - return Err(FlowError::InvalidArgument( - "Gemini functionResponse is missing required 'response'".into(), - )); - } - Some(r) if !r.is_object() => { - return Err(FlowError::InvalidArgument( - "Gemini functionResponse.response must be an object".into(), - )); - } - _ => {} + Some("functionCall") => { + validate_gemini_function_call_part(part)?; + calls.push(part); } - if let Some(nested_parts) = fr_obj.get("parts") { - let nested_parts = nested_parts.as_array().ok_or_else(|| { - FlowError::InvalidArgument( - "Gemini functionResponse.parts must be an array".into(), - ) - })?; - for nested_part in nested_parts { - validate_gemini_nested_function_response_part(nested_part)?; - } - } - parse_optional_id(fr_obj, "functionResponse")?; - fr_parts.push(part); - } else if data_key == Some("functionCall") { - let fc = part.get("functionCall").unwrap(); - let fc_obj = fc.as_object().ok_or_else(|| { - FlowError::InvalidArgument("Gemini functionCall must be an object".into()) - })?; - let name = fc_obj.get("name").and_then(Json::as_str).unwrap_or(""); - if name.is_empty() { + Some("text") if part.get("text").is_some_and(|value| !value.is_string()) => { return Err(FlowError::InvalidArgument( - "Gemini functionCall is missing a non-empty 'name'".into(), + "Gemini parts item 'text' must be a string".into(), )); } - if fc_obj.get("args").is_some_and(|a| !a.is_object()) { - return Err(FlowError::InvalidArgument( - "Gemini functionCall.args must be an object".into(), - )); - } - fn_call_parts.push(part); - } else if data_key == Some("text") && part.get("text").is_some_and(|v| !v.is_string()) { - // A plain text part with a non-string text value has no lossless encoding. + _ => {} + } + } + if !responses.is_empty() && !calls.is_empty() { + return Err(FlowError::InvalidArgument( + "Gemini contents item must not contain both functionResponse and functionCall parts" + .into(), + )); + } + Ok((responses, calls)) +} + +fn validate_gemini_function_response_part(part: &Json) -> Result<()> { + let response = part + .get("functionResponse") + .and_then(Json::as_object) + .ok_or_else(|| { + FlowError::InvalidArgument("Gemini functionResponse must be an object".into()) + })?; + if response + .get("name") + .and_then(Json::as_str) + .is_none_or(str::is_empty) + { + return Err(FlowError::InvalidArgument( + "Gemini functionResponse is missing a non-empty 'name'".into(), + )); + } + match response.get("response") { + None => { + return Err(FlowError::InvalidArgument( + "Gemini functionResponse is missing required 'response'".into(), + )); + } + Some(value) if !value.is_object() => { return Err(FlowError::InvalidArgument( - "Gemini parts item 'text' must be a string".into(), + "Gemini functionResponse.response must be an object".into(), )); } + Some(_) => {} + } + if let Some(parts) = response.get("parts") { + for part in parts.as_array().ok_or_else(|| { + FlowError::InvalidArgument("Gemini functionResponse.parts must be an array".into()) + })? { + validate_gemini_nested_function_response_part(part)?; + } } + parse_optional_id(response, "functionResponse")?; + Ok(()) +} - // A content item must not mix functionResponse and functionCall parts. - if !fr_parts.is_empty() && !fn_call_parts.is_empty() { +fn validate_gemini_function_call_part(part: &Json) -> Result<()> { + let call = part + .get("functionCall") + .and_then(Json::as_object) + .ok_or_else(|| { + FlowError::InvalidArgument("Gemini functionCall must be an object".into()) + })?; + if call + .get("name") + .and_then(Json::as_str) + .is_none_or(str::is_empty) + { return Err(FlowError::InvalidArgument( - "Gemini contents item must not contain both functionResponse and functionCall parts" - .into(), + "Gemini functionCall is missing a non-empty 'name'".into(), )); } - // functionResponse belongs in user-role turns only. - if !fr_parts.is_empty() && role != "user" { + if call.get("args").is_some_and(|args| !args.is_object()) { + return Err(FlowError::InvalidArgument( + "Gemini functionCall.args must be an object".into(), + )); + } + Ok(()) +} + +fn validate_gemini_content_roles(role: &str, responses: &[&Json], calls: &[&Json]) -> Result<()> { + if !responses.is_empty() && role != "user" { return Err(FlowError::InvalidArgument(format!( "Gemini functionResponse parts must be in a 'user' role content item, got '{role}'" ))); } - // functionCall belongs in model-role turns only. - if !fn_call_parts.is_empty() && role != "model" { + if !calls.is_empty() && role != "model" { return Err(FlowError::InvalidArgument(format!( "Gemini functionCall parts must be in a 'model' role content item, got '{role}'" ))); } + Ok(()) +} - // --- functionResponse parts (tool results sent as a user turn) --- - if !fr_parts.is_empty() { - // Reject sibling visible/native parts — they would be silently lost because we - // return only tool messages from this branch. - let has_visible_or_native_content = parts.iter().any(|p| { - p.get("functionResponse").is_none() - && p.get("thought").and_then(Json::as_bool) != Some(true) - }); - if has_visible_or_native_content { - return Err(FlowError::InvalidArgument( - "Gemini contents item must not mix functionResponse with visible/native parts" - .into(), - )); - } - let mut msgs = Vec::with_capacity(fr_parts.len()); - for fr_part in fr_parts { - let fr = fr_part.get("functionResponse").unwrap(); - let name = fr.get("name").and_then(|v| v.as_str()).unwrap().to_string(); - let fr_obj = fr.as_object().unwrap(); // validated above - let id = parse_optional_id(fr_obj, "functionResponse")?.unwrap_or_else(|| name.clone()); - let content = gemini_function_response_to_message_content(fr)?; - msgs.push(Message::Tool { - content, - tool_call_id: id, - }); - } - return Ok(msgs); +fn gemini_function_response_messages(parts: &[Json], responses: &[&Json]) -> Result> { + if parts.iter().any(|part| { + part.get("functionResponse").is_none() + && part.get("thought").and_then(Json::as_bool) != Some(true) + }) { + return Err(FlowError::InvalidArgument( + "Gemini contents item must not mix functionResponse with visible/native parts".into(), + )); } - - // Thought parts carry internal reasoning — exclude from visible content. - let content_opt = gemini_parts_to_message_content(parts, "request")?; - - // --- functionCall parts (model invoking a tool) --- - if !fn_call_parts.is_empty() { - let mut tool_calls: Vec = Vec::with_capacity(fn_call_parts.len()); - for p in &fn_call_parts { - let fc = p.get("functionCall").unwrap(); // guaranteed by validation loop above - let fc_map = fc.as_object().unwrap(); - let name = fc_map + responses + .iter() + .map(|part| { + let response = part + .get("functionResponse") + .and_then(Json::as_object) + .unwrap(); + let name = response .get("name") .and_then(Json::as_str) .unwrap() .to_string(); - let id = parse_optional_id(fc_map, "functionCall")?.unwrap_or_else(|| name.clone()); - let args = fc_map + let id = parse_optional_id(response, "functionResponse")?.unwrap_or(name); + Ok(Message::Tool { + content: gemini_function_response_to_message_content( + part.get("functionResponse").unwrap(), + )?, + tool_call_id: id, + }) + }) + .collect() +} + +fn gemini_function_call_messages( + content: Option, + calls: &[&Json], +) -> Result> { + let tool_calls = calls + .iter() + .map(|part| { + let call = part.get("functionCall").and_then(Json::as_object).unwrap(); + let name = call.get("name").and_then(Json::as_str).unwrap().to_string(); + let id = parse_optional_id(call, "functionCall")?.unwrap_or_else(|| name.clone()); + let args = call .get("args") .cloned() .unwrap_or_else(|| Json::Object(Default::default())); - let arguments = serde_json::to_string(&args).unwrap_or_else(|_| "{}".into()); - tool_calls.push(ToolCall { + Ok(ToolCall { id, call_type: "function".into(), - function: FunctionCall { name, arguments }, - }); - } - - return Ok(vec![Message::Assistant { - content: content_opt, - tool_calls: Some(tool_calls), - name: None, - }]); - } + function: FunctionCall { + name, + arguments: serde_json::to_string(&args).unwrap_or_else(|_| "{}".into()), + }, + }) + }) + .collect::>>()?; + Ok(vec![Message::Assistant { + content, + tool_calls: Some(tool_calls), + name: None, + }]) +} - // --- plain text message (user or model) --- - let content = content_opt.unwrap_or_else(|| MessageContent::Text(String::new())); - let msg = if role == "model" { +fn gemini_plain_message(role: &str, content: Option) -> Message { + let content = content.unwrap_or_else(|| MessageContent::Text(String::new())); + if role == "model" { Message::Assistant { content: Some(content), tool_calls: None, @@ -778,8 +795,7 @@ fn gemini_content_to_messages(content: &Json) -> Result> { content, name: None, } - }; - Ok(vec![msg]) + } } /// Build the `functionCall` JSON object for a single normalized tool call. @@ -1362,186 +1378,10 @@ impl LlmCodec for GeminiGenerateContentCodec { .as_object() .ok_or_else(|| FlowError::Internal("request content is not an object".into()))?; - let mut messages: Vec = Vec::new(); - - // Validate and decode systemInstruction when present. - if let Some(sys_val) = obj.get("systemInstruction") { - validate_system_instruction(sys_val)?; - if let Some(text) = system_instruction_text(sys_val) { - let msg = serde_json::from_value::( - serde_json::json!({"role": "system", "content": text}), - ) - .map_err(|e| { - FlowError::Internal(format!("Gemini system instruction decode: {e}")) - })?; - messages.push(msg); - } - } + let messages = decode_gemini_messages(obj)?; + let params = decode_gemini_generation_params(obj)?; - // `contents` is required; a present-but-non-array value is malformed. - match obj.get("contents") { - None => { - return Err(FlowError::InvalidArgument( - "Gemini request is missing contents".into(), - )); - } - Some(v) if !v.is_array() => { - return Err(FlowError::InvalidArgument( - "Gemini request contents must be an array".into(), - )); - } - Some(arr) => { - for content in arr.as_array().unwrap() { - messages.extend(gemini_content_to_messages(content)?); - } - } - } - - // generationConfig → GenerationParams - let gen_config = match obj.get("generationConfig") { - Some(v) if !v.is_object() => { - return Err(FlowError::InvalidArgument( - "Gemini generationConfig must be an object".into(), - )); - } - other => other, - }; - let temperature = match gen_config.and_then(|c| c.get("temperature")) { - None => None, - Some(v) => Some(v.as_f64().ok_or_else(|| { - FlowError::InvalidArgument("Gemini temperature must be a number".into()) - })?), - }; - let top_p = match gen_config.and_then(|c| c.get("topP")) { - None => None, - Some(v) => Some(v.as_f64().ok_or_else(|| { - FlowError::InvalidArgument("Gemini topP must be a number".into()) - })?), - }; - let max_tokens = match gen_config.and_then(|c| c.get("maxOutputTokens")) { - None => None, - Some(v) => Some(v.as_u64().ok_or_else(|| { - FlowError::InvalidArgument( - "Gemini maxOutputTokens must be a non-negative integer".into(), - ) - })?), - }; - let stop = match gen_config.and_then(|c| c.get("stopSequences")) { - None => None, - Some(v) => { - let parsed = serde_json::from_value::>(v.clone()).ok(); - if parsed.is_none() { - return Err(FlowError::InvalidArgument( - "Gemini stopSequences must be an array of strings".into(), - )); - } - parsed - } - }; - - let params = - if temperature.is_some() || max_tokens.is_some() || top_p.is_some() || stop.is_some() { - Some(GenerationParams { - temperature, - max_tokens, - top_p, - stop, - }) - } else { - None - }; - - // tools[].functionDeclarations → Vec - // Non-modeled fields (parametersJsonSchema, responseJsonSchema, response, behavior, …) - // are captured into FunctionDefinition.extra so they survive encode. - const MODELED_FD_KEYS: &[&str] = &["name", "description", "parameters"]; - - let tools: Option> = match obj.get("tools") { - None => None, - Some(v) if !v.is_array() => { - return Err(FlowError::InvalidArgument( - "Gemini tools must be an array".into(), - )); - } - Some(v) => { - let arr = v.as_array().unwrap(); - let mut defs: Vec = Vec::new(); - for group in arr { - let group_obj = group.as_object().ok_or_else(|| { - FlowError::InvalidArgument("Gemini tools[] entry must be an object".into()) - })?; - let has_function_declarations = group_obj.contains_key("functionDeclarations"); - if let Some(fds_val) = group_obj.get("functionDeclarations") { - let fds = fds_val.as_array().ok_or_else(|| { - FlowError::InvalidArgument( - "Gemini functionDeclarations must be an array".into(), - ) - })?; - for fd in fds { - if !fd.is_object() { - return Err(FlowError::InvalidArgument( - "Gemini functionDeclaration entry must be an object".into(), - )); - } - let name = fd - .get("name") - .and_then(|n| n.as_str()) - .filter(|s| !s.is_empty()) - .ok_or_else(|| { - FlowError::InvalidArgument( - "Gemini functionDeclaration must have a non-empty 'name'" - .into(), - ) - })? - .to_string(); - let description = match fd.get("description") { - None => None, - Some(Json::String(s)) => Some(s.clone()), - Some(_) => { - return Err(FlowError::InvalidArgument( - "Gemini functionDeclaration.description must be a string" - .into(), - )); - } - }; - let parameters = fd.get("parameters").cloned(); - let extra: serde_json::Map = fd - .as_object() - .map(|o| { - o.iter() - .filter(|(k, _)| !MODELED_FD_KEYS.contains(&k.as_str())) - .map(|(k, v)| (k.clone(), v.clone())) - .collect() - }) - .unwrap_or_default(); - defs.push(ToolDefinition::Function { - function: FunctionDefinition { - name, - description, - parameters, - strict: None, - extra, - }, - extra: Default::default(), - }); - } - } - let native_value = if has_function_declarations { - gemini_native_tool_fields(group_obj) - } else { - Some(group.clone()) - }; - if let Some(value) = native_value { - defs.push(ToolDefinition::ProviderNative { - provider: GEMINI_PROVIDER.into(), - kind: gemini_native_tool_kind(&value), - value, - }); - } - } - if defs.is_empty() { None } else { Some(defs) } - } - }; + let tools = decode_gemini_tools(obj)?; // All unrecognized top-level keys go into extra. let extra: serde_json::Map = obj @@ -1550,15 +1390,7 @@ impl LlmCodec for GeminiGenerateContentCodec { .map(|(k, v)| (k.clone(), v.clone())) .collect(); - let model = match obj.get("model") { - None => None, - Some(Json::String(s)) => Some(s.clone()), - Some(_) => { - return Err(FlowError::InvalidArgument( - "Gemini request 'model' must be a string".into(), - )); - } - }; + let model = decode_gemini_model(obj)?; Ok(AnnotatedLlmRequest { messages, @@ -2318,33 +2150,13 @@ fn merge_gemini_original_parts( let mut content_emitted = false; for orig_part in orig_parts { - let is_thought = orig_part.get("thought").and_then(Json::as_bool) == Some(true); - let is_content_part = !is_thought - && orig_part.get("functionCall").is_none() - && orig_part.get("functionResponse").is_none(); - - if orig_part.get("functionCall").is_some() { - if let Some(rebuilt) = rebuilt_fn_calls.next() { - parts.push(rebuilt); - } - } else if is_content_part { - if let Some(replacement) = new_content_parts.next() { - parts.push(replacement_content_part( - orig_part, - replacement, - content_is_parts_form, - )); - content_emitted = true; - } else if !content_is_parts_form - && (orig_part.get("text").is_none() - || (orig_part.get("text").is_some() - && orig_part.get("thoughtSignature").is_some())) - { - parts.push(orig_part.clone()); - } - } else { - parts.push(orig_part.clone()); - } + content_emitted |= merge_gemini_original_part( + orig_part, + &mut new_content_parts, + &mut rebuilt_fn_calls, + content_is_parts_form, + &mut parts, + ); } let remaining_content_parts: Vec = new_content_parts.collect(); @@ -2362,6 +2174,218 @@ fn merge_gemini_original_parts( parts } +fn merge_gemini_original_part( + orig_part: &Json, + new_content_parts: &mut impl Iterator, + rebuilt_fn_calls: &mut impl Iterator, + content_is_parts_form: bool, + parts: &mut Vec, +) -> bool { + if orig_part.get("functionCall").is_some() { + if let Some(rebuilt) = rebuilt_fn_calls.next() { + parts.push(rebuilt); + } + return false; + } + let is_thought = orig_part.get("thought").and_then(Json::as_bool) == Some(true); + let is_content_part = !is_thought && orig_part.get("functionResponse").is_none(); + if !is_content_part { + parts.push(orig_part.clone()); + return false; + } + if let Some(replacement) = new_content_parts.next() { + parts.push(replacement_content_part( + orig_part, + replacement, + content_is_parts_form, + )); + return true; + } + if !content_is_parts_form + && (orig_part.get("text").is_none() + || (orig_part.get("text").is_some() && orig_part.get("thoughtSignature").is_some())) + { + parts.push(orig_part.clone()); + } + false +} + +fn decode_gemini_messages(obj: &serde_json::Map) -> Result> { + let mut messages = Vec::new(); + if let Some(system) = obj.get("systemInstruction") { + validate_system_instruction(system)?; + if let Some(text) = system_instruction_text(system) { + messages.push( + serde_json::from_value(serde_json::json!({ + "role": "system", + "content": text, + })) + .map_err(|e| { + FlowError::Internal(format!("Gemini system instruction decode: {e}")) + })?, + ); + } + } + let contents = obj + .get("contents") + .ok_or_else(|| FlowError::InvalidArgument("Gemini request is missing contents".into()))?; + let contents = contents.as_array().ok_or_else(|| { + FlowError::InvalidArgument("Gemini request contents must be an array".into()) + })?; + for content in contents { + messages.extend(gemini_content_to_messages(content)?); + } + Ok(messages) +} + +fn decode_gemini_generation_params( + obj: &serde_json::Map, +) -> Result> { + let config = match obj.get("generationConfig") { + Some(value) if !value.is_object() => { + return Err(FlowError::InvalidArgument( + "Gemini generationConfig must be an object".into(), + )); + } + value => value, + }; + let temperature = + decode_gemini_f64(config, "temperature", "Gemini temperature must be a number")?; + let top_p = decode_gemini_f64(config, "topP", "Gemini topP must be a number")?; + let max_tokens = config + .and_then(|value| value.get("maxOutputTokens")) + .map(|value| { + value.as_u64().ok_or_else(|| { + FlowError::InvalidArgument( + "Gemini maxOutputTokens must be a non-negative integer".into(), + ) + }) + }) + .transpose()?; + let stop = config + .and_then(|value| value.get("stopSequences")) + .map(|value| { + serde_json::from_value::>(value.clone()).map_err(|_| { + FlowError::InvalidArgument( + "Gemini stopSequences must be an array of strings".into(), + ) + }) + }) + .transpose()?; + if temperature.is_none() && top_p.is_none() && max_tokens.is_none() && stop.is_none() { + Ok(None) + } else { + Ok(Some(GenerationParams { + temperature, + max_tokens, + top_p, + stop, + })) + } +} + +fn decode_gemini_f64(config: Option<&Json>, key: &str, error: &str) -> Result> { + config + .and_then(|value| value.get(key)) + .map(|value| { + value + .as_f64() + .ok_or_else(|| FlowError::InvalidArgument(error.into())) + }) + .transpose() +} + +fn decode_gemini_tools(obj: &serde_json::Map) -> Result>> { + let Some(value) = obj.get("tools") else { + return Ok(None); + }; + let groups = value + .as_array() + .ok_or_else(|| FlowError::InvalidArgument("Gemini tools must be an array".into()))?; + let mut definitions = Vec::new(); + for group in groups { + decode_gemini_tool_group(group, &mut definitions)?; + } + Ok((!definitions.is_empty()).then_some(definitions)) +} + +fn decode_gemini_tool_group(group: &Json, definitions: &mut Vec) -> Result<()> { + let group_obj = group.as_object().ok_or_else(|| { + FlowError::InvalidArgument("Gemini tools[] entry must be an object".into()) + })?; + let has_declarations = group_obj.contains_key("functionDeclarations"); + if let Some(value) = group_obj.get("functionDeclarations") { + let declarations = value.as_array().ok_or_else(|| { + FlowError::InvalidArgument("Gemini functionDeclarations must be an array".into()) + })?; + for declaration in declarations { + definitions.push(decode_gemini_function_definition(declaration)?); + } + } + let native_value = has_declarations + .then(|| gemini_native_tool_fields(group_obj)) + .flatten() + .or_else(|| (!has_declarations).then(|| group.clone())); + if let Some(value) = native_value { + definitions.push(ToolDefinition::ProviderNative { + provider: GEMINI_PROVIDER.into(), + kind: gemini_native_tool_kind(&value), + value, + }); + } + Ok(()) +} + +fn decode_gemini_function_definition(fd: &Json) -> Result { + const MODELED_FUNCTION_DEFINITION_KEYS: &[&str] = &["name", "description", "parameters"]; + let fd_obj = fd.as_object().ok_or_else(|| { + FlowError::InvalidArgument("Gemini functionDeclaration entry must be an object".into()) + })?; + let name = fd_obj + .get("name") + .and_then(Json::as_str) + .filter(|s| !s.is_empty()) + .ok_or_else(|| { + FlowError::InvalidArgument( + "Gemini functionDeclaration must have a non-empty 'name'".into(), + ) + })?; + let description = match fd_obj.get("description") { + None => None, + Some(Json::String(value)) => Some(value.clone()), + Some(_) => { + return Err(FlowError::InvalidArgument( + "Gemini functionDeclaration.description must be a string".into(), + )); + } + }; + let extra = fd_obj + .iter() + .filter(|(key, _)| !MODELED_FUNCTION_DEFINITION_KEYS.contains(&key.as_str())) + .map(|(key, value)| (key.clone(), value.clone())) + .collect(); + Ok(ToolDefinition::Function { + function: FunctionDefinition { + name: name.to_string(), + description, + parameters: fd_obj.get("parameters").cloned(), + strict: None, + extra, + }, + extra: Default::default(), + }) +} + +fn decode_gemini_model(obj: &serde_json::Map) -> Result> { + match obj.get("model") { + None => Ok(None), + Some(Json::String(value)) => Ok(Some(value.clone())), + Some(_) => Err(FlowError::InvalidArgument( + "Gemini request 'model' must be a string".into(), + )), + } +} + fn patch_gemini_visible_content( orig_parts: &[Json], msg_obj: &serde_json::Map, @@ -2476,75 +2500,8 @@ fn patch_gemini_tools( obj: &mut serde_json::Map, tools: Option<&Vec>, ) -> Result<()> { - // Build functionDeclaration objects and provider-native Gemini tool groups. - let fn_declarations: Vec = { - let mut out = Vec::new(); - if let Some(ts) = tools { - for td in ts { - match td { - ToolDefinition::Function { function: fd, .. } => { - if fd.name.is_empty() { - return Err(FlowError::InvalidArgument( - "Gemini encoder: FunctionDefinition.name must be non-empty".into(), - )); - } - if fd.strict.is_some() { - return Err(FlowError::InvalidArgument( - "Gemini encoder: FunctionDefinition.strict is not supported; \ - remove it or use a provider-native extra field" - .into(), - )); - } - // Start with extra (provider-native fields), then overlay modeled fields. - let mut fdobj: serde_json::Map = fd.extra.clone(); - fdobj.insert("name".into(), Json::String(fd.name.clone())); - if let Some(ref desc) = fd.description { - fdobj.insert("description".into(), Json::String(desc.clone())); - } - if let Some(ref params) = fd.parameters { - fdobj.insert("parameters".into(), params.clone()); - } - out.push(Json::Object(fdobj)); - } - ToolDefinition::ProviderNative { - provider, kind: _, .. - } if provider == GEMINI_PROVIDER => {} - ToolDefinition::ProviderNative { provider, kind, .. } => { - return Err(FlowError::InvalidArgument(format!( - "Gemini encoder: ProviderNative tool '{kind}' (provider '{provider}') \ - cannot be represented on the Gemini surface" - ))); - } - } - } - } - out - }; - let native_groups: Vec = tools - .into_iter() - .flatten() - .filter_map(|td| match td { - ToolDefinition::ProviderNative { - provider, value, .. - } if provider == GEMINI_PROVIDER => Some(value), - _ => None, - }) - .map(|value| { - if !value.is_object() { - return Err(FlowError::InvalidArgument( - "Gemini encoder: ProviderNative tool value must be an object".into(), - )); - } - if value.get("functionDeclarations").is_some() { - return Err(FlowError::InvalidArgument( - "Gemini encoder: ProviderNative tool value must not contain \ - functionDeclarations; use ToolDefinition::Function instead" - .into(), - )); - } - Ok(value.clone()) - }) - .collect::>>()?; + let fn_declarations = gemini_function_declarations(tools)?; + let native_groups = gemini_native_tool_groups(tools)?; // Walk the original tools array in order: replace the FIRST functionDeclarations // group with the rebuilt list, merge any native sibling fields through the @@ -2573,45 +2530,110 @@ fn patch_gemini_tools( recovered. Use provider-native extra fields to manage multi-group tools." ))); } - let mut new_groups: Vec = Vec::with_capacity(orig_tools.len()); + let new_groups = rebuild_gemini_tool_groups(&orig_tools, &fn_declarations, &native_groups); + + if new_groups.is_empty() { + obj.remove("tools"); + } else { + obj.insert("tools".into(), Json::Array(new_groups)); + } + Ok(()) +} + +fn gemini_function_declarations(tools: Option<&Vec>) -> Result> { + tools.into_iter().flatten().filter_map(|tool| match tool { + ToolDefinition::Function { function, .. } => Some(gemini_function_declaration(function)), + ToolDefinition::ProviderNative { provider, .. } if provider == GEMINI_PROVIDER => None, + ToolDefinition::ProviderNative { provider, kind, .. } => Some(Err(FlowError::InvalidArgument( + format!("Gemini encoder: ProviderNative tool '{kind}' (provider '{provider}') cannot be represented on the Gemini surface") + ))), + }).collect() +} + +fn gemini_function_declaration(fd: &FunctionDefinition) -> Result { + if fd.name.is_empty() { + return Err(FlowError::InvalidArgument( + "Gemini encoder: FunctionDefinition.name must be non-empty".into(), + )); + } + if fd.strict.is_some() { + return Err(FlowError::InvalidArgument( + "Gemini encoder: FunctionDefinition.strict is not supported; remove it or use a provider-native extra field".into(), + )); + } + let mut object = fd.extra.clone(); + object.insert("name".into(), Json::String(fd.name.clone())); + if let Some(description) = &fd.description { + object.insert("description".into(), Json::String(description.clone())); + } + if let Some(parameters) = &fd.parameters { + object.insert("parameters".into(), parameters.clone()); + } + Ok(Json::Object(object)) +} + +fn gemini_native_tool_groups(tools: Option<&Vec>) -> Result> { + tools + .into_iter() + .flatten() + .filter_map(|tool| match tool { + ToolDefinition::ProviderNative { + provider, value, .. + } if provider == GEMINI_PROVIDER => Some(validate_gemini_native_tool_group(value)), + _ => None, + }) + .collect() +} + +fn validate_gemini_native_tool_group(value: &Json) -> Result { + if !value.is_object() { + return Err(FlowError::InvalidArgument( + "Gemini encoder: ProviderNative tool value must be an object".into(), + )); + } + if value.get("functionDeclarations").is_some() { + return Err(FlowError::InvalidArgument( + "Gemini encoder: ProviderNative tool value must not contain functionDeclarations; use ToolDefinition::Function instead".into(), + )); + } + Ok(value.clone()) +} + +fn rebuild_gemini_tool_groups( + orig_tools: &[Json], + fn_declarations: &[Json], + native_groups: &[Json], +) -> Vec { + let mut new_groups = Vec::with_capacity(orig_tools.len()); let mut fn_group_placed = false; let mut native_used = vec![false; native_groups.len()]; - for orig_group in &orig_tools { + for orig_group in orig_tools { if orig_group.get("functionDeclarations").is_some() { if !fn_group_placed { - let native_sibling_keys = gemini_native_tool_keys(orig_group); let mut group = serde_json::Map::new(); if !fn_declarations.is_empty() { group.insert( "functionDeclarations".into(), - Json::Array(fn_declarations.clone()), + Json::Array(fn_declarations.to_vec()), ); } - if !native_sibling_keys.is_empty() - && let Some(native_group) = take_matching_native_group( - &native_groups, - &mut native_used, - &native_sibling_keys, - ) - && let Some(native_obj) = native_group.as_object() - { - group.extend(native_obj.clone()); - } + merge_gemini_native_sibling_group( + orig_group, + native_groups, + &mut native_used, + &mut group, + ); if !group.is_empty() { new_groups.push(Json::Object(group)); } fn_group_placed = true; } - // Unreachable: the fn_decl_group_count > 1 guard above returns Err before - // this loop when tools changed, and this function is only called when tools - // changed. A second functionDeclarations group can never be reached here. - } else { - let native_keys = gemini_native_tool_keys(orig_group); - if let Some(native_group) = - take_matching_native_group(&native_groups, &mut native_used, &native_keys) - { - new_groups.push(native_group); - } + } else if let Some(native_group) = take_matching_native_group( + native_groups, + &mut native_used, + &gemini_native_tool_keys(orig_group), + ) { + new_groups.push(native_group); } } if !fn_group_placed && !fn_declarations.is_empty() { @@ -2621,16 +2643,24 @@ fn patch_gemini_tools( native_groups .iter() .enumerate() - .filter(|(idx, _)| !native_used[*idx]) + .filter(|(index, _)| !native_used[*index]) .map(|(_, group)| group.clone()), ); + new_groups +} - if new_groups.is_empty() { - obj.remove("tools"); - } else { - obj.insert("tools".into(), Json::Array(new_groups)); +fn merge_gemini_native_sibling_group( + orig_group: &Json, + native_groups: &[Json], + native_used: &mut [bool], + group: &mut serde_json::Map, +) { + let keys = gemini_native_tool_keys(orig_group); + if let Some(native_group) = take_matching_native_group(native_groups, native_used, &keys) + && let Some(native_obj) = native_group.as_object() + { + group.extend(native_obj.clone()); } - Ok(()) } /// Overlay extra-field changes from `annotated` onto `obj`, guided by `baseline`. @@ -2745,113 +2775,8 @@ impl GeminiGenerateContentStreamingState { } fn observe(&mut self, event: &Json) -> Result<()> { - if let Some(candidates) = event.get("candidates").and_then(Json::as_array) - && let Some(candidate) = candidates.first() - { - if candidates.len() > 1 { - return Err(FlowError::InvalidArgument( - "Gemini streaming chunks with multiple candidates are not supported".into(), - )); - } - let candidate_obj = candidate.as_object().ok_or_else(|| { - FlowError::InvalidArgument("Gemini streaming candidate must be an object".into()) - })?; - let index = candidate_obj - .get("index") - .ok_or_else(|| { - FlowError::InvalidArgument( - "Gemini streaming candidate index is required".into(), - ) - })? - .as_u64() - .ok_or_else(|| { - FlowError::InvalidArgument( - "Gemini streaming candidate index must be an unsigned integer".into(), - ) - })?; - if let Some(previous_index) = self.candidate_index { - if previous_index != index { - return Err(FlowError::InvalidArgument( - "Gemini streaming candidate index changed across chunks".into(), - )); - } - } else { - if index != 0 { - return Err(FlowError::InvalidArgument( - "Gemini streaming only supports candidate index 0".into(), - )); - } - self.candidate_index = Some(index); - } - - if let Some(parts) = candidate_obj - .get("content") - .and_then(|c| c.get("parts")) - .and_then(Json::as_array) - { - for part in parts { - if !part.is_object() { - return Err(FlowError::InvalidArgument( - "Gemini streaming parts entry must be an object".into(), - )); - } - let part_obj = part.as_object().unwrap(); - let data_key = validate_single_gemini_part_data_field(part_obj, "streaming")?; - // Preserve thought parts in the provider-native aggregate; the response - // decoder filters them out of the normalized message. - if part.get("thought").and_then(Json::as_bool) == Some(true) { - self.parts.push(part.clone()); - continue; - } - - match data_key { - Some("text") => { - let text_val = part.get("text").unwrap(); - match text_val.as_str() { - Some(s) => { - self.push_text_part(s, part_obj); - } - None => { - return Err(FlowError::InvalidArgument( - "Gemini streaming parts[].text must be a string".into(), - )); - } - } - } - Some("functionCall") => { - self.parts.push(part.clone()); - } - Some("functionResponse") => { - return Err(FlowError::InvalidArgument( - "Gemini streaming response parts must not contain functionResponse" - .into(), - )); - } - Some(_) | None => { - self.parts.push(part.clone()); - } - } - } - } - if let Some(reason_val) = candidate.get("finishReason") { - match reason_val.as_str() { - Some(s) => self.finish_reason = Some(s.to_string()), - None => { - return Err(FlowError::InvalidArgument( - "Gemini streaming candidate finishReason must be a string".into(), - )); - } - } - } - // Collect candidate-level metadata fields (safetyRatings, groundingMetadata, - // citationMetadata, avgLogprobs, etc.) that non-streaming decode preserves - // in ApiSpecificResponse::GeminiGenerateContent. Later chunks overwrite earlier ones for - // the same key (last-wins), matching the non-streaming behaviour. - for (k, v) in candidate_obj { - if !matches!(k.as_str(), "content" | "finishReason" | "index") { - self.candidate_extra.insert(k.clone(), v.clone()); - } - } + if let Some(candidates) = event.get("candidates").and_then(Json::as_array) { + self.observe_candidate(candidates)?; } if let Some(usage) = event.get("usageMetadata") { self.usage_metadata = Some(usage.clone()); @@ -2879,6 +2804,118 @@ impl GeminiGenerateContentStreamingState { Ok(()) } + fn observe_candidate(&mut self, candidates: &[Json]) -> Result<()> { + let Some(candidate) = candidates.first() else { + return Ok(()); + }; + if candidates.len() > 1 { + return Err(FlowError::InvalidArgument( + "Gemini streaming chunks with multiple candidates are not supported".into(), + )); + } + let candidate_obj = candidate.as_object().ok_or_else(|| { + FlowError::InvalidArgument("Gemini streaming candidate must be an object".into()) + })?; + self.observe_candidate_index(candidate_obj)?; + if let Some(parts) = candidate_obj + .get("content") + .and_then(|content| content.get("parts")) + .and_then(Json::as_array) + { + self.observe_parts(parts)?; + } + self.observe_finish_reason(candidate)?; + for (key, value) in candidate_obj { + if !matches!(key.as_str(), "content" | "finishReason" | "index") { + self.candidate_extra.insert(key.clone(), value.clone()); + } + } + Ok(()) + } + + fn observe_candidate_index(&mut self, candidate: &serde_json::Map) -> Result<()> { + let index = candidate + .get("index") + .ok_or_else(|| { + FlowError::InvalidArgument("Gemini streaming candidate index is required".into()) + })? + .as_u64() + .ok_or_else(|| { + FlowError::InvalidArgument( + "Gemini streaming candidate index must be an unsigned integer".into(), + ) + })?; + match self.candidate_index { + Some(previous) if previous != index => Err(FlowError::InvalidArgument( + "Gemini streaming candidate index changed across chunks".into(), + )), + Some(_) => Ok(()), + None if index == 0 => { + self.candidate_index = Some(index); + Ok(()) + } + None => Err(FlowError::InvalidArgument( + "Gemini streaming only supports candidate index 0".into(), + )), + } + } + + fn observe_parts(&mut self, parts: &[Json]) -> Result<()> { + for part in parts { + let part_obj = part.as_object().ok_or_else(|| { + FlowError::InvalidArgument("Gemini streaming parts entry must be an object".into()) + })?; + let data_key = validate_single_gemini_part_data_field(part_obj, "streaming")?; + if part.get("thought").and_then(Json::as_bool) == Some(true) { + self.parts.push(part.clone()); + } else { + self.observe_part(part, part_obj, data_key)?; + } + } + Ok(()) + } + + fn observe_part( + &mut self, + part: &Json, + part_obj: &serde_json::Map, + data_key: Option<&str>, + ) -> Result<()> { + match data_key { + Some("text") => { + let text = part.get("text").and_then(Json::as_str).ok_or_else(|| { + FlowError::InvalidArgument( + "Gemini streaming parts[].text must be a string".into(), + ) + })?; + self.push_text_part(text, part_obj); + } + Some("functionResponse") => { + return Err(FlowError::InvalidArgument( + "Gemini streaming response parts must not contain functionResponse".into(), + )); + } + _ => self.parts.push(part.clone()), + } + Ok(()) + } + + fn observe_finish_reason(&mut self, candidate: &Json) -> Result<()> { + if let Some(reason) = candidate.get("finishReason") { + self.finish_reason = Some( + reason + .as_str() + .ok_or_else(|| { + FlowError::InvalidArgument( + "Gemini streaming candidate finishReason must be a string".into(), + ) + })? + .to_string(), + ); + } + Ok(()) + } + fn finalize(self) -> Json { let mut candidate_obj = serde_json::Map::new(); candidate_obj.insert( diff --git a/crates/core/src/observability/plugin_component.rs b/crates/core/src/observability/plugin_component.rs index 8485561a4..b1c48bc4b 100644 --- a/crates/core/src/observability/plugin_component.rs +++ b/crates/core/src/observability/plugin_component.rs @@ -943,63 +943,70 @@ fn register_atif_dispatcher( "observability", ctx.qualify_name("atif.shutdown"), Box::new(move || { - let work = match (|| -> PluginResult<_> { - let work = { - let mut guard = manager.lock().map_err(|err| { - PluginError::Internal(format!("ATIF dispatcher lock poisoned: {err}")) - })?; - guard.flush_open_agents() - }; - for (scope_uuid, name) in &work.scope_subscribers { - deregister_atif_shutdown_subscriber(scope_uuid, name)?; - } - Ok(work) - })() { - Ok(work) => work, - Err(error) => return PluginRegistrationCleanupOutcome::NotRemoved(error), - }; - - let delivery = (|| -> PluginResult<()> { - for export in work.exports { - let write = prepare_atif_shutdown_file(&export, Arc::clone(&manager)) - .map_err(observability_registration_error)?; - let agent_uuid = write.agent_uuid; - let targets = { - let guard = manager.lock().map_err(|err| { - PluginError::Internal(format!("ATIF dispatcher lock poisoned: {err}")) - })?; - guard.sink_targets() - }; - let results = write_atif(&write, shutdown_storage.as_slice(), &targets); - let mut guard = manager.lock().map_err(|err| { - PluginError::Internal(format!("ATIF dispatcher lock poisoned: {err}")) - })?; - let _ = guard.complete_scope_write(agent_uuid, results); - } - Ok(()) - })(); - if let Err(error) = delivery { - return PluginRegistrationCleanupOutcome::RemovedWithError(error); - } - let guard = match manager.lock() { - Ok(guard) => guard, - Err(error) => { - return PluginRegistrationCleanupOutcome::RemovedWithError( - PluginError::Internal(format!("ATIF dispatcher lock poisoned: {error}")), - ); - } - }; - match guard.last_error_result() { - Ok(()) => PluginRegistrationCleanupOutcome::Removed, - Err(error) => PluginRegistrationCleanupOutcome::RemovedWithError( - observability_registration_error(error), - ), - } + atif_shutdown_cleanup(Arc::clone(&manager), Arc::clone(&shutdown_storage)) }), )); Ok(()) } +fn atif_shutdown_cleanup( + manager: Arc>, + shutdown_storage: AtifStorageList, +) -> PluginRegistrationCleanupOutcome { + let work = match flush_atif_shutdown_work(&manager) { + Ok(work) => work, + Err(error) => return PluginRegistrationCleanupOutcome::NotRemoved(error), + }; + if let Err(error) = write_atif_shutdown_exports(&manager, &shutdown_storage, work.exports) { + return PluginRegistrationCleanupOutcome::RemovedWithError(error); + } + match manager.lock() { + Ok(guard) => match guard.last_error_result() { + Ok(()) => PluginRegistrationCleanupOutcome::Removed, + Err(error) => PluginRegistrationCleanupOutcome::RemovedWithError( + observability_registration_error(error), + ), + }, + Err(error) => PluginRegistrationCleanupOutcome::RemovedWithError(PluginError::Internal( + format!("ATIF dispatcher lock poisoned: {error}"), + )), + } +} + +fn flush_atif_shutdown_work(manager: &Arc>) -> PluginResult { + let work = { + let mut guard = manager.lock().map_err(|err| { + PluginError::Internal(format!("ATIF dispatcher lock poisoned: {err}")) + })?; + guard.flush_open_agents() + }; + for (scope_uuid, name) in &work.scope_subscribers { + deregister_atif_shutdown_subscriber(scope_uuid, name)?; + } + Ok(work) +} + +fn write_atif_shutdown_exports( + manager: &Arc>, + shutdown_storage: &AtifStorageList, + exports: Vec, +) -> PluginResult<()> { + for export in exports { + let write = prepare_atif_shutdown_file(&export, Arc::clone(manager)) + .map_err(observability_registration_error)?; + let targets = manager + .lock() + .map_err(|err| PluginError::Internal(format!("ATIF dispatcher lock poisoned: {err}")))? + .sink_targets(); + let results = write_atif(&write, shutdown_storage.as_slice(), &targets); + let mut guard = manager.lock().map_err(|err| { + PluginError::Internal(format!("ATIF dispatcher lock poisoned: {err}")) + })?; + let _ = guard.complete_scope_write(write.agent_uuid, results); + } + Ok(()) +} + fn deregister_atif_shutdown_subscriber(scope_uuid: &Uuid, name: &str) -> PluginResult<()> { match scope_deregister_subscriber(scope_uuid, name) { Ok(_) | Err(FlowError::NotFound(_)) => Ok(()), diff --git a/crates/core/src/plugin.rs b/crates/core/src/plugin.rs index 2d6078a3a..ee7bb88a2 100644 --- a/crates/core/src/plugin.rs +++ b/crates/core/src/plugin.rs @@ -1578,102 +1578,175 @@ async fn initialize_plugins_exact_inner( guard.take() }; - if let Some(mut previous_state) = previous { - // Keep the previous report installed while teardown callbacks run so - // runtime diagnostics emitted by teardown remain observable. - { - let mut guard = ACTIVE_PLUGIN_CONFIGURATION.lock().map_err(|err| { - PluginError::Internal(format!("active plugin configuration lock poisoned: {err}")) - })?; - *guard = Some(ActivePluginConfiguration { - config: previous_state.config.clone(), - report: previous_state.report.clone(), - registrations: Vec::new(), - }); - } - let teardown = rollback_registrations_checked(&mut previous_state.registrations); - let teardown_report = ACTIVE_PLUGIN_CONFIGURATION - .lock() - .map_err(|err| { - PluginError::Internal(format!("active plugin configuration lock poisoned: {err}")) - })? - .take() - .map(|state| state.report); - if !teardown.errors.is_empty() { - if let Some(report) = - teardown_report.filter(|report| !report.runtime_diagnostics.is_empty()) - && let Ok(mut guard) = LAST_FAILED_RUNTIME_DIAGNOSTICS_REPORT.lock() - { - *guard = Some(report); - } - if !teardown.callbacks_cleared { - record_rollback_failures(rollback_failures.as_ref(), teardown.errors.clone()); - } - return Err(PluginError::RegistrationFailed(format!( - "previous plugin configuration could not be cleared: {}", - teardown.errors.join("; ") - ))); + match previous { + Some(previous_state) => { + replace_plugin_configuration( + config, + report, + previous_state, + rollback_failures, + enabled_component_count, + ) + .await } - match initialize_plugin_components_catching_panics( - config.clone(), - rollback_failures.clone(), - ) - .await - { - Ok(registrations) => { - store_active_plugin_configuration(config, report.clone(), registrations)?; - log::info!( - target: "nemo_relay.plugin", - event = "plugin_configuration_replaced", - component_count = enabled_component_count; - "Plugin configuration replaced" - ); - Ok(report) - } - Err(err) => match initialize_plugin_components_catching_panics( - previous_state.config.clone(), - rollback_failures.clone(), + None => { + activate_initial_plugin_configuration( + config, + report, + rollback_failures, + enabled_component_count, ) .await - { - Ok(registrations) => { - store_active_plugin_configuration( - previous_state.config, - previous_state.report, - registrations, - )?; - log::warn!( - target: "nemo_relay.plugin", - event = "plugin_configuration_restored", - recovery = "previous_configuration"; - "Plugin activation failed; previous configuration restored" - ); - Err(err) - } - Err(restore_err) => { - log::error!( - target: "nemo_relay.plugin", - event = "plugin_rollback_failed", - recovery = "previous_configuration"; - "Plugin activation failed and the previous configuration could not be restored" - ); - Err(PluginError::RegistrationFailed(format!( - "{err}; previous plugin configuration could not be restored: {restore_err}" - ))) - } - }, } - } else { - let registrations = - initialize_plugin_components_catching_panics(config.clone(), rollback_failures).await?; - store_active_plugin_configuration(config, report.clone(), registrations)?; - log::info!( - target: "nemo_relay.plugin", - event = "plugin_configuration_activated", - component_count = enabled_component_count; - "Plugin configuration activated" - ); - Ok(report) + } +} + +async fn activate_initial_plugin_configuration( + config: PluginConfig, + report: ConfigReport, + rollback_failures: Option>>>, + enabled_component_count: usize, +) -> Result { + let registrations = + initialize_plugin_components_catching_panics(config.clone(), rollback_failures).await?; + store_active_plugin_configuration(config, report.clone(), registrations)?; + log::info!( + target: "nemo_relay.plugin", + event = "plugin_configuration_activated", + component_count = enabled_component_count; + "Plugin configuration activated" + ); + Ok(report) +} + +async fn replace_plugin_configuration( + config: PluginConfig, + report: ConfigReport, + mut previous_state: ActivePluginConfiguration, + rollback_failures: Option>>>, + enabled_component_count: usize, +) -> Result { + install_previous_configuration_for_teardown(&previous_state)?; + let teardown = rollback_registrations_checked(&mut previous_state.registrations); + let teardown_report = take_active_runtime_diagnostics_report()?; + if !teardown.errors.is_empty() { + record_failed_teardown(&teardown, teardown_report, rollback_failures.as_ref()); + return Err(PluginError::RegistrationFailed(format!( + "previous plugin configuration could not be cleared: {}", + teardown.errors.join("; ") + ))); + } + activate_replacement_or_restore( + config, + report, + previous_state, + rollback_failures, + enabled_component_count, + ) + .await +} + +fn install_previous_configuration_for_teardown( + previous_state: &ActivePluginConfiguration, +) -> Result<()> { + let mut guard = ACTIVE_PLUGIN_CONFIGURATION.lock().map_err(|err| { + PluginError::Internal(format!("active plugin configuration lock poisoned: {err}")) + })?; + *guard = Some(ActivePluginConfiguration { + config: previous_state.config.clone(), + report: previous_state.report.clone(), + registrations: Vec::new(), + }); + Ok(()) +} + +fn take_active_runtime_diagnostics_report() -> Result> { + Ok(ACTIVE_PLUGIN_CONFIGURATION + .lock() + .map_err(|err| { + PluginError::Internal(format!("active plugin configuration lock poisoned: {err}")) + })? + .take() + .map(|state| state.report)) +} + +fn record_failed_teardown( + teardown: &PluginRollbackOutcome, + teardown_report: Option, + rollback_failures: Option<&Arc>>>, +) { + if let Some(report) = teardown_report.filter(|report| !report.runtime_diagnostics.is_empty()) + && let Ok(mut guard) = LAST_FAILED_RUNTIME_DIAGNOSTICS_REPORT.lock() + { + *guard = Some(report); + } + if !teardown.callbacks_cleared { + record_rollback_failures(rollback_failures, teardown.errors.clone()); + } +} + +async fn activate_replacement_or_restore( + config: PluginConfig, + report: ConfigReport, + previous_state: ActivePluginConfiguration, + rollback_failures: Option>>>, + enabled_component_count: usize, +) -> Result { + match initialize_plugin_components_catching_panics(config.clone(), rollback_failures.clone()) + .await + { + Ok(registrations) => { + store_active_plugin_configuration(config, report.clone(), registrations)?; + log::info!( + target: "nemo_relay.plugin", + event = "plugin_configuration_replaced", + component_count = enabled_component_count; + "Plugin configuration replaced" + ); + Ok(report) + } + Err(err) => { + restore_previous_plugin_configuration(previous_state, rollback_failures, err).await + } + } +} + +async fn restore_previous_plugin_configuration( + previous_state: ActivePluginConfiguration, + rollback_failures: Option>>>, + err: PluginError, +) -> Result { + match initialize_plugin_components_catching_panics( + previous_state.config.clone(), + rollback_failures, + ) + .await + { + Ok(registrations) => { + store_active_plugin_configuration( + previous_state.config, + previous_state.report, + registrations, + )?; + log::warn!( + target: "nemo_relay.plugin", + event = "plugin_configuration_restored", + recovery = "previous_configuration"; + "Plugin activation failed; previous configuration restored" + ); + Err(err) + } + Err(restore_err) => { + log::error!( + target: "nemo_relay.plugin", + event = "plugin_rollback_failed", + recovery = "previous_configuration"; + "Plugin activation failed and the previous configuration could not be restored" + ); + Err(PluginError::RegistrationFailed(format!( + "{err}; previous plugin configuration could not be restored: {restore_err}" + ))) + } } } diff --git a/crates/core/src/plugins/nemo_guardrails/python.rs b/crates/core/src/plugins/nemo_guardrails/python.rs index 78ebf92a2..baa2ef915 100644 --- a/crates/core/src/plugins/nemo_guardrails/python.rs +++ b/crates/core/src/plugins/nemo_guardrails/python.rs @@ -1308,69 +1308,64 @@ fn streaming_output_blocked(message: String) -> FlowError { fn extract_stream_text(codec: LocalGuardrailsCodec, chunk: &Json) -> Option { let chunk = chunk.as_object()?; match codec { - LocalGuardrailsCodec::OpenAIChat => { - let choices = chunk.get("choices")?.as_array()?; - let mut parts = vec![]; - for choice in choices { - let content = choice - .get("delta") - .and_then(Json::as_object) - .and_then(|delta| delta.get("content")) - .and_then(Json::as_str); - if let Some(content) = content - && !content.is_empty() - { - parts.push(content); - } - } - (!parts.is_empty()).then(|| parts.join("")) - } - LocalGuardrailsCodec::OpenAIResponses => { - if chunk.get("type").and_then(Json::as_str) == Some("response.output_text.delta") { - chunk - .get("delta") - .and_then(Json::as_str) - .filter(|delta| !delta.is_empty()) - .map(str::to_string) - } else { - None - } - } - LocalGuardrailsCodec::AnthropicMessages => { - if chunk.get("type").and_then(Json::as_str) != Some("content_block_delta") { - return None; - } - let delta = chunk.get("delta")?.as_object()?; - if delta.get("type").and_then(Json::as_str) != Some("text_delta") { - return None; - } - delta - .get("text") + LocalGuardrailsCodec::OpenAIChat => extract_openai_chat_stream_text(chunk), + LocalGuardrailsCodec::OpenAIResponses => extract_openai_response_stream_text(chunk), + LocalGuardrailsCodec::AnthropicMessages => extract_anthropic_stream_text(chunk), + LocalGuardrailsCodec::GeminiGenerateContent => extract_gemini_stream_text(chunk), + } +} + +fn extract_openai_chat_stream_text(chunk: &serde_json::Map) -> Option { + let choices = chunk.get("choices")?.as_array()?; + let parts = choices + .iter() + .filter_map(|choice| { + choice + .get("delta") + .and_then(Json::as_object) + .and_then(|delta| delta.get("content")) .and_then(Json::as_str) - .filter(|text| !text.is_empty()) - .map(str::to_string) - } - LocalGuardrailsCodec::GeminiGenerateContent => { - let candidates = chunk.get("candidates")?.as_array()?; - let parts = candidates - .first()? - .get("content")? - .get("parts")? - .as_array()?; - let mut texts = vec![]; - for part in parts { - if part.get("thought").and_then(Json::as_bool) == Some(true) { - continue; - } - if let Some(text) = part.get("text").and_then(Json::as_str) - && !text.is_empty() - { - texts.push(text); - } - } - (!texts.is_empty()).then(|| texts.join("")) - } + .filter(|content| !content.is_empty()) + }) + .collect::>(); + (!parts.is_empty()).then(|| parts.join("")) +} + +fn extract_openai_response_stream_text(chunk: &serde_json::Map) -> Option { + (chunk.get("type").and_then(Json::as_str) == Some("response.output_text.delta")) + .then(|| chunk.get("delta").and_then(Json::as_str)) + .flatten() + .filter(|delta| !delta.is_empty()) + .map(str::to_string) +} + +fn extract_anthropic_stream_text(chunk: &serde_json::Map) -> Option { + if chunk.get("type").and_then(Json::as_str) != Some("content_block_delta") { + return None; } + let delta = chunk.get("delta")?.as_object()?; + (delta.get("type").and_then(Json::as_str) == Some("text_delta")) + .then(|| delta.get("text").and_then(Json::as_str)) + .flatten() + .filter(|text| !text.is_empty()) + .map(str::to_string) +} + +fn extract_gemini_stream_text(chunk: &serde_json::Map) -> Option { + let parts = chunk + .get("candidates")? + .as_array()? + .first()? + .get("content")? + .get("parts")? + .as_array()?; + let texts = parts + .iter() + .filter(|part| part.get("thought").and_then(Json::as_bool) != Some(true)) + .filter_map(|part| part.get("text").and_then(Json::as_str)) + .filter(|text| !text.is_empty()) + .collect::>(); + (!texts.is_empty()).then(|| texts.join("")) } async fn monitor_guardrails_stream( diff --git a/crates/node/src/callable.rs b/crates/node/src/callable.rs index e807979be..422b7862f 100644 --- a/crates/node/src/callable.rs +++ b/crates/node/src/callable.rs @@ -1227,48 +1227,7 @@ pub fn wrap_js_event_subscriber( ) -> napi::Result { let callback = safe_subscriber_callback(env, &callback)?; let queue_error_name = name.clone(); - let mut func = callback.create_threadsafe_function::< - JsSubscriberCallbackCall, - JsUnknown, - _, - ErrorStrategy::CalleeHandled, - >(0, move |ctx: ThreadSafeCallContext| { - let JsSubscriberCallbackCall { event, callback_id } = ctx.value; - let completed = Arc::new(AtomicBool::new(false)); - let completed_callback = Arc::clone(&completed); - let callback_name = name.clone(); - let complete = ctx.env.create_function_from_closure( - "__nemo_relay_complete_subscriber_callback", - move |ctx| { - if completed_callback - .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) - .is_ok() - { - if ctx.length > 0 { - match ctx.get::(0) { - Ok(message) => record_callback_error(format!( - "nemo_relay: JS event subscriber '{callback_name}' failed: {message}" - )), - Err(error) => record_callback_error(format!( - "nemo_relay: failed to read JS event subscriber \ - '{callback_name}' failure: {error}" - )), - } - } - complete_js_subscriber_callback(callback_id); - } - ctx.env.get_undefined() - }, - )?; - let event = unsafe { - JsUnknown::from_raw_unchecked( - ctx.env.raw(), - Json::to_napi_value(ctx.env.raw(), event)?, - ) - }; - let complete = unsafe { JsUnknown::from_raw_unchecked(ctx.env.raw(), complete.raw()) }; - Ok(vec![event, complete]) - })?; + let mut func = create_js_event_subscriber_function(&callback, name)?; func.unref(env)?; let func = Arc::new(func); Ok(Arc::new(move |event: &Event| { @@ -1298,6 +1257,69 @@ pub fn wrap_js_event_subscriber( })) } +fn create_js_event_subscriber_function( + callback: &JsFunction, + name: String, +) -> napi::Result> { + callback.create_threadsafe_function::< + JsSubscriberCallbackCall, + JsUnknown, + _, + ErrorStrategy::CalleeHandled, + >(0, move |ctx: ThreadSafeCallContext| { + let JsSubscriberCallbackCall { event, callback_id } = ctx.value; + let completed = Arc::new(AtomicBool::new(false)); + let completed_callback = Arc::clone(&completed); + let callback_name = name.clone(); + let complete = create_js_subscriber_completion_callback( + &ctx.env, + callback_name, + callback_id, + completed_callback, + )?; + let event = unsafe { + JsUnknown::from_raw_unchecked( + ctx.env.raw(), + Json::to_napi_value(ctx.env.raw(), event)?, + ) + }; + let complete = unsafe { JsUnknown::from_raw_unchecked(ctx.env.raw(), complete.raw()) }; + Ok(vec![event, complete]) + }) +} + +fn create_js_subscriber_completion_callback( + env: &Env, + callback_name: String, + callback_id: u64, + completed_callback: Arc, +) -> napi::Result { + env.create_function_from_closure("__nemo_relay_complete_subscriber_callback", move |ctx| { + if completed_callback + .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) + .is_ok() + { + record_js_subscriber_callback_error(&ctx, &callback_name); + complete_js_subscriber_callback(callback_id); + } + ctx.env.get_undefined() + }) +} + +fn record_js_subscriber_callback_error(ctx: &napi::CallContext, callback_name: &str) { + if ctx.length == 0 { + return; + } + match ctx.get::(0) { + Ok(message) => record_callback_error(format!( + "nemo_relay: JS event subscriber '{callback_name}' failed: {message}" + )), + Err(error) => record_callback_error(format!( + "nemo_relay: failed to read JS event subscriber '{callback_name}' failure: {error}" + )), + } +} + // --------------------------------------------------------------------------- // Codec wrappers // --------------------------------------------------------------------------- diff --git a/crates/pii-redaction/src/overlay.rs b/crates/pii-redaction/src/overlay.rs index dc896053d..5f935a834 100644 --- a/crates/pii-redaction/src/overlay.rs +++ b/crates/pii-redaction/src/overlay.rs @@ -99,12 +99,21 @@ fn overlay_gemini_response(mut payload: Json, annotated: &AnnotatedLlmResponse) return payload; }; - if let Some(message_parts) = gemini_message_parts_for_overlay(annotated.message.as_ref()) { + overlay_gemini_message_parts(parts, annotated.message.as_ref()); + + // Overlay functionCall parts. + overlay_gemini_tool_calls(parts, annotated.tool_calls.as_deref()); + + payload +} + +fn overlay_gemini_message_parts(parts: &mut Vec, message: Option<&MessageContent>) { + if let Some(message_parts) = gemini_message_parts_for_overlay(message) { let mut sanitized = message_parts.into_iter(); parts.retain_mut(|part| { - let is_thought = part.get("thought").and_then(Json::as_bool) == Some(true); - let is_tool_call = part.get("functionCall").is_some(); - if is_thought || is_tool_call { + if part.get("thought").and_then(Json::as_bool) == Some(true) + || part.get("functionCall").is_some() + { return true; } let Some(next) = sanitized.next() else { @@ -114,63 +123,59 @@ fn overlay_gemini_response(mut payload: Json, annotated: &AnnotatedLlmResponse) true }); parts.extend(sanitized); - } else { - // Overlay fallback text into the first visible text part. The normalized - // text may contain embedded newlines, so splitting would confuse text - // content with Gemini part boundaries. - let message_text = annotated_message_text(annotated.message.as_ref()); - let mut wrote_text = false; - parts.retain_mut(|part| { - let is_thought = part.get("thought").and_then(Json::as_bool) == Some(true); - if part.get("text").is_none() || is_thought { - return true; - } - let Some(p) = part.as_object_mut() else { - return false; - }; - if wrote_text { - return false; - } - let Some(text) = message_text.as_deref() else { - return false; - }; - set_optional_string_field(p, "text", Some(text)); - wrote_text = true; - true - }); + return; } - // Overlay functionCall parts. - if let Some(tool_calls) = annotated.tool_calls.as_deref() { - let mut sanitized = tool_calls.iter(); - parts.retain_mut(|part| { - let Some(fc) = part - .as_object_mut() - .and_then(|p| p.get_mut("functionCall")) - .and_then(Json::as_object_mut) - else { - return true; // not a functionCall part — keep - }; - let Some(sc) = sanitized.next() else { - return false; // no sanitized call left — remove this part - }; - if fc.contains_key("id") { - set_optional_string_field(fc, "id", Some(sc.id.as_str())); - } - set_optional_string_field(fc, "name", Some(sc.name.as_str())); - fc.insert("args".into(), sc.arguments.clone()); - true - }); - } else { - // No tool calls in sanitized response: remove all functionCall parts. + let message_text = annotated_message_text(message); + let mut wrote_text = false; + parts.retain_mut(|part| { + if part.get("text").is_none() || part.get("thought").and_then(Json::as_bool) == Some(true) { + return true; + } + let Some(part) = part.as_object_mut() else { + return false; + }; + let Some(text) = message_text.as_deref() else { + return false; + }; + if wrote_text { + return false; + } + set_optional_string_field(part, "text", Some(text)); + wrote_text = true; + true + }); +} + +fn overlay_gemini_tool_calls(parts: &mut Vec, tool_calls: Option<&[ResponseToolCall]>) { + let Some(tool_calls) = tool_calls else { parts.retain(|part| { part.as_object() - .map(|p| !p.contains_key("functionCall")) + .map(|object| !object.contains_key("functionCall")) .unwrap_or(true) }); - } + return; + }; - payload + let mut sanitized = tool_calls.iter(); + parts.retain_mut(|part| { + let Some(function_call) = part + .as_object_mut() + .and_then(|object| object.get_mut("functionCall")) + .and_then(Json::as_object_mut) + else { + return true; + }; + let Some(call) = sanitized.next() else { + return false; + }; + if function_call.contains_key("id") { + set_optional_string_field(function_call, "id", Some(call.id.as_str())); + } + set_optional_string_field(function_call, "name", Some(call.name.as_str())); + function_call.insert("args".into(), call.arguments.clone()); + true + }); } fn overlay_openai_chat_response(mut payload: Json, annotated: &AnnotatedLlmResponse) -> Json { diff --git a/go/nemo_relay/logging_test.go b/go/nemo_relay/logging_test.go index 09af8e799..1c47b02d5 100644 --- a/go/nemo_relay/logging_test.go +++ b/go/nemo_relay/logging_test.go @@ -20,6 +20,8 @@ var loggingEnvironmentNames = map[string]struct{}{ "NEMO_RELAY_LOG_CONFIG_PATH": {}, } +const loggingEnvironmentTestRunArg = "-test.run=TestBindingLoggingEnvironment" + func loggingTestEnvironment(values ...string) []string { environment := make([]string, 0, len(os.Environ())+len(values)) for _, value := range os.Environ() { @@ -41,42 +43,47 @@ func TestBindingLoggingEnvironment(t *testing.T) { return } - t.Run("initializes from environment", func(t *testing.T) { - command := exec.Command(os.Args[0], "-test.run=TestBindingLoggingEnvironment") - command.Env = loggingTestEnvironment( - loggingHelperEnvironment+"=shutdown", - "NEMO_RELAY_LOG=info", - "NEMO_RELAY_LOG_STDERR_FORMAT=jsonl", - ) - output, err := command.CombinedOutput() - if err != nil { - t.Fatalf("binding import failed: %v\n%s", err, output) - } - if !strings.Contains(string(output), `"event":"logging_initialized"`) { - t.Fatalf("logging initialization event missing from output:\n%s", output) - } - }) + t.Run("initializes from environment", testLoggingInitialization) + t.Run("rejects invalid environment", testLoggingInvalidEnvironment) + t.Run("flushes file sink during shutdown", testLoggingFileSinkShutdown) +} - t.Run("rejects invalid environment", func(t *testing.T) { - command := exec.Command(os.Args[0], "-test.run=TestBindingLoggingEnvironment") - command.Env = loggingTestEnvironment( - loggingHelperEnvironment+"=1", - "NEMO_RELAY_LOG=", - ) - output, err := command.CombinedOutput() - if err == nil { - t.Fatalf("binding initialization unexpectedly succeeded:\n%s", output) - } - if !strings.Contains(string(output), "NEMO_RELAY_LOG must not be empty") { - t.Fatalf("logging initialization error missing from output:\n%s", output) - } - }) +func testLoggingInitialization(t *testing.T) { + command := exec.Command(os.Args[0], loggingEnvironmentTestRunArg) + command.Env = loggingTestEnvironment( + loggingHelperEnvironment+"=shutdown", + "NEMO_RELAY_LOG=info", + "NEMO_RELAY_LOG_STDERR_FORMAT=jsonl", + ) + output, err := command.CombinedOutput() + if err != nil { + t.Fatalf("binding import failed: %v\n%s", err, output) + } + if !strings.Contains(string(output), `"event":"logging_initialized"`) { + t.Fatalf("logging initialization event missing from output:\n%s", output) + } +} - t.Run("flushes file sink during shutdown", func(t *testing.T) { - directory := t.TempDir() - configPath := filepath.Join(directory, "logging.toml") - logPath := filepath.Join(directory, "operational.jsonl") - config := `[logging] +func testLoggingInvalidEnvironment(t *testing.T) { + command := exec.Command(os.Args[0], loggingEnvironmentTestRunArg) + command.Env = loggingTestEnvironment( + loggingHelperEnvironment+"=1", + "NEMO_RELAY_LOG=", + ) + output, err := command.CombinedOutput() + if err == nil { + t.Fatalf("binding initialization unexpectedly succeeded:\n%s", output) + } + if !strings.Contains(string(output), "NEMO_RELAY_LOG must not be empty") { + t.Fatalf("logging initialization error missing from output:\n%s", output) + } +} + +func testLoggingFileSinkShutdown(t *testing.T) { + directory := t.TempDir() + configPath := filepath.Join(directory, "logging.toml") + logPath := filepath.Join(directory, "operational.jsonl") + config := `[logging] level = "info" stderr_format = "human" flush_interval_millis = 0 @@ -87,25 +94,24 @@ level = "info" format = "jsonl" queue_capacity = 16 ` - if err := os.WriteFile(configPath, []byte(config), 0o600); err != nil { - t.Fatalf("write logging config: %v", err) - } + if err := os.WriteFile(configPath, []byte(config), 0o600); err != nil { + t.Fatalf("write logging config: %v", err) + } - command := exec.Command(os.Args[0], "-test.run=TestBindingLoggingEnvironment") - command.Env = loggingTestEnvironment( - loggingHelperEnvironment+"=shutdown", - "NEMO_RELAY_LOG_CONFIG_PATH="+configPath, - ) - output, err := command.CombinedOutput() - if err != nil { - t.Fatalf("binding logging shutdown failed: %v\n%s", err, output) - } - contents, err := os.ReadFile(logPath) - if err != nil { - t.Fatalf("read operational log: %v", err) - } - if !strings.Contains(string(contents), `"event":"logging_shutdown_started"`) { - t.Fatalf("logging shutdown event missing from file:\n%s", contents) - } - }) + command := exec.Command(os.Args[0], loggingEnvironmentTestRunArg) + command.Env = loggingTestEnvironment( + loggingHelperEnvironment+"=shutdown", + "NEMO_RELAY_LOG_CONFIG_PATH="+configPath, + ) + output, err := command.CombinedOutput() + if err != nil { + t.Fatalf("binding logging shutdown failed: %v\n%s", err, output) + } + contents, err := os.ReadFile(logPath) + if err != nil { + t.Fatalf("read operational log: %v", err) + } + if !strings.Contains(string(contents), `"event":"logging_shutdown_started"`) { + t.Fatalf("logging shutdown event missing from file:\n%s", contents) + } } diff --git a/go/nemo_relay/plugin_activation_test.go b/go/nemo_relay/plugin_activation_test.go index fe80328ba..61c70d0c1 100644 --- a/go/nemo_relay/plugin_activation_test.go +++ b/go/nemo_relay/plugin_activation_test.go @@ -19,6 +19,15 @@ import ( "unsafe" ) +const ( + deferredCloseErrorFmt = "deferred Close() error = %v" + goNativeToolName = "go-native-tool" + toolInterceptErrorFmt = "ToolRequestIntercepts() error = %v" + invalidToolArgsFmt = "transformed tool args are invalid JSON: %v" + pluginsTOMLName = "plugins.toml" + staticFixtureKind = "go.fixture.static_base" +) + const ( pluginFixtureManifest = "/tmp/relay-plugin.toml" initializePluginsErrorFmt = "InitializeWithDynamicPlugins() error = %v" @@ -593,7 +602,7 @@ func TestInitializeWithDynamicPluginsLoadsNativePluginThroughCgo(t *testing.T) { } defer func() { if err := activation.Close(); err != nil { - t.Errorf("deferred Close() error = %v", err) + t.Errorf(deferredCloseErrorFmt, err) } }() if len(report.Diagnostics) != 1 { @@ -634,7 +643,7 @@ func TestInitializeWithDynamicPluginsIgnoresProjectPluginConfig(t *testing.T) { } defer func() { if err := activation.Close(); err != nil { - t.Errorf("deferred Close() error = %v", err) + t.Errorf(deferredCloseErrorFmt, err) } }() if len(report.Diagnostics) != 0 { @@ -644,13 +653,13 @@ func TestInitializeWithDynamicPluginsIgnoresProjectPluginConfig(t *testing.T) { t.Fatalf("static registrations = %d, want 0", staticRegistrations.Load()) } - transformed, err := ToolRequestIntercepts("go-native-tool", json.RawMessage(`{"input":true}`)) + transformed, err := ToolRequestIntercepts(goNativeToolName, json.RawMessage(`{"input":true}`)) if err != nil { - t.Fatalf("ToolRequestIntercepts() error = %v", err) + t.Fatalf(toolInterceptErrorFmt, err) } var transformedObject map[string]any if err := json.Unmarshal(transformed, &transformedObject); err != nil { - t.Fatalf("transformed tool args are invalid JSON: %v", err) + t.Fatalf(invalidToolArgsFmt, err) } if transformedObject["native_plugin"] != true || transformedObject["go_static_base"] != nil { t.Fatalf("transformed tool args = %s, want only native-plugin marker", transformed) @@ -663,7 +672,7 @@ func TestInitializeWithDynamicPluginsIgnoresProjectPluginConfig(t *testing.T) { func configureNativePluginUserConfig(t *testing.T) string { t.Helper() projectDir := t.TempDir() - legacyPluginsTOML := filepath.Join(projectDir, ".nemo-relay", "plugins.toml") + legacyPluginsTOML := filepath.Join(projectDir, ".nemo-relay", pluginsTOMLName) if err := os.MkdirAll(filepath.Dir(legacyPluginsTOML), 0o700); err != nil { t.Fatalf("MkdirAll(legacy project config) error = %v", err) } @@ -675,8 +684,8 @@ func configureNativePluginUserConfig(t *testing.T) string { if err := os.MkdirAll(userConfigDir, 0o700); err != nil { t.Fatalf("MkdirAll(user config) error = %v", err) } - pluginsTOML := filepath.Join(userConfigDir, "plugins.toml") - const staticKind = "go.fixture.static_base" + pluginsTOML := filepath.Join(userConfigDir, pluginsTOMLName) + const staticKind = staticFixtureKind fileConfig := fmt.Sprintf(`version = 1 [[components]] @@ -708,11 +717,11 @@ source = "user-file" func configureNativePluginProjectConfig(t *testing.T) { t.Helper() projectDir := t.TempDir() - projectPluginsTOML := filepath.Join(projectDir, ".nemo-relay", "plugins.toml") + projectPluginsTOML := filepath.Join(projectDir, ".nemo-relay", pluginsTOMLName) if err := os.MkdirAll(filepath.Dir(projectPluginsTOML), 0o700); err != nil { t.Fatalf("MkdirAll(project config) error = %v", err) } - const staticKind = "go.fixture.static_base" + const staticKind = staticFixtureKind projectConfig := fmt.Sprintf(`version = 1 [[components]] @@ -746,7 +755,7 @@ source = "project-file" func registerStaticFixturePlugin(t *testing.T) (*atomic.Int32, *atomic.Int32) { t.Helper() - const staticKind = "go.fixture.static_base" + const staticKind = staticFixtureKind staticRegistrations := &atomic.Int32{} staticCallbacks := &atomic.Int32{} if err := RegisterPlugin(staticKind, PluginFuncs{ @@ -792,13 +801,13 @@ func assertNativePluginInterception(t *testing.T, pluginsTOML string, staticCall t.Fatalf("mutate plugins.toml error = %v", err) } - transformed, err := ToolRequestIntercepts("go-native-tool", json.RawMessage(`{"input":true}`)) + transformed, err := ToolRequestIntercepts(goNativeToolName, json.RawMessage(`{"input":true}`)) if err != nil { - t.Fatalf("ToolRequestIntercepts() error = %v", err) + t.Fatalf(toolInterceptErrorFmt, err) } var transformedObject map[string]any if err := json.Unmarshal(transformed, &transformedObject); err != nil { - t.Fatalf("transformed tool args are invalid JSON: %v", err) + t.Fatalf(invalidToolArgsFmt, err) } if transformedObject["native_plugin"] != true { t.Fatalf("transformed tool args = %s, want native_plugin marker", transformed) @@ -819,7 +828,7 @@ func assertNativePluginCleanup(t *testing.T, activation *PluginActivation, plugi if err := activation.Close(); err != nil { t.Fatalf(closeErrorFmt, err) } - afterClose, err := ToolRequestIntercepts("go-native-tool", json.RawMessage(`{"input":true}`)) + afterClose, err := ToolRequestIntercepts(goNativeToolName, json.RawMessage(`{"input":true}`)) if err != nil { t.Fatalf("ToolRequestIntercepts() after Close error = %v", err) } @@ -879,7 +888,7 @@ func testInitializeWithDynamicWorkerPlugin(t *testing.T, manifest string) { } defer func() { if err := activation.Close(); err != nil { - t.Errorf("deferred Close() error = %v", err) + t.Errorf(deferredCloseErrorFmt, err) } }() if len(report.Diagnostics) != 0 { @@ -888,11 +897,11 @@ func testInitializeWithDynamicWorkerPlugin(t *testing.T, manifest string) { transformed, err := ToolRequestIntercepts("go-worker-tool", json.RawMessage(`{"input":true}`)) if err != nil { - t.Fatalf("ToolRequestIntercepts() error = %v", err) + t.Fatalf(toolInterceptErrorFmt, err) } var transformedObject map[string]any if err := json.Unmarshal(transformed, &transformedObject); err != nil { - t.Fatalf("transformed tool args are invalid JSON: %v", err) + t.Fatalf(invalidToolArgsFmt, err) } if transformedObject["worker_plugin"] != true { t.Fatalf("transformed tool args = %s, want worker_plugin marker", transformed) diff --git a/scripts/latency_benchmark/src/report/report.js b/scripts/latency_benchmark/src/report/report.js index 38e09c203..67631541f 100644 --- a/scripts/latency_benchmark/src/report/report.js +++ b/scripts/latency_benchmark/src/report/report.js @@ -74,7 +74,12 @@ function formatMs(value) { return "—"; } const magnitude = Math.abs(value); - const digits = magnitude >= 100 ? 1 : magnitude >= 10 ? 2 : 3; + let digits = 3; + if (magnitude >= 100) { + digits = 1; + } else if (magnitude >= 10) { + digits = 2; + } return `${value.toFixed(digits)} ms`; } @@ -237,7 +242,7 @@ function drawLineChart(svg, series, payloads, includeZero) { class: "line-series", points, stroke: color, - "stroke-dasharray": seriesIndex === 1 ? "9 5" : seriesIndex === 2 ? "3 5" : "none", + "stroke-dasharray": strokeDasharray(seriesIndex), }); item.values.forEach((point, index) => { const circle = addSvg(svg, "circle", { @@ -252,6 +257,16 @@ function drawLineChart(svg, series, payloads, includeZero) { }); } +function strokeDasharray(seriesIndex) { + if (seriesIndex === 1) { + return "9 5"; + } + if (seriesIndex === 2) { + return "3 5"; + } + return "none"; +} + function renderLegend(series) { const legend = byId("gateway-legend"); legend.replaceChildren(); @@ -321,11 +336,12 @@ function renderGateway() { renderLegend(series); const metricLabel = metric === "first_content" ? "time to first content" : "total response time"; - const viewLabel = view === "absolute" - ? "absolute latency" - : view === "minimal" - ? "variant overhead relative to minimal Relay" - : "Relay overhead relative to direct provider calls"; + let viewLabel = "Relay overhead relative to direct provider calls"; + if (view === "absolute") { + viewLabel = "absolute latency"; + } else if (view === "minimal") { + viewLabel = "variant overhead relative to minimal Relay"; + } byId("gateway-chart-description").textContent = `${statistic.replace("_ms", "")} ${metricLabel}; ${viewLabel}. ` + `Provider ${byId("gateway-provider").value}, ${byId("gateway-mode").value}, ` + diff --git a/scripts/latency_benchmark/src/report/template.html b/scripts/latency_benchmark/src/report/template.html index 2ba6e4b71..86495b5bb 100644 --- a/scripts/latency_benchmark/src/report/template.html +++ b/scripts/latency_benchmark/src/report/template.html @@ -104,7 +104,6 @@

Gateway Latency by Payload Size

@@ -136,7 +135,6 @@

Hook Absolute p50

@@ -168,7 +166,6 @@

Startup Absolute p50

From 82463082197cb46cec6e0189d25413e29624aa21 Mon Sep 17 00:00:00 2001 From: Will Killian Date: Fri, 7 Aug 2026 16:47:23 -0400 Subject: [PATCH 2/4] chore: reduce Gemini tool group complexity Signed-off-by: Will Killian --- .../core/src/codec/gemini_generate_content.rs | 61 +++++++++++-------- 1 file changed, 37 insertions(+), 24 deletions(-) diff --git a/crates/core/src/codec/gemini_generate_content.rs b/crates/core/src/codec/gemini_generate_content.rs index dd3b06a47..0d8c1fec0 100644 --- a/crates/core/src/codec/gemini_generate_content.rs +++ b/crates/core/src/codec/gemini_generate_content.rs @@ -2608,32 +2608,16 @@ fn rebuild_gemini_tool_groups( let mut fn_group_placed = false; let mut native_used = vec![false; native_groups.len()]; for orig_group in orig_tools { - if orig_group.get("functionDeclarations").is_some() { - if !fn_group_placed { - let mut group = serde_json::Map::new(); - if !fn_declarations.is_empty() { - group.insert( - "functionDeclarations".into(), - Json::Array(fn_declarations.to_vec()), - ); - } - merge_gemini_native_sibling_group( - orig_group, - native_groups, - &mut native_used, - &mut group, - ); - if !group.is_empty() { - new_groups.push(Json::Object(group)); - } - fn_group_placed = true; - } - } else if let Some(native_group) = take_matching_native_group( + let (placed, group) = rebuild_gemini_tool_group( + orig_group, + fn_declarations, native_groups, &mut native_used, - &gemini_native_tool_keys(orig_group), - ) { - new_groups.push(native_group); + fn_group_placed, + ); + fn_group_placed |= placed; + if let Some(group) = group { + new_groups.push(group); } } if !fn_group_placed && !fn_declarations.is_empty() { @@ -2649,6 +2633,35 @@ fn rebuild_gemini_tool_groups( new_groups } +fn rebuild_gemini_tool_group( + orig_group: &Json, + fn_declarations: &[Json], + native_groups: &[Json], + native_used: &mut [bool], + fn_group_placed: bool, +) -> (bool, Option) { + if orig_group.get("functionDeclarations").is_some() { + if fn_group_placed { + return (false, None); + } + let mut group = serde_json::Map::new(); + if !fn_declarations.is_empty() { + group.insert( + "functionDeclarations".into(), + Json::Array(fn_declarations.to_vec()), + ); + } + merge_gemini_native_sibling_group(orig_group, native_groups, native_used, &mut group); + return (true, (!group.is_empty()).then_some(Json::Object(group))); + } + let group = take_matching_native_group( + native_groups, + native_used, + &gemini_native_tool_keys(orig_group), + ); + (false, group) +} + fn merge_gemini_native_sibling_group( orig_group: &Json, native_groups: &[Json], From 97ea3d855ae45718177b90572cd7d48ace6877e0 Mon Sep 17 00:00:00 2001 From: Will Killian Date: Fri, 7 Aug 2026 16:57:33 -0400 Subject: [PATCH 3/4] fix: release pending Node subscriber callbacks Signed-off-by: Will Killian --- crates/node/src/callable.rs | 53 +++++++++++++++++++++++++++++++------ 1 file changed, 45 insertions(+), 8 deletions(-) diff --git a/crates/node/src/callable.rs b/crates/node/src/callable.rs index 422b7862f..a9181b970 100644 --- a/crates/node/src/callable.rs +++ b/crates/node/src/callable.rs @@ -1271,23 +1271,35 @@ fn create_js_event_subscriber_function( let completed = Arc::new(AtomicBool::new(false)); let completed_callback = Arc::clone(&completed); let callback_name = name.clone(); - let complete = create_js_subscriber_completion_callback( - &ctx.env, - callback_name, + let complete = complete_subscriber_callback_on_error( callback_id, - completed_callback, + create_js_subscriber_completion_callback( + &ctx.env, + callback_name, + callback_id, + completed_callback, + ), )?; - let event = unsafe { - JsUnknown::from_raw_unchecked( + let event = complete_subscriber_callback_on_error(callback_id, unsafe { + Ok(JsUnknown::from_raw_unchecked( ctx.env.raw(), Json::to_napi_value(ctx.env.raw(), event)?, - ) - }; + )) + })?; let complete = unsafe { JsUnknown::from_raw_unchecked(ctx.env.raw(), complete.raw()) }; Ok(vec![event, complete]) }) } +fn complete_subscriber_callback_on_error( + callback_id: u64, + result: napi::Result, +) -> napi::Result { + result.inspect_err(|_| { + complete_js_subscriber_callback(callback_id); + }) +} + fn create_js_subscriber_completion_callback( env: &Env, callback_name: String, @@ -1589,3 +1601,28 @@ pub fn wrap_js_llm_stream_exec_intercept_fn( }, ) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn subscriber_callback_completion_is_idempotent() { + let callback_id = reserve_js_subscriber_callback(); + complete_js_subscriber_callback(callback_id); + complete_js_subscriber_callback(callback_id); + + assert!(flush_js_subscriber_callbacks().is_ok()); + } + + #[test] + fn subscriber_callback_preparation_errors_release_pending_ids() { + for message in ["completion callback", "event conversion"] { + let callback_id = reserve_js_subscriber_callback(); + let result: napi::Result<()> = Err(napi::Error::from_reason(message)); + + assert!(complete_subscriber_callback_on_error(callback_id, result).is_err()); + assert!(flush_js_subscriber_callbacks().is_ok()); + } + } +} From 89fecac63745c6f33abe34b78f05254b30843477 Mon Sep 17 00:00:00 2001 From: Will Killian Date: Fri, 7 Aug 2026 17:17:49 -0400 Subject: [PATCH 4/4] fix: address CodeRabbit review feedback Signed-off-by: Will Killian --- crates/cli/src/diagnostics/mod.rs | 4 +- .../cli/tests/coverage/shared/doctor_tests.rs | 19 ++++--- .../src/observability/plugin_component.rs | 54 +++++++++++++------ crates/node/src/callable.rs | 25 --------- crates/node/tests/scope_tests.mjs | 12 +++++ 5 files changed, 65 insertions(+), 49 deletions(-) diff --git a/crates/cli/src/diagnostics/mod.rs b/crates/cli/src/diagnostics/mod.rs index e8de8e53b..943091456 100644 --- a/crates/cli/src/diagnostics/mod.rs +++ b/crates/cli/src/diagnostics/mod.rs @@ -921,9 +921,7 @@ async fn observability_http_exporter_check( } else { probe_http_endpoint(label, url, probe_mode).await }; - if !probe_mode.is_offline() { - check.details = format!("endpoints[{index}] ({endpoint_type}): {}", check.details); - } + check.details = format!("endpoints[{index}] ({endpoint_type}): {}", check.details); check } diff --git a/crates/cli/tests/coverage/shared/doctor_tests.rs b/crates/cli/tests/coverage/shared/doctor_tests.rs index 96f161d8f..4dd8fe117 100644 --- a/crates/cli/tests/coverage/shared/doctor_tests.rs +++ b/crates/cli/tests/coverage/shared/doctor_tests.rs @@ -1201,11 +1201,16 @@ async fn opentelemetry_doctor_skips_live_network_probes_offline() { assert_eq!(checks.len(), 2); assert!(checks.iter().all(|check| check.status == Status::Info)); - assert!(checks.iter().all(|check| { - check + assert!( + checks[0] .details - .contains("live network probe skipped (--offline)") - })); + .contains("endpoints[0] (gen_ai): live network probe skipped (--offline)") + ); + assert!( + checks[1] + .details + .contains("endpoints[1] (openinference): live network probe skipped (--offline)") + ); } #[tokio::test] @@ -1229,8 +1234,10 @@ async fn opentelemetry_doctor_offline_still_rejects_malformed_endpoints() { .await; assert_eq!(checks.len(), 2); - assert!(checks[0].details.contains("invalid gRPC endpoint")); - assert!(checks[1].details.contains("invalid OTLP HTTP endpoint")); + assert!(checks[0].details.starts_with("endpoints[0] (gen_ai): ")); + assert!(checks[0].details.contains("gRPC endpoint")); + assert!(checks[1].details.starts_with("endpoints[1] (full): ")); + assert!(checks[1].details.contains("OTLP HTTP endpoint")); assert!(checks.iter().all(|check| check.status == Status::Fail)); } diff --git a/crates/core/src/observability/plugin_component.rs b/crates/core/src/observability/plugin_component.rs index b1c48bc4b..5bb451590 100644 --- a/crates/core/src/observability/plugin_component.rs +++ b/crates/core/src/observability/plugin_component.rs @@ -953,13 +953,16 @@ fn atif_shutdown_cleanup( manager: Arc>, shutdown_storage: AtifStorageList, ) -> PluginRegistrationCleanupOutcome { - let work = match flush_atif_shutdown_work(&manager) { + let (work, deregistration_error) = match flush_atif_shutdown_work(&manager) { Ok(work) => work, Err(error) => return PluginRegistrationCleanupOutcome::NotRemoved(error), }; if let Err(error) = write_atif_shutdown_exports(&manager, &shutdown_storage, work.exports) { return PluginRegistrationCleanupOutcome::RemovedWithError(error); } + if let Some(error) = deregistration_error { + return PluginRegistrationCleanupOutcome::NotRemoved(error); + } match manager.lock() { Ok(guard) => match guard.last_error_result() { Ok(()) => PluginRegistrationCleanupOutcome::Removed, @@ -973,17 +976,24 @@ fn atif_shutdown_cleanup( } } -fn flush_atif_shutdown_work(manager: &Arc>) -> PluginResult { +fn flush_atif_shutdown_work( + manager: &Arc>, +) -> PluginResult<(AtifFlushWork, Option)> { let work = { let mut guard = manager.lock().map_err(|err| { PluginError::Internal(format!("ATIF dispatcher lock poisoned: {err}")) })?; guard.flush_open_agents() }; + let mut deregistration_error = None; for (scope_uuid, name) in &work.scope_subscribers { - deregister_atif_shutdown_subscriber(scope_uuid, name)?; + if let Err(error) = deregister_atif_shutdown_subscriber(scope_uuid, name) + && deregistration_error.is_none() + { + deregistration_error = Some(error); + } } - Ok(work) + Ok((work, deregistration_error)) } fn write_atif_shutdown_exports( @@ -991,19 +1001,33 @@ fn write_atif_shutdown_exports( shutdown_storage: &AtifStorageList, exports: Vec, ) -> PluginResult<()> { + let mut first_error = None; for export in exports { - let write = prepare_atif_shutdown_file(&export, Arc::clone(manager)) - .map_err(observability_registration_error)?; - let targets = manager - .lock() - .map_err(|err| PluginError::Internal(format!("ATIF dispatcher lock poisoned: {err}")))? - .sink_targets(); - let results = write_atif(&write, shutdown_storage.as_slice(), &targets); - let mut guard = manager.lock().map_err(|err| { - PluginError::Internal(format!("ATIF dispatcher lock poisoned: {err}")) - })?; - let _ = guard.complete_scope_write(write.agent_uuid, results); + if let Err(error) = write_atif_shutdown_export(manager, shutdown_storage, &export) + && first_error.is_none() + { + first_error = Some(error); + } } + first_error.map_or(Ok(()), Err) +} + +fn write_atif_shutdown_export( + manager: &Arc>, + shutdown_storage: &AtifStorageList, + export: &PendingAtifExport, +) -> PluginResult<()> { + let write = prepare_atif_shutdown_file(export, Arc::clone(manager)) + .map_err(observability_registration_error)?; + let targets = manager + .lock() + .map_err(|err| PluginError::Internal(format!("ATIF dispatcher lock poisoned: {err}")))? + .sink_targets(); + let results = write_atif(&write, shutdown_storage.as_slice(), &targets); + let mut guard = manager + .lock() + .map_err(|err| PluginError::Internal(format!("ATIF dispatcher lock poisoned: {err}")))?; + let _ = guard.complete_scope_write(write.agent_uuid, results); Ok(()) } diff --git a/crates/node/src/callable.rs b/crates/node/src/callable.rs index a9181b970..5b0521d5c 100644 --- a/crates/node/src/callable.rs +++ b/crates/node/src/callable.rs @@ -1601,28 +1601,3 @@ pub fn wrap_js_llm_stream_exec_intercept_fn( }, ) } - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn subscriber_callback_completion_is_idempotent() { - let callback_id = reserve_js_subscriber_callback(); - complete_js_subscriber_callback(callback_id); - complete_js_subscriber_callback(callback_id); - - assert!(flush_js_subscriber_callbacks().is_ok()); - } - - #[test] - fn subscriber_callback_preparation_errors_release_pending_ids() { - for message in ["completion callback", "event conversion"] { - let callback_id = reserve_js_subscriber_callback(); - let result: napi::Result<()> = Err(napi::Error::from_reason(message)); - - assert!(complete_subscriber_callback_on_error(callback_id, result).is_err()); - assert!(flush_js_subscriber_callbacks().is_ok()); - } - } -} diff --git a/crates/node/tests/scope_tests.mjs b/crates/node/tests/scope_tests.mjs index e2f849301..5c0ce470b 100644 --- a/crates/node/tests/scope_tests.mjs +++ b/crates/node/tests/scope_tests.mjs @@ -482,6 +482,18 @@ describe('Subscribers', () => { } }); + it('flushSubscribers settles after a subscriber callback failure', async () => { + registerSubscriber('node_flush_js_failure', () => { + throw new Error('flush failure'); + }); + try { + event('node_flush_js_failure_mark', null, null, null); + await flushSubscribers(); + } finally { + deregisterSubscriber('node_flush_js_failure'); + } + }); + it('isolates a synchronous global subscriber throw', () => { runSubscriberFailureChild({ callback: "throw new Error('sync subscriber boom');",