diff --git a/CLAUDE.md b/CLAUDE.md index ea3ec78..73da5b5 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -57,7 +57,7 @@ The loop: stream assistant response → extract tool calls → execute tools (pa - **`Message`** — enum: `User`, `Assistant`, `ToolResult` — each variant carries its own fields - **`AgentMessage`** — `Llm(Message)` | `Extension(ExtensionMessage)` — extension messages (`role`, `kind`, `data`) don't enter LLM context - **`AgentEvent`** — full event stream emitted to callers: `AgentStart`, `TurnStart`, `MessageStart/Update/End`, `ToolExecutionStart/Update/End`, `ProgressMessage`, `InputRejected`, `TurnEnd`, `AgentEnd` -- **`StopReason`** — `Stop`, `Length`, `ToolUse`, `Error`, `Aborted` +- **`StopReason`** — `Stop`, `Length`, `ToolUse`, `Error`, `Aborted`, `Refusal` ### Context Management (`context.rs`) diff --git a/docs/concepts/messages-events.md b/docs/concepts/messages-events.md index 335f1fc..60690f0 100644 --- a/docs/concepts/messages-events.md +++ b/docs/concepts/messages-events.md @@ -74,12 +74,19 @@ pub enum Content { Text { text: String }, Image { data: String, mime_type: String }, Thinking { thinking: String, signature: Option }, - ToolCall { id: String, name: String, arguments: serde_json::Value }, + ToolCall { + id: String, + name: String, + arguments: serde_json::Value, + provider_metadata: Option, // e.g. Gemini thought signatures + }, } ``` An assistant message can contain multiple content blocks — e.g., thinking + text + tool calls. +`Content` is `#[non_exhaustive]` (match with a wildcard arm), and the `ToolCall` and `Thinking` variants are separately `#[non_exhaustive]` — construct them via `Content::tool_call()` / `tool_call_with_metadata()` / `thinking()` / `thinking_signed()`. `Message::Assistant` is likewise `#[non_exhaustive]`; custom providers construct it via `Message::assistant()`. + ## StopReason ```rust @@ -89,6 +96,7 @@ pub enum StopReason { ToolUse, // Wants to call tools Error, // Provider error Aborted, // Cancelled by user + Refusal, // Declined by the provider's safety system } ``` diff --git a/src/provider/google.rs b/src/provider/google.rs index b07da98..33e57dd 100644 --- a/src/provider/google.rs +++ b/src/provider/google.rs @@ -62,6 +62,7 @@ impl StreamProvider for GoogleProvider { let mut content: Vec = Vec::new(); let mut usage = Usage::default(); let mut stop_reason = StopReason::Stop; + let mut error_message: Option = None; let _ = tx.send(StreamEvent::Start); @@ -78,31 +79,31 @@ impl StreamProvider for GoogleProvider { match chunk { None => break, Some(Err(e)) => { - warn!("Google stream error: {}", e); - break; + // Match the other providers: a transport failure is an + // error (and retryable), not a silently truncated turn. + let provider_err = ProviderError::Network(e.to_string()); + warn!("Google stream error: {}", provider_err); + return Err(provider_err); } Some(Ok(bytes)) => { buffer.push_str(&String::from_utf8_lossy(&bytes)); // Process complete SSE events (handle both \n\n and \r\n\r\n) - while let Some(pos) = buffer.find("\n\n").or_else(|| buffer.find("\r\n\r\n")) { - let sep_len = if buffer[pos..].starts_with("\r\n\r\n") { 4 } else { 2 }; - let event_str = buffer[..pos].to_string(); - buffer = buffer[pos + sep_len..].to_string(); - - // Parse SSE data line (strip trailing \r) - let data = event_str - .lines() - .map(|l| l.trim_end_matches('\r')) - .find(|l| l.starts_with("data: ")) - .map(|l| &l[6..]) - .unwrap_or(""); - + while let Some(data) = next_sse_data(&mut buffer) { if data.is_empty() { continue; } - let chunk: GoogleChunk = match serde_json::from_str(data) { + // Google reports mid-stream failures as + // {"error": {...}} payloads, which would otherwise + // deserialize into an empty chunk and vanish. + if is_error_payload(&data) { + let provider_err = classify_sse_error_event(&data); + warn!("Google in-stream error: {}", provider_err); + return Err(provider_err); + } + + let chunk: GoogleChunk = match serde_json::from_str(&data) { Ok(c) => c, Err(e) => { warn!("Failed to parse Google chunk: {}", e); @@ -114,11 +115,7 @@ impl StreamProvider for GoogleProvider { for candidate in &chunk.candidates.unwrap_or_default() { if let Some(c) = &candidate.content { for part in &c.parts { - if let Some(text) = &part.text { - // Skip empty text parts (sent during thinking) - if text.is_empty() { - continue; - } + if let Some(text) = part_text(part) { let text_idx = content.iter().position(|c| matches!(c, Content::Text { .. })); let idx = match text_idx { Some(i) => i, @@ -132,7 +129,7 @@ impl StreamProvider for GoogleProvider { } let _ = tx.send(StreamEvent::TextDelta { content_index: idx, - delta: text.clone(), + delta: text.to_string(), }); } if let Some(fc) = &part.function_call { @@ -165,6 +162,18 @@ impl StreamProvider for GoogleProvider { stop_reason = match reason.as_str() { "STOP" => StopReason::Stop, "MAX_TOKENS" | "RECITATION" => StopReason::Length, + "SAFETY" | "PROHIBITED_CONTENT" | "BLOCKLIST" + | "SPII" => { + warn!( + "Gemini blocked the response (finishReason={})", + reason + ); + error_message = Some(format!( + "Response blocked by Gemini safety filters (finishReason: {})", + reason + )); + StopReason::Refusal + } _ => StopReason::Stop, }; } @@ -192,7 +201,7 @@ impl StreamProvider for GoogleProvider { provider: model_config.provider.clone(), usage, timestamp: now_ms(), - error_message: None, + error_message, }; let _ = tx.send(StreamEvent::Done { @@ -202,6 +211,45 @@ impl StreamProvider for GoogleProvider { } } +/// Pop the next complete SSE event from `buffer` and return its `data:` +/// payload (empty string when the event carries no data line). Handles both +/// `\n\n` and `\r\n\r\n` event separators, splitting at whichever occurs +/// first. Returns `None` until a complete event is buffered. Only the first +/// `data:` line of an event is returned. +fn next_sse_data(buffer: &mut String) -> Option { + let lf = buffer.find("\n\n"); + let crlf = buffer.find("\r\n\r\n"); + let (pos, sep_len) = match (lf, crlf) { + (Some(l), Some(c)) if c < l => (c, 4), + (Some(l), _) => (l, 2), + (None, Some(c)) => (c, 4), + (None, None) => return None, + }; + let event_str = buffer[..pos].to_string(); + *buffer = buffer[pos + sep_len..].to_string(); + let data = event_str + .lines() + .map(|l| l.trim_end_matches('\r')) + .find(|l| l.starts_with("data: ")) + .map(|l| l[6..].to_string()) + .unwrap_or_default(); + Some(data) +} + +/// Whether an SSE data payload is a Google error envelope +/// (`{"error": {...}}`) rather than a content chunk. +fn is_error_payload(data: &str) -> bool { + serde_json::from_str::(data) + .map(|v| v.get("error").is_some()) + .unwrap_or(false) +} + +/// Non-empty text of a part. Gemini streams empty text parts while thinking; +/// those must be skipped. +fn part_text(part: &GooglePart) -> Option<&str> { + part.text.as_deref().filter(|t| !t.is_empty()) +} + fn build_request_body(config: &StreamConfig) -> serde_json::Value { let mut contents: Vec = Vec::new(); @@ -475,41 +523,75 @@ mod tests { #[test] fn test_parse_chunk_with_empty_text() { - // Gemini sends empty text parts during thinking -- these should parse fine - let data = r#"{"candidates": [{"content": {"parts": [{"text": ""}], "role": "model"}, "index": 0}]}"#; + // Gemini sends empty text parts during thinking -- part_text (used by + // the streaming loop) must skip them and keep non-empty ones. + let data = r#"{"candidates": [{"content": {"parts": [{"text": ""}, {"text": "Hello"}], "role": "model"}, "index": 0}]}"#; let chunk: GoogleChunk = serde_json::from_str(data).unwrap(); let candidates = chunk.candidates.unwrap(); let parts = &candidates[0].content.as_ref().unwrap().parts; - assert_eq!(parts[0].text.as_deref(), Some("")); + assert_eq!(part_text(&parts[0]), None, "empty text parts are skipped"); + assert_eq!(part_text(&parts[1]), Some("Hello")); } #[test] fn test_parse_chunk_with_crlf_sse() { - // Simulate \r\n line endings from Google's API - let raw = "data: {\"candidates\": [{\"content\": {\"parts\": [{\"text\": \"Blue\"}], \"role\": \"model\"}, \"finishReason\": \"STOP\", \"index\": 0}]}\r\n\r\n"; - - // Verify our separator detection works - let pos = raw.find("\n\n").or_else(|| raw.find("\r\n\r\n")); - assert!(pos.is_some(), "Should find separator"); - let sep_start = pos.unwrap(); - let is_crlf = raw[sep_start..].starts_with("\r\n\r\n"); - assert!(is_crlf, "Should detect \\r\\n\\r\\n separator"); - - let event_str = &raw[..sep_start]; - let data = event_str - .lines() - .map(|l| l.trim_end_matches('\r')) - .find(|l| l.starts_with("data: ")) - .map(|l| &l[6..]) - .unwrap(); + // Full pipeline: next_sse_data (the production splitter) on a CRLF + // stream, then chunk parsing. + let mut buf = "data: {\"candidates\": [{\"content\": {\"parts\": [{\"text\": \"Blue\"}], \"role\": \"model\"}, \"finishReason\": \"STOP\", \"index\": 0}]}\r\n\r\n".to_string(); - let chunk: GoogleChunk = serde_json::from_str(data).unwrap(); + let data = next_sse_data(&mut buf).expect("complete CRLF event"); + assert!(buf.is_empty(), "event consumed from buffer"); + assert_eq!(next_sse_data(&mut buf), None); + + let chunk: GoogleChunk = serde_json::from_str(&data).unwrap(); let candidates = chunk.candidates.unwrap(); let text = &candidates[0].content.as_ref().unwrap().parts[0].text; assert_eq!(text.as_deref(), Some("Blue")); } + #[test] + fn test_next_sse_data_partial_events_stay_buffered() { + let mut buf = "data: {\"a\":1}\n\ndata: partial".to_string(); + assert_eq!(next_sse_data(&mut buf).as_deref(), Some("{\"a\":1}")); + assert_eq!(next_sse_data(&mut buf), None, "incomplete event waits"); + assert_eq!(buf, "data: partial"); + buf.push_str("\n\n"); + assert_eq!(next_sse_data(&mut buf).as_deref(), Some("partial")); + } + + #[test] + fn test_next_sse_data_consumes_events_without_data_lines() { + // SSE comments/keepalives and leading separators must be CONSUMED + // (returning Some("")), never None — returning None would wedge the + // buffer and drop every subsequent event. + let mut buf = ": keepalive\n\n\n\ndata: x\n\n".to_string(); + assert_eq!(next_sse_data(&mut buf).as_deref(), Some("")); + assert_eq!(next_sse_data(&mut buf).as_deref(), Some("")); + assert_eq!(next_sse_data(&mut buf).as_deref(), Some("x")); + assert_eq!(next_sse_data(&mut buf), None); + } + + #[test] + fn test_is_error_payload() { + assert!(is_error_payload( + r#"{"error": {"code": 429, "status": "RESOURCE_EXHAUSTED"}}"# + )); + assert!(!is_error_payload(r#"{"candidates": []}"#)); + assert!(!is_error_payload("not json")); + } + + #[test] + fn test_next_sse_data_splits_earliest_separator_first() { + // A CRLF-separated event earlier in the buffer must split before a + // later LF separator (the old inline logic preferred the LF match and + // merged the two events). + let mut buf = "data: one\r\n\r\ndata: two\n\n".to_string(); + assert_eq!(next_sse_data(&mut buf).as_deref(), Some("one")); + assert_eq!(next_sse_data(&mut buf).as_deref(), Some("two")); + assert_eq!(next_sse_data(&mut buf), None); + } + #[test] fn test_thought_signature_round_trip() { let content = vec![Content::ToolCall { diff --git a/src/provider/model.rs b/src/provider/model.rs index 676cff1..523f9bf 100644 --- a/src/provider/model.rs +++ b/src/provider/model.rs @@ -275,7 +275,17 @@ impl OpenCodeGateway { } /// Full model configuration. Knows everything needed to make API calls. +/// +/// Marked `#[non_exhaustive]`: fields may be added in minor releases (e.g. +/// the `anthropic` compat flags, slated for 0.9.0). Construct via the +/// `ModelConfig::*` preset constructors — or [`ModelConfig::custom`] for +/// protocols without a preset — and mutate fields to customize. Note that +/// downstream struct literals and functional-record-update +/// (`ModelConfig { .. }`) no longer compile; field mutation is the supported +/// pattern. New fields must carry `#[serde(default)]` so previously +/// persisted configs keep deserializing. #[derive(Debug, Clone, Serialize, Deserialize)] +#[non_exhaustive] pub struct ModelConfig { /// Model identifier sent to the API (e.g. "gpt-4o", "claude-sonnet-4-20250514"). pub id: String, @@ -309,6 +319,35 @@ pub struct ModelConfig { } impl ModelConfig { + /// Create a config for any protocol without a dedicated preset + /// (Bedrock, Vertex, Azure, or future protocols). + /// + /// Since `ModelConfig` is `#[non_exhaustive]`, this is the construction + /// path when no `ModelConfig::*` preset fits. Defaults: 128K context, + /// 16K max output, no compat flags — mutate fields to adjust. + pub fn custom( + api: ApiProtocol, + provider: impl Into, + base_url: impl Into, + model_id: impl Into, + name: impl Into, + ) -> Self { + Self { + id: model_id.into(), + name: name.into(), + api, + provider: provider.into(), + base_url: base_url.into(), + reasoning: false, + context_window: 128_000, + max_tokens: 16_000, + cost: CostConfig::default(), + headers: HashMap::new(), + compat: None, + anthropic: None, + } + } + /// Create a new Anthropic model config. pub fn anthropic(id: impl Into, name: impl Into) -> Self { Self { diff --git a/src/types.rs b/src/types.rs index caaed95..1e0bc51 100644 --- a/src/types.rs +++ b/src/types.rs @@ -6,8 +6,20 @@ use std::sync::Arc; // Content types // --------------------------------------------------------------------------- +/// Content block of a message. +/// +/// Exhaustiveness policy (two separate levers): +/// - The **enum** is `#[non_exhaustive]`: new content kinds may be added in +/// minor releases, so downstream `match` arms need a wildcard. +/// - The `ToolCall` and `Thinking` **variants** are separately +/// `#[non_exhaustive]`: their fields grow with provider features (PR #32 +/// added `provider_metadata`), so downstream constructs them via the +/// `Content::tool_call*` / `Content::thinking*` constructors and uses `..` +/// in patterns. `Text` and `Image` stay literally constructible — they are +/// user-facing shapes that do not grow. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] #[serde(tag = "type")] +#[non_exhaustive] pub enum Content { #[serde(rename = "text")] Text { text: String }, @@ -18,12 +30,14 @@ pub enum Content { mime_type: String, }, #[serde(rename = "thinking")] + #[non_exhaustive] Thinking { thinking: String, #[serde(skip_serializing_if = "Option::is_none")] signature: Option, }, #[serde(rename = "toolCall")] + #[non_exhaustive] ToolCall { id: String, name: String, @@ -36,6 +50,58 @@ pub enum Content { }, } +impl Content { + /// Construct a tool-call content block. + /// + /// The `ToolCall` variant is `#[non_exhaustive]` so provider-specific + /// fields can be added without breaking downstream crates — use this + /// constructor instead of a struct literal. + pub fn tool_call( + id: impl Into, + name: impl Into, + arguments: serde_json::Value, + ) -> Self { + Self::ToolCall { + id: id.into(), + name: name.into(), + arguments, + provider_metadata: None, + } + } + + /// Construct a thinking content block without a signature. + pub fn thinking(text: impl Into) -> Self { + Self::Thinking { + thinking: text.into(), + signature: None, + } + } + + /// Construct a thinking content block with a provider signature. + pub fn thinking_signed(text: impl Into, signature: impl Into) -> Self { + Self::Thinking { + thinking: text.into(), + signature: Some(signature.into()), + } + } + + /// Construct a tool-call content block carrying provider metadata + /// (e.g. a Gemini thought signature). + pub fn tool_call_with_metadata( + id: impl Into, + name: impl Into, + arguments: serde_json::Value, + provider_metadata: serde_json::Value, + ) -> Self { + Self::ToolCall { + id: id.into(), + name: name.into(), + arguments, + provider_metadata: Some(provider_metadata), + } + } +} + // --------------------------------------------------------------------------- // Messages // --------------------------------------------------------------------------- @@ -49,6 +115,7 @@ pub enum Message { timestamp: u64, }, #[serde(rename = "assistant")] + #[non_exhaustive] Assistant { content: Vec, #[serde(rename = "stopReason")] @@ -81,6 +148,50 @@ impl Message { } } + /// Construct an assistant message. + /// + /// The `Assistant` variant is `#[non_exhaustive]` — its fields grow with + /// provider features (`error_message` was itself a later addition), so + /// custom `StreamProvider` implementations construct it here instead of + /// with a struct literal. `timestamp` is set to now and `error_message` + /// to `None`; use [`Message::with_error_message`] / + /// [`Message::with_timestamp`] to override. + pub fn assistant( + content: Vec, + stop_reason: StopReason, + model: impl Into, + provider: impl Into, + usage: Usage, + ) -> Self { + Self::Assistant { + content, + stop_reason, + model: model.into(), + provider: provider.into(), + usage, + timestamp: now_ms(), + error_message: None, + } + } + + /// Set the error message (no-op on non-assistant messages). + pub fn with_error_message(mut self, msg: impl Into) -> Self { + if let Self::Assistant { error_message, .. } = &mut self { + *error_message = Some(msg.into()); + } + self + } + + /// Override the timestamp (applies to all message kinds). + pub fn with_timestamp(mut self, ts: u64) -> Self { + match &mut self { + Self::User { timestamp, .. } + | Self::Assistant { timestamp, .. } + | Self::ToolResult { timestamp, .. } => *timestamp = ts, + } + self + } + pub fn role(&self) -> &str { match self { Self::User { .. } => "user", @@ -162,6 +273,12 @@ impl From for AgentMessage { // Stop reasons & usage // --------------------------------------------------------------------------- +/// Why the model stopped generating. +/// +/// Deliberately NOT `#[non_exhaustive]`: this is a control-flow enum, and a +/// new variant (like `Refusal`, added for the 0.9.0 breaking release) +/// should be a compile error for +/// downstream matches rather than silently falling into a wildcard arm. #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] #[serde(rename_all = "camelCase")] pub enum StopReason { diff --git a/tests/agent_loop_test.rs b/tests/agent_loop_test.rs index b2e0d86..518fa5a 100644 --- a/tests/agent_loop_test.rs +++ b/tests/agent_loop_test.rs @@ -714,17 +714,15 @@ impl StreamProvider for UsageProvider { ) -> Result { self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst); - let message = Message::Assistant { - content: vec![Content::Text { + let message = Message::assistant( + vec![Content::Text { text: "usage response".into(), }], - stop_reason: StopReason::Stop, - model: "usage-test".into(), - provider: "usage-test".into(), - usage: self.usage.clone(), - timestamp: now_ms(), - error_message: None, - }; + StopReason::Stop, + "usage-test", + "usage-test", + self.usage.clone(), + ); let _ = tx.send(StreamEvent::Start); let _ = tx.send(StreamEvent::Done { diff --git a/tests/agent_test.rs b/tests/agent_test.rs index e74ac9a..8fbd6d3 100644 --- a/tests/agent_test.rs +++ b/tests/agent_test.rs @@ -126,17 +126,15 @@ async fn test_agent_builder_pattern() { async fn test_with_messages_builder() { let saved = vec![ AgentMessage::Llm(Message::user("Hello")), - AgentMessage::Llm(Message::Assistant { - content: vec![Content::Text { + AgentMessage::Llm(Message::assistant( + vec![Content::Text { text: "Hi there!".into(), }], - stop_reason: StopReason::Stop, - model: "mock".into(), - provider: "mock".into(), - usage: Usage::default(), - timestamp: 0, - error_message: None, - }), + StopReason::Stop, + "mock", + "mock", + Usage::default(), + )), ]; let provider = MockProvider::text("ok"); @@ -305,15 +303,16 @@ async fn test_continue_loop_with_sender() { // First, add some messages to continue from (last must not be assistant) agent.append_message(AgentMessage::Llm(Message::user("Hello"))); - agent.append_message(AgentMessage::Llm(Message::Assistant { - content: vec![Content::Text { text: "Hi!".into() }], - stop_reason: StopReason::Error, - model: "mock".into(), - provider: "mock".into(), - usage: Usage::default(), - timestamp: 0, - error_message: Some("rate limited".into()), - })); + agent.append_message(AgentMessage::Llm( + Message::assistant( + vec![Content::Text { text: "Hi!".into() }], + StopReason::Error, + "mock", + "mock", + Usage::default(), + ) + .with_error_message("rate limited"), + )); agent.append_message(AgentMessage::Llm(Message::user("Please try again"))); let (tx, mut rx) = mpsc::unbounded_channel(); diff --git a/tests/google_stream_test.rs b/tests/google_stream_test.rs new file mode 100644 index 0000000..39affc9 --- /dev/null +++ b/tests/google_stream_test.rs @@ -0,0 +1,316 @@ +//! Behavioral tests for `GoogleProvider` against a local mock server. +//! +//! These give CI coverage for the PR #32 regression scenarios (Gemini +//! thought-signature round-trip and multi-turn function calling), previously +//! exercised only by the key-gated `integration_gemini.rs` live tests +//! (issue #33). + +use tokio::sync::mpsc; +use tokio_util::sync::CancellationToken; +use wiremock::matchers::{method, path}; +use wiremock::{Mock, MockServer, ResponseTemplate}; +use yoagent::provider::{GoogleProvider, ModelConfig, StreamConfig, StreamProvider}; +use yoagent::types::*; + +const MODEL: &str = "gemini-2.5-flash"; + +fn sse(events: &[&str]) -> String { + events + .iter() + .map(|data| format!("data: {}\r\n\r\n", data)) + .collect() +} + +fn stream_config(base_url: &str, messages: Vec) -> StreamConfig { + let mut mc = ModelConfig::google(MODEL, "Gemini 2.5 Flash"); + mc.base_url = base_url.to_string(); + StreamConfig { + model: MODEL.into(), + system_prompt: "test".into(), + messages, + tools: vec![], + thinking_level: ThinkingLevel::Off, + api_key: "test-key".into(), + max_tokens: Some(256), + temperature: None, + model_config: Some(mc), + cache_config: CacheConfig::default(), + } +} + +async fn run_stream(config: StreamConfig) -> Message { + let (tx, _rx) = mpsc::unbounded_channel(); + GoogleProvider + .stream(config, tx, CancellationToken::new()) + .await + .expect("stream should succeed") +} + +/// A streamed function call with a thoughtSignature must surface as a +/// ToolCall with the signature preserved in provider_metadata and a +/// synthetic id when Gemini sends none. +#[tokio::test] +async fn function_call_with_thought_signature_is_captured() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path(format!( + "/v1beta/models/{}:streamGenerateContent", + MODEL + ))) + .respond_with(ResponseTemplate::new(200).set_body_raw( + sse(&[ + r#"{"candidates":[{"content":{"parts":[{"functionCall":{"name":"get_weather","args":{"city":"Paris"}},"thoughtSignature":"sig-abc"}],"role":"model"},"finishReason":"STOP","index":0}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":5,"totalTokenCount":15}}"#, + ]), + "text/event-stream", + )) + .mount(&server) + .await; + + let message = run_stream(stream_config( + &server.uri(), + vec![Message::user("weather?")], + )) + .await; + + let Message::Assistant { + content, + stop_reason, + .. + } = &message + else { + panic!("expected assistant message"); + }; + assert_eq!(*stop_reason, StopReason::ToolUse); + + let Some(Content::ToolCall { + id, + name, + arguments, + provider_metadata, + .. + }) = content.first() + else { + panic!("expected a tool call, got {content:?}"); + }; + assert_eq!(name, "get_weather"); + assert_eq!(arguments["city"], "Paris"); + assert_eq!(id, "google-fc-0", "missing id must be synthesized"); + assert_eq!( + provider_metadata + .as_ref() + .and_then(|m| m["thought_signature"].as_str()), + Some("sig-abc"), + "thought signature must be preserved in provider_metadata" + ); + + let Message::Assistant { usage, .. } = &message else { + unreachable!() + }; + assert_eq!((usage.input, usage.output, usage.total_tokens), (10, 5, 15)); +} + +/// Multi-turn: when the history contains a prior tool call carrying a +/// thought signature, the next request must echo the signature back to +/// Gemini and must NOT leak the synthetic `google-fc-` id. +#[tokio::test] +async fn thought_signature_round_trips_and_synthetic_id_is_stripped() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path(format!( + "/v1beta/models/{}:streamGenerateContent", + MODEL + ))) + .respond_with(ResponseTemplate::new(200).set_body_raw( + sse(&[ + r#"{"candidates":[{"content":{"parts":[{"text":"It is 22C in Paris."}],"role":"model"},"finishReason":"STOP","index":0}]}"#, + ]), + "text/event-stream", + )) + .expect(1) + .mount(&server) + .await; + + let history = vec![ + Message::user("weather in Paris?"), + Message::assistant( + vec![Content::tool_call_with_metadata( + "google-fc-0", + "get_weather", + serde_json::json!({"city": "Paris"}), + serde_json::json!({"thought_signature": "sig-abc"}), + )], + StopReason::ToolUse, + MODEL, + "google", + Usage::default(), + ), + Message::ToolResult { + tool_call_id: "google-fc-0".into(), + tool_name: "get_weather".into(), + content: vec![Content::Text { text: "22C".into() }], + is_error: false, + timestamp: 2, + }, + ]; + + let message = run_stream(stream_config(&server.uri(), history)).await; + + // The follow-up turn parses normally + let Message::Assistant { content, .. } = &message else { + panic!("expected assistant message"); + }; + assert!(matches!( + content.first(), + Some(Content::Text { text }) if text.contains("22C") + )); + + // Inspect the actual request body sent to the gateway + let requests = server.received_requests().await.expect("recording enabled"); + assert_eq!(requests.len(), 1); + let body: serde_json::Value = requests[0].body_json().expect("json body"); + // Structural: the signature must sit on the functionCall part of the + // assistant turn (contents[1]), not merely appear somewhere in the body. + assert_eq!( + body["contents"][1]["parts"][0]["thoughtSignature"], "sig-abc", + "thought signature must be echoed on the functionCall part, body: {body}" + ); + let raw = serde_json::to_string(&body).unwrap(); + assert!( + !raw.contains("google-fc-0"), + "synthetic tool-call id must not be sent to Gemini, body: {raw}" + ); +} + +/// A part carrying BOTH empty text and a functionCall must still produce the +/// tool call (the old loop `continue`d past it while Gemini was thinking). +#[tokio::test] +async fn empty_text_part_does_not_swallow_function_call() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path(format!( + "/v1beta/models/{}:streamGenerateContent", + MODEL + ))) + .respond_with(ResponseTemplate::new(200).set_body_raw( + sse(&[ + r#"{"candidates":[{"content":{"parts":[{"text":"","functionCall":{"name":"get_weather","args":{"city":"Oslo"}}}],"role":"model"},"finishReason":"STOP","index":0}]}"#, + ]), + "text/event-stream", + )) + .mount(&server) + .await; + + let message = run_stream(stream_config(&server.uri(), vec![Message::user("hi")])).await; + + let Message::Assistant { + content, + stop_reason, + .. + } = &message + else { + panic!("expected assistant message"); + }; + assert_eq!(*stop_reason, StopReason::ToolUse); + assert!( + matches!(content.first(), Some(Content::ToolCall { name, .. }) if name == "get_weather"), + "functionCall in an empty-text part must not be dropped, got {content:?}" + ); +} + +/// Text deltas across multiple SSE events accumulate into ONE Content::Text. +#[tokio::test] +async fn text_deltas_accumulate_across_events() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path(format!( + "/v1beta/models/{}:streamGenerateContent", + MODEL + ))) + .respond_with(ResponseTemplate::new(200).set_body_raw( + sse(&[ + r#"{"candidates":[{"content":{"parts":[{"text":"Hello, "}],"role":"model"},"index":0}]}"#, + r#"{"candidates":[{"content":{"parts":[{"text":"world!"}],"role":"model"},"finishReason":"STOP","index":0}]}"#, + ]), + "text/event-stream", + )) + .mount(&server) + .await; + + let message = run_stream(stream_config(&server.uri(), vec![Message::user("hi")])).await; + + let Message::Assistant { content, .. } = &message else { + panic!("expected assistant message"); + }; + assert_eq!(content.len(), 1, "deltas must merge into one text block"); + assert!(matches!(content.first(), Some(Content::Text { text }) if text == "Hello, world!")); +} + +/// A mid-stream {"error": ...} payload must fail the stream, not vanish +/// into an empty chunk and a fake successful turn. +#[tokio::test] +async fn in_stream_error_payload_fails_the_stream() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path(format!( + "/v1beta/models/{}:streamGenerateContent", + MODEL + ))) + .respond_with(ResponseTemplate::new(200).set_body_raw( + sse(&[ + r#"{"error":{"code":429,"message":"Resource has been exhausted","status":"RESOURCE_EXHAUSTED"}}"#, + ]), + "text/event-stream", + )) + .mount(&server) + .await; + + let (tx, _rx) = tokio::sync::mpsc::unbounded_channel(); + let result = yoagent::provider::GoogleProvider + .stream( + stream_config(&server.uri(), vec![Message::user("hi")]), + tx, + CancellationToken::new(), + ) + .await; + + let err = result.expect_err("in-stream error must surface as Err"); + assert!( + err.to_string().contains("RESOURCE_EXHAUSTED"), + "error should carry the provider payload, got: {err}" + ); +} + +/// A SAFETY finish reason maps to StopReason::Refusal with an explanation. +#[tokio::test] +async fn safety_finish_reason_maps_to_refusal() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path(format!( + "/v1beta/models/{}:streamGenerateContent", + MODEL + ))) + .respond_with(ResponseTemplate::new(200).set_body_raw( + sse(&[ + r#"{"candidates":[{"content":{"parts":[],"role":"model"},"finishReason":"SAFETY","index":0}]}"#, + ]), + "text/event-stream", + )) + .mount(&server) + .await; + + let message = run_stream(stream_config(&server.uri(), vec![Message::user("hi")])).await; + + let Message::Assistant { + stop_reason, + error_message, + .. + } = &message + else { + panic!("expected assistant message"); + }; + assert_eq!(*stop_reason, StopReason::Refusal); + assert!( + error_message.as_deref().unwrap_or("").contains("SAFETY"), + "error_message should explain the block, got {error_message:?}" + ); +} diff --git a/tests/serialization_test.rs b/tests/serialization_test.rs index ba705fa..a1c3627 100644 --- a/tests/serialization_test.rs +++ b/tests/serialization_test.rs @@ -27,31 +27,25 @@ fn test_message_user_roundtrip() { #[test] fn test_message_assistant_roundtrip() { - let msg = Message::Assistant { - content: vec![ + let msg = Message::assistant( + vec![ Content::Text { text: "Hi there".into(), }, - Content::ToolCall { - provider_metadata: None, - id: "tc-1".into(), - name: "read_file".into(), - arguments: serde_json::json!({"path": "foo.rs"}), - }, + Content::tool_call("tc-1", "read_file", serde_json::json!({"path": "foo.rs"})), ], - stop_reason: StopReason::ToolUse, - model: "claude-sonnet".into(), - provider: "anthropic".into(), - usage: Usage { + StopReason::ToolUse, + "claude-sonnet", + "anthropic", + Usage { input: 100, output: 50, cache_read: 10, cache_write: 5, total_tokens: 165, }, - timestamp: 789, - error_message: None, - }; + ) + .with_timestamp(789); roundtrip(&msg); } @@ -101,16 +95,13 @@ fn test_content_variants_roundtrip() { data: "base64data".into(), mime_type: "image/png".into(), }); - roundtrip(&Content::Thinking { - thinking: "let me think...".into(), - signature: Some("sig123".into()), - }); - roundtrip(&Content::ToolCall { - provider_metadata: None, - id: "tc-1".into(), - name: "bash".into(), - arguments: serde_json::json!({"command": "ls"}), - }); + roundtrip(&Content::thinking_signed("let me think...", "sig123")); + roundtrip(&Content::thinking("unsigned thought")); + roundtrip(&Content::tool_call( + "tc-1", + "bash", + serde_json::json!({"command": "ls"}), + )); } // --------------------------------------------------------------------------- @@ -121,20 +112,20 @@ fn test_content_variants_roundtrip() { fn test_full_conversation_roundtrip() { let conversation: Vec = vec![ AgentMessage::Llm(Message::user("Read the file")), - AgentMessage::Llm(Message::Assistant { - content: vec![Content::ToolCall { - provider_metadata: None, - id: "tc-1".into(), - name: "read_file".into(), - arguments: serde_json::json!({"path": "main.rs"}), - }], - stop_reason: StopReason::ToolUse, - model: "mock".into(), - provider: "mock".into(), - usage: Usage::default(), - timestamp: 100, - error_message: None, - }), + AgentMessage::Llm( + Message::assistant( + vec![Content::tool_call( + "tc-1", + "read_file", + serde_json::json!({"path": "main.rs"}), + )], + StopReason::ToolUse, + "mock", + "mock", + Usage::default(), + ) + .with_timestamp(100), + ), AgentMessage::Llm(Message::ToolResult { tool_call_id: "tc-1".into(), tool_name: "read_file".into(), @@ -144,17 +135,18 @@ fn test_full_conversation_roundtrip() { is_error: false, timestamp: 200, }), - AgentMessage::Llm(Message::Assistant { - content: vec![Content::Text { - text: "The file contains a main function.".into(), - }], - stop_reason: StopReason::Stop, - model: "mock".into(), - provider: "mock".into(), - usage: Usage::default(), - timestamp: 300, - error_message: None, - }), + AgentMessage::Llm( + Message::assistant( + vec![Content::Text { + text: "The file contains a main function.".into(), + }], + StopReason::Stop, + "mock", + "mock", + Usage::default(), + ) + .with_timestamp(300), + ), AgentMessage::Extension(ExtensionMessage::new( "ui_event", serde_json::json!({"action": "scroll"}), @@ -196,18 +188,88 @@ fn test_tool_execution_strategy_roundtrip() { fn test_refusal_stop_reason_round_trip() { use yoagent::types::*; - let message = Message::Assistant { - content: vec![], - stop_reason: StopReason::Refusal, - model: "claude-fable-5".into(), - provider: "anthropic".into(), - usage: Usage::default(), - timestamp: 1, - error_message: Some("declined".into()), - }; + let message = Message::assistant( + vec![], + StopReason::Refusal, + "claude-fable-5", + "anthropic", + Usage::default(), + ) + .with_timestamp(1) + .with_error_message("declined"); let json = serde_json::to_string(&message).unwrap(); let back: Message = serde_json::from_str(&json).unwrap(); assert_eq!(message, back); assert_eq!(StopReason::Refusal.to_string(), "refusal"); } + +#[test] +fn test_constructors_pin_fields() { + // The non_exhaustive variants make these the mandated construction path + // for downstream crates — pin every field explicitly. + let tc = Content::tool_call("id-1", "bash", serde_json::json!({"cmd": "ls"})); + let Content::ToolCall { + id, + name, + arguments, + provider_metadata, + .. + } = &tc + else { + panic!("expected tool call"); + }; + assert_eq!(id, "id-1"); + assert_eq!(name, "bash"); + assert_eq!(arguments["cmd"], "ls"); + assert!(provider_metadata.is_none()); + + let think = Content::thinking_signed("hmm", "sig"); + let Content::Thinking { + thinking, + signature, + .. + } = &think + else { + panic!("expected thinking"); + }; + assert_eq!(thinking, "hmm"); + assert_eq!(signature.as_deref(), Some("sig")); + + let msg = Message::assistant( + vec![Content::Text { text: "hi".into() }], + StopReason::Stop, + "model-x", + "provider-y", + Usage::default(), + ) + .with_timestamp(42) + .with_error_message("oops"); + let Message::Assistant { + model, + provider, + timestamp, + error_message, + stop_reason, + .. + } = &msg + else { + panic!("expected assistant"); + }; + assert_eq!(model, "model-x"); + assert_eq!(provider, "provider-y"); + assert_eq!(*timestamp, 42); + assert_eq!(error_message.as_deref(), Some("oops")); + assert_eq!(*stop_reason, StopReason::Stop); +} + +#[test] +fn test_tool_call_with_metadata_roundtrip() { + // Thought signatures must survive session persistence. + roundtrip(&Content::tool_call_with_metadata( + "tc-9", + "get_weather", + serde_json::json!({"city": "Paris"}), + serde_json::json!({"thought_signature": "sig-xyz"}), + )); +} diff --git a/tests/sub_agent_test.rs b/tests/sub_agent_test.rs index 371b02e..9785b50 100644 --- a/tests/sub_agent_test.rs +++ b/tests/sub_agent_test.rs @@ -300,17 +300,15 @@ async fn test_sub_agent_parallel() { content_index: 0, delta: self.text.clone(), }); - let msg = Message::Assistant { - content: vec![Content::Text { + let msg = Message::assistant( + vec![Content::Text { text: self.text.clone(), }], - stop_reason: StopReason::Stop, - model: "slow".into(), - provider: "slow".into(), - usage: Usage::default(), - timestamp: yoagent::now_ms(), - error_message: None, - }; + StopReason::Stop, + "slow", + "slow", + Usage::default(), + ); let _ = tx.send(yoagent::provider::StreamEvent::Done { message: msg.clone(), }); @@ -500,17 +498,15 @@ impl yoagent::provider::StreamProvider for CapturingProvider { ) -> Result { *self.captured.lock().unwrap() = config.system_prompt.clone(); let _ = tx.send(yoagent::provider::StreamEvent::Start); - let msg = Message::Assistant { - content: vec![Content::Text { + let msg = Message::assistant( + vec![Content::Text { text: "done".into(), }], - stop_reason: StopReason::Stop, - model: "mock".into(), - provider: "mock".into(), - usage: Usage::default(), - timestamp: yoagent::now_ms(), - error_message: None, - }; + StopReason::Stop, + "mock", + "mock", + Usage::default(), + ); let _ = tx.send(yoagent::provider::StreamEvent::Done { message: msg.clone(), });