Add Rig native model providers

crates/galaxy_agent_rig/Cargo.toml

@@ -14,6 +14,7 @@ futures.workspace = true
 galaxy_agent_core.workspace = true
 rig-core.workspace = true
 rig-bedrock.workspace = true
+rig-vertexai.workspace = true
 serde_json.workspace = true
 uuid.workspace = true
 

crates/galaxy_agent_rig/src/bedrock_tests.rs

@@ -6,7 +6,7 @@ use galaxy_agent_core::{
     StopReason, ToolDefinition, TurnCommand, TurnRequest, Usage,
 };
 use rig_bedrock::streaming::{BedrockStreamingResponse, BedrockUsage};
-use rig_core::completion::{AssistantContent, CompletionError, GetTokenUsage, Message};
+use rig_core::completion::{AssistantContent, CompletionError, Message};
 use rig_core::message::{DocumentSourceKind, ToolResultContent, UserContent};
 
 use super::*;
@@ -78,9 +78,10 @@ fn normalizes_bedrock_usage_and_max_token_stop() {
             cache_read_input_tokens: Some(40),
             cache_write_input_tokens: Some(10),
         }),
+        stop_reason: None,
     };
     assert_eq!(
-        map_usage(response.token_usage()),
+        map_usage((&response).into()),
         Usage {
             input_tokens: 100,
             output_tokens: 25,

crates/galaxy_agent_rig/src/lib.rs

@@ -2,10 +2,12 @@
 
 mod bedrock;
 mod chatgpt;
+mod native;
 mod openai_compatible;
 mod request;
 mod stream;
 
 pub use bedrock::*;
 pub use chatgpt::*;
+pub use native::*;
 pub use openai_compatible::*;

crates/galaxy_agent_rig/src/native.rs

@@ -0,0 +1,256 @@
+use async_trait::async_trait;
+use galaxy_agent_core::{
+    AgentError, AgentErrorKind, AgentEventStream, AgentRuntime, RuntimeCapabilities,
+    RuntimeDescriptor, RuntimeKind, TurnControl, TurnRequest,
+};
+use rig_core::client::{CompletionClient, ModelListingClient};
+use rig_core::providers::{anthropic, gemini};
+
+use crate::request::build_completion_request;
+use crate::stream::{start_model_completion, start_model_turn};
+
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct RigModelInfo {
+    pub id: String,
+    pub display_name: String,
+    pub context_size: Option<u32>,
+}
+
+pub async fn discover_anthropic_models(api_key: &str) -> Result<Vec<RigModelInfo>, String> {
+    let client = anthropic::Client::new(api_key).map_err(|error| error.to_string())?;
+    let models = client
+        .list_models()
+        .await
+        .map_err(|error| error.to_string())?;
+    Ok(models
+        .into_iter()
+        .map(|model| RigModelInfo {
+            display_name: model.display_name().to_string(),
+            id: model.id,
+            context_size: model.context_length,
+        })
+        .collect())
+}
+
+pub async fn discover_gemini_models(api_key: &str) -> Result<Vec<RigModelInfo>, String> {
+    let client = gemini::Client::new(api_key).map_err(|error| error.to_string())?;
+    let models = client
+        .list_models()
+        .await
+        .map_err(|error| error.to_string())?;
+    Ok(models
+        .into_iter()
+        .filter(|model| !model.id.contains("embedding"))
+        .map(|model| RigModelInfo {
+            display_name: model.display_name().to_string(),
+            id: model.id,
+            context_size: model.context_length,
+        })
+        .collect())
+}
+
+pub fn vertex_ai_model_catalog() -> Vec<RigModelInfo> {
+    [
+        (rig_vertexai::completion::GEMINI_2_5_PRO, "Gemini 2.5 Pro"),
+        (
+            rig_vertexai::completion::GEMINI_2_5_FLASH,
+            "Gemini 2.5 Flash",
+        ),
+        (
+            rig_vertexai::completion::GEMINI_2_5_FLASH_LITE,
+            "Gemini 2.5 Flash Lite",
+        ),
+        (rig_vertexai::completion::GEMINI_1_5_PRO, "Gemini 1.5 Pro"),
+        (
+            rig_vertexai::completion::GEMINI_1_5_FLASH,
+            "Gemini 1.5 Flash",
+        ),
+    ]
+    .into_iter()
+    .map(|(id, display_name)| RigModelInfo {
+        id: id.to_string(),
+        display_name: display_name.to_string(),
+        context_size: None,
+    })
+    .collect()
+}
+
+pub fn validate_vertex_ai_credentials(project_id: &str, location: &str) -> Result<(), String> {
+    rig_vertexai::Client::builder()
+        .with_project(project_id)
+        .with_location(location)
+        .build()
+        .map(|_| ())
+        .map_err(|error| error.to_string())
+}
+
+#[derive(Clone, Debug)]
+pub struct AnthropicRuntimeConfig {
+    pub api_key: String,
+    pub model: String,
+    pub max_output_tokens: Option<u64>,
+}
+
+#[derive(Clone, Debug)]
+pub struct AnthropicRuntime {
+    config: AnthropicRuntimeConfig,
+    descriptor: RuntimeDescriptor,
+}
+
+impl AnthropicRuntime {
+    pub fn new(config: AnthropicRuntimeConfig) -> Self {
+        let descriptor = native_descriptor("anthropic", &config.model);
+        Self { config, descriptor }
+    }
+}
+
+#[async_trait]
+impl AgentRuntime for AnthropicRuntime {
+    fn descriptor(&self) -> &RuntimeDescriptor {
+        &self.descriptor
+    }
+
+    async fn start_turn(
+        &self,
+        request: TurnRequest,
+        control: TurnControl,
+    ) -> Result<AgentEventStream, AgentError> {
+        if self.config.api_key.trim().is_empty() {
+            return Err(AgentError::new(
+                AgentErrorKind::Configuration,
+                "An Anthropic API key is required",
+            ));
+        }
+        let client = anthropic::Client::new(&self.config.api_key)
+            .map_err(|error| AgentError::new(AgentErrorKind::Configuration, error.to_string()))?;
+        let model = client.completion_model(&self.config.model);
+        let max_output_tokens = request.max_output_tokens.or(self.config.max_output_tokens);
+        let completion_request =
+            build_completion_request(request, max_output_tokens, true, true, None)?;
+        start_model_turn(model, completion_request, control, max_output_tokens).await
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct GeminiRuntimeConfig {
+    pub api_key: String,
+    pub model: String,
+    pub max_output_tokens: Option<u64>,
+}
+
+#[derive(Clone, Debug)]
+pub struct GeminiRuntime {
+    config: GeminiRuntimeConfig,
+    descriptor: RuntimeDescriptor,
+}
+
+impl GeminiRuntime {
+    pub fn new(config: GeminiRuntimeConfig) -> Self {
+        let descriptor = native_descriptor("gemini", &config.model);
+        Self { config, descriptor }
+    }
+}
+
+#[async_trait]
+impl AgentRuntime for GeminiRuntime {
+    fn descriptor(&self) -> &RuntimeDescriptor {
+        &self.descriptor
+    }
+
+    async fn start_turn(
+        &self,
+        request: TurnRequest,
+        control: TurnControl,
+    ) -> Result<AgentEventStream, AgentError> {
+        if self.config.api_key.trim().is_empty() {
+            return Err(AgentError::new(
+                AgentErrorKind::Configuration,
+                "A Gemini API key is required",
+            ));
+        }
+        let client = gemini::Client::new(&self.config.api_key)
+            .map_err(|error| AgentError::new(AgentErrorKind::Configuration, error.to_string()))?;
+        let model = client.completion_model(&self.config.model);
+        let max_output_tokens = request.max_output_tokens.or(self.config.max_output_tokens);
+        let completion_request =
+            build_completion_request(request, max_output_tokens, true, true, None)?;
+        start_model_turn(model, completion_request, control, max_output_tokens).await
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct VertexAiRuntimeConfig {
+    pub project_id: String,
+    pub location: String,
+    pub model: String,
+    pub max_output_tokens: Option<u64>,
+}
+
+#[derive(Clone, Debug)]
+pub struct VertexAiRuntime {
+    config: VertexAiRuntimeConfig,
+    descriptor: RuntimeDescriptor,
+}
+
+impl VertexAiRuntime {
+    pub fn new(config: VertexAiRuntimeConfig) -> Self {
+        let descriptor = native_descriptor("vertex-ai", &config.model);
+        Self { config, descriptor }
+    }
+}
+
+#[async_trait]
+impl AgentRuntime for VertexAiRuntime {
+    fn descriptor(&self) -> &RuntimeDescriptor {
+        &self.descriptor
+    }
+
+    async fn start_turn(
+        &self,
+        request: TurnRequest,
+        control: TurnControl,
+    ) -> Result<AgentEventStream, AgentError> {
+        if self.config.project_id.trim().is_empty() {
+            return Err(AgentError::new(
+                AgentErrorKind::Configuration,
+                "A Google Cloud project ID is required for Vertex AI",
+            ));
+        }
+        let client = rig_vertexai::Client::builder()
+            .with_project(&self.config.project_id)
+            .with_location(if self.config.location.trim().is_empty() {
+                "global"
+            } else {
+                &self.config.location
+            })
+            .build()
+            .map_err(|error| AgentError::new(AgentErrorKind::Configuration, error.to_string()))?;
+        let model = client.completion_model(&self.config.model);
+        let max_output_tokens = request.max_output_tokens.or(self.config.max_output_tokens);
+        let completion_request =
+            build_completion_request(request, max_output_tokens, true, true, None)?;
+        start_model_completion(model, completion_request, control, max_output_tokens).await
+    }
+}
+
+fn native_descriptor(provider: &str, model: &str) -> RuntimeDescriptor {
+    RuntimeDescriptor {
+        id: format!("rig-{provider}:{model}"),
+        display_name: format!("Rig / {provider} / {model}"),
+        kind: RuntimeKind::Provider,
+        capabilities: RuntimeCapabilities::provider(),
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::native_descriptor;
+
+    #[test]
+    fn native_descriptors_are_provider_specific() {
+        let descriptor = native_descriptor("anthropic", "claude-sonnet");
+
+        assert_eq!(descriptor.id, "rig-anthropic:claude-sonnet");
+        assert_eq!(descriptor.display_name, "Rig / anthropic / claude-sonnet");
+    }
+}

crates/galaxy_agent_rig/src/openai_compatible.rs

@@ -76,7 +76,6 @@ async fn start_model_turn<M>(
 ) -> Result<AgentEventStream, AgentError>
 where
     M: CompletionModel + Send + Sync + 'static,
-    M::StreamingResponse: Send + Sync + 'static,
 {
     let max_output_tokens = request.max_output_tokens.or(configured_max_output_tokens);
     let completion_request = build_completion_request(

crates/galaxy_agent_rig/src/stream.rs

@@ -3,7 +3,7 @@ use galaxy_agent_core::{
     AgentError, AgentErrorKind, AgentEvent, AgentEventStream, StopReason, ToolCall, TurnCommand,
     TurnControl, Usage,
 };
-use rig_core::completion::{CompletionError, CompletionModel, CompletionRequest, GetTokenUsage};
+use rig_core::completion::{CompletionError, CompletionModel, CompletionRequest};
 use rig_core::streaming::StreamedAssistantContent;
 use uuid::Uuid;
 
@@ -15,7 +15,6 @@ pub(crate) async fn start_model_turn<M>(
 ) -> Result<AgentEventStream, AgentError>
 where
     M: CompletionModel + Send + Sync + 'static,
-    M::StreamingResponse: Send + Sync + 'static,
 {
     let runtime_request_id = Uuid::new_v4().to_string();
     let stream_future = model.stream(completion_request).fuse();
@@ -112,7 +111,7 @@ where
                             // is the canonical event Galaxy consumes.
                         }
                         Ok(StreamedAssistantContent::Final(response)) => {
-                            let mapped_usage = map_usage(response.token_usage());
+                            let mapped_usage = map_usage(response.usage);
                             last_output_tokens = mapped_usage.output_tokens;
                             yield Ok(AgentEvent::UsageUpdated {
                                 usage: mapped_usage,
@@ -142,6 +141,82 @@ where
     Ok(Box::pin(events))
 }
 
+pub(crate) async fn start_model_completion<M>(
+    model: M,
+    completion_request: CompletionRequest,
+    control: TurnControl,
+    max_output_tokens: Option<u64>,
+) -> Result<AgentEventStream, AgentError>
+where
+    M: CompletionModel + Send + Sync + 'static,
+{
+    let runtime_request_id = Uuid::new_v4().to_string();
+    let completion_future = model.completion(completion_request).fuse();
+    let initial_control = control.clone();
+    let control_future = initial_control.receive().fuse();
+    futures::pin_mut!(completion_future, control_future);
+
+    let response = futures::select_biased! {
+        command = control_future => match command {
+            Ok(TurnCommand::Cancel) => {
+                return Ok(stopped_before_stream(runtime_request_id));
+            }
+            Ok(TurnCommand::Steer { .. }) | Err(_) => {
+                completion_future.await.map_err(map_completion_error)?
+            }
+        },
+        result = completion_future => result.map_err(map_completion_error)?,
+    };
+
+    let events = async_stream::stream! {
+        yield Ok(AgentEvent::TurnStarted { runtime_request_id });
+
+        for content in response.choice {
+            match content {
+                rig_core::completion::AssistantContent::Text(text) => {
+                    if !text.text.is_empty() {
+                        yield Ok(AgentEvent::TextDelta { text: text.text });
+                    }
+                }
+                rig_core::completion::AssistantContent::Reasoning(reasoning) => {
+                    yield Ok(AgentEvent::ReasoningCompleted {
+                        text: reasoning.display_text(),
+                        signature: reasoning.first_signature().map(str::to_string),
+                    });
+                }
+                rig_core::completion::AssistantContent::ToolCall(tool_call) => {
+                    yield Ok(AgentEvent::Tool {
+                        event: galaxy_agent_core::ToolEvent::Proposed {
+                            call: domain_tool_call(tool_call),
+                        },
+                    });
+                }
+                rig_core::completion::AssistantContent::Image(_) => {
+                    yield Err(AgentError::new(
+                        AgentErrorKind::Protocol,
+                        "Rig returned an unsupported image completion content block",
+                    ));
+                    return;
+                }
+            }
+        }
+
+        let usage = map_usage(response.usage);
+        let reached_max_tokens = max_output_tokens
+            .is_some_and(|max| usage.output_tokens >= max);
+        yield Ok(AgentEvent::UsageUpdated { usage });
+        yield Ok(AgentEvent::TurnStopped {
+            reason: if reached_max_tokens {
+                StopReason::MaxTokens
+            } else {
+                StopReason::Completed
+            },
+        });
+    };
+
+    Ok(Box::pin(events))
+}
+
 fn domain_tool_call(tool_call: rig_core::message::ToolCall) -> ToolCall {
     ToolCall {
         // OpenAI Responses uses a separate `call_id` for function-call output