Fix SSE parsing in Gemini provider

This fixes tool calling issues on some gemini models. The bug here won't
be observed if the server only returns small, completed chunks.

We should not process SSE *byte*stream line-by-line. Therefore, we need to
have correct buffering. Instead of reinventing the wheels, we introduce
new dependencies (reqwest-eventsource) for parsing SSE chunks.

We also set id in tool_call to its function name instead of empty string

Disclaimer: I did vibe code this (reviewed, tested internally though)

Author: quangIO

Cargo.lock

@@ -90,6 +90,7 @@ tracing-subscriber = "0.3.19"
 uuid = "1.17.0"
 worker = "0.6"
 zerocopy = "0.8.26"
+reqwest-eventsource = "0.6.0"
 
 [workspace.metadata.cargo-autoinherit]
 # Skip cargo-autoinherit for these packages

Cargo.toml

@@ -38,6 +38,7 @@ base64 = { workspace = true }
 as-any = { workspace = true }
 rmcp = { version = "0.3", optional = true, features = ["client"] }
 url = { workspace = true }
+reqwest-eventsource = { workspace = true }
 
 [dev-dependencies]
 anyhow = { workspace = true }

rig-core/Cargo.toml

@@ -0,0 +1,291 @@
+use futures::{Stream, StreamExt};
+use rig::providers::gemini;
+use rig::{
+    OneOrMany,
+    agent::Agent,
+    client::{CompletionClient, ProviderClient},
+    completion::{self, CompletionError, CompletionModel, PromptError, ToolDefinition},
+    message::{AssistantContent, Message, Text, ToolResultContent, UserContent},
+    streaming::{StreamedAssistantContent, StreamingCompletion},
+    tool::{Tool, ToolSetError},
+};
+use schemars::{JsonSchema, schema_for};
+use serde::{Deserialize, Serialize};
+
+use std::pin::Pin;
+use thiserror::Error;
+
+#[derive(Debug, Error)]
+enum StreamingError {
+    #[error("CompletionError: {0}")]
+    Completion(#[from] CompletionError),
+    #[error("PromptError: {0}")]
+    Prompt(#[from] PromptError),
+    #[error("ToolSetError: {0}")]
+    Tool(#[from] ToolSetError),
+}
+
+type StreamingResult = Pin<Box<dyn Stream<Item = Result<Text, StreamingError>> + Send>>;
+
+#[tokio::main]
+async fn main() -> anyhow::Result<()> {
+    // tracing_subscriber::registry()
+    //     .with(
+    //         tracing_subscriber::EnvFilter::try_from_default_env()
+    //             .unwrap_or_else(|_| "info".into()),
+    //     )
+    //     .with(tracing_subscriber::fmt::layer())
+    //     .init();
+
+    // Create gemini client
+    let llm_client = gemini::Client::from_env();
+
+    // Create agent with a single context prompt and a calculator tools
+    let calculator_agent = llm_client
+        .agent("gemini-2.5-flash")
+        .preamble("You are an calculator. You must use tools to get the user result")
+        .tool(Add)
+        .tool(Subtract)
+        .tool(Multiply)
+        .tool(Divide)
+        .build();
+
+    // Prompt the agent and get the stream
+    let mut stream = multi_turn_prompt(
+        calculator_agent,
+        "Calculate 2 * (3 + 5) / 9  = ?. Describe the result to me.",
+        Vec::new(),
+    )
+    .await;
+
+    custom_stream_to_stdout(&mut stream).await?;
+
+    Ok(())
+}
+
+async fn multi_turn_prompt<M>(
+    agent: Agent<M>,
+    prompt: impl Into<Message> + Send,
+    mut chat_history: Vec<completion::Message>,
+) -> StreamingResult
+where
+    M: CompletionModel + 'static,
+    <M as CompletionModel>::StreamingResponse: std::marker::Send,
+{
+    let prompt: Message = prompt.into();
+
+    (Box::pin(async_stream::stream! {
+        let mut current_prompt = prompt;
+        let mut did_call_tool = false;
+
+        'outer: loop {
+            let mut stream = agent
+                .stream_completion(current_prompt.clone(), chat_history.clone())
+                .await?
+                .stream()
+                .await?;
+
+            chat_history.push(current_prompt.clone());
+
+            let mut tool_calls = vec![];
+            let mut tool_results = vec![];
+
+            while let Some(content) = stream.next().await {
+                match content {
+                    Ok(StreamedAssistantContent::Text(text)) => {
+                        yield Ok(Text { text: text.text });
+                        did_call_tool = false;
+                    },
+                    Ok(StreamedAssistantContent::ToolCall(tool_call)) => {
+                        let tool_result =
+                            agent.tools.call(&tool_call.function.name, tool_call.function.arguments.to_string()).await?;
+
+                        let tool_call_msg = AssistantContent::ToolCall(tool_call.clone());
+
+                        tool_calls.push(tool_call_msg);
+                        tool_results.push((tool_call.id, tool_call.call_id, tool_result));
+
+                        did_call_tool = true;
+                        // break;
+                    },
+                    Ok(StreamedAssistantContent::Reasoning(rig::message::Reasoning { reasoning })) => {
+                        yield Ok(Text { text: reasoning });
+                        did_call_tool = false;
+                    },
+                    Ok(_) => {
+                        // do nothing here as we don't need to accumulate token usage
+                    }
+                    Err(e) => {
+                        yield Err(e.into());
+                        break 'outer;
+                    }
+                }
+            }
+
+            // Add (parallel) tool calls to chat history
+            if !tool_calls.is_empty() {
+                chat_history.push(Message::Assistant {
+                    id: None,
+                    content: OneOrMany::many(tool_calls).expect("Impossible EmptyListError"),
+                });
+            }
+
+            // Add tool results to chat history
+            for (id, call_id, tool_result) in tool_results {
+                if let Some(call_id) = call_id {
+                    chat_history.push(Message::User {
+                        content: OneOrMany::one(UserContent::tool_result_with_call_id(
+                            id,
+                            call_id,
+                            OneOrMany::one(ToolResultContent::text(tool_result)),
+                        )),
+                    });
+                } else {
+                    chat_history.push(Message::User {
+                        content: OneOrMany::one(UserContent::tool_result(
+                            id,
+                            OneOrMany::one(ToolResultContent::text(tool_result)),
+                        )),
+                    });
+
+                }
+
+            }
+
+            // Set the current prompt to the last message in the chat history
+            current_prompt = match chat_history.pop() {
+                Some(prompt) => prompt,
+                None => unreachable!("Chat history should never be empty at this point"),
+            };
+
+            if !did_call_tool {
+                break;
+            }
+        }
+
+    })) as _
+}
+
+/// helper function to stream a completion request to stdout
+async fn custom_stream_to_stdout(stream: &mut StreamingResult) -> Result<(), std::io::Error> {
+    print!("Response: ");
+    while let Some(content) = stream.next().await {
+        match content {
+            Ok(Text { text }) => {
+                print!("{text}");
+                std::io::Write::flush(&mut std::io::stdout())?;
+            }
+            Err(err) => {
+                eprintln!("Error: {err:#?}");
+            }
+        }
+    }
+    println!(); // New line after streaming completes
+
+    Ok(())
+}
+
+#[derive(Deserialize, JsonSchema)]
+struct OperationArgs {
+    x: i32,
+    y: i32,
+}
+
+#[derive(Debug, thiserror::Error)]
+#[error("Math error")]
+struct MathError;
+
+#[derive(Deserialize, Serialize)]
+struct Add;
+impl Tool for Add {
+    const NAME: &'static str = "add";
+
+    type Error = MathError;
+    type Args = OperationArgs;
+    type Output = i32;
+
+    async fn definition(&self, _prompt: String) -> ToolDefinition {
+        ToolDefinition {
+            name: "add".to_string(),
+            description: "Add x and y together".to_string(),
+            parameters: serde_json::to_value(schema_for!(OperationArgs))
+                .expect("converting JSON schema to JSON value should never fail"),
+        }
+    }
+
+    async fn call(&self, args: Self::Args) -> Result<Self::Output, Self::Error> {
+        let result = args.x + args.y;
+        Ok(result)
+    }
+}
+
+#[derive(Deserialize, Serialize)]
+struct Subtract;
+impl Tool for Subtract {
+    const NAME: &'static str = "subtract";
+
+    type Error = MathError;
+    type Args = OperationArgs;
+    type Output = i32;
+
+    async fn definition(&self, _prompt: String) -> ToolDefinition {
+        ToolDefinition {
+            name: "subtract".to_string(),
+            description: "Subtract y from x (i.e.: x - y)".to_string(),
+            parameters: serde_json::to_value(schema_for!(OperationArgs))
+                .expect("converting JSON schema to JSON value should never fail"),
+        }
+    }
+
+    async fn call(&self, args: Self::Args) -> Result<Self::Output, Self::Error> {
+        let result = args.x - args.y;
+        Ok(result)
+    }
+}
+
+struct Multiply;
+impl Tool for Multiply {
+    const NAME: &'static str = "multiply";
+
+    type Error = MathError;
+    type Args = OperationArgs;
+    type Output = i32;
+
+    async fn definition(&self, _prompt: String) -> ToolDefinition {
+        ToolDefinition {
+            name: "multiply".to_string(),
+            description: "Compute the product of x and y (i.e.: x * y)".to_string(),
+            parameters: serde_json::to_value(schema_for!(OperationArgs))
+                .expect("converting JSON schema to JSON value should never fail"),
+        }
+    }
+
+    async fn call(&self, args: Self::Args) -> Result<Self::Output, Self::Error> {
+        let result = args.x * args.y;
+        Ok(result)
+    }
+}
+
+struct Divide;
+impl Tool for Divide {
+    const NAME: &'static str = "divide";
+
+    type Error = MathError;
+    type Args = OperationArgs;
+    type Output = i32;
+
+    async fn definition(&self, _prompt: String) -> ToolDefinition {
+        ToolDefinition {
+            name: "divide".to_string(),
+            description: "Compute the Quotient of x and y (i.e.: x / y). Useful for ratios."
+                .to_string(),
+            parameters: serde_json::to_value(schema_for!(OperationArgs))
+                .expect("converting JSON schema to JSON value should never fail"),
+        }
+    }
+
+    async fn call(&self, args: Self::Args) -> Result<Self::Output, Self::Error> {
+        let result = args.x / args.y;
+        Ok(result)
+    }
+}