From e43ed3a6fddf076876b24dde012b9850d091fd9c Mon Sep 17 00:00:00 2001 From: Jack Amadeo Date: Wed, 22 Jul 2026 14:48:30 -0400 Subject: [PATCH] enhance the uniffi API layer (#10427) --- Cargo.lock | 14 +- crates/goose-cli/src/session/mod.rs | 1 + crates/goose-download-manager/Cargo.toml | 2 +- crates/goose-local-inference/Cargo.toml | 8 +- crates/goose-provider-types/Cargo.toml | 2 +- crates/goose-provider-types/src/base.rs | 32 +- .../src/formats/databricks.rs | 9 + .../src/formats/openai_responses.rs | 20 + crates/goose-provider-types/src/model.rs | 30 + crates/goose-providers/Cargo.toml | 6 +- crates/goose-providers/src/anthropic.rs | 5 +- crates/goose-providers/src/api_client.rs | 72 ++ crates/goose-providers/src/databricks.rs | 21 +- crates/goose-providers/src/databricks_v2.rs | 12 +- crates/goose-providers/src/google.rs | 16 +- crates/goose-providers/src/ollama.rs | 4 +- crates/goose-providers/src/openai.rs | 19 +- .../goose-providers/src/openai_compatible.rs | 4 +- crates/goose-providers/src/snowflake.rs | 12 +- crates/goose-sdk-types/Cargo.toml | 2 +- crates/goose-sdk/Cargo.toml | 10 +- crates/goose-sdk/README.md | 6 +- .../uniffi/kotlin/src/main/kotlin/Main.kt | 16 +- crates/goose-sdk/examples/uniffi/provider.py | 22 +- crates/goose-sdk/justfile | 6 + crates/goose-sdk/maven/README.md | 5 +- crates/goose-sdk/maven/build.gradle.kts | 4 - .../io/github/aaif_goose/ProviderFlow.kt | 13 +- .../providers/anthropic/Anthropic.kt | 9 + .../github/aaif_goose/providers/groq/Groq.kt | 5 + crates/goose-sdk/python/pyproject.toml | 2 +- .../scripts/maven-resource-prefix.sh | 1 + .../scripts/prepare-maven-package.sh | 2 +- crates/goose-sdk/src/bindings.rs | 925 ++++++++++++++++-- crates/goose/src/context_mgmt/mod.rs | 1 + crates/goose/src/model_config.rs | 2 + crates/goose/src/providers/bedrock.rs | 1 + crates/goose/src/providers/gcpvertexai.rs | 11 +- crates/goose/src/providers/githubcopilot.rs | 8 +- crates/goose/src/providers/litellm.rs | 12 +- crates/goose/src/providers/nanogpt.rs | 4 +- crates/goose/src/providers/openrouter.rs | 5 +- crates/goose/src/providers/tetrate.rs | 4 +- 43 files changed, 1207 insertions(+), 158 deletions(-) create mode 100644 crates/goose-sdk/maven/src/support/kotlin/io/github/aaif_goose/providers/anthropic/Anthropic.kt create mode 100644 crates/goose-sdk/maven/src/support/kotlin/io/github/aaif_goose/providers/groq/Groq.kt diff --git a/Cargo.lock b/Cargo.lock index 239c77cd2..3eb60f173 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5108,7 +5108,7 @@ dependencies = [ [[package]] name = "goose-download-manager" -version = "0.1.0-alpha.0" +version = "0.1.0-alpha.5" dependencies = [ "anyhow", "once_cell", @@ -5120,7 +5120,7 @@ dependencies = [ [[package]] name = "goose-local-inference" -version = "0.1.0-alpha.0" +version = "0.1.0-alpha.5" dependencies = [ "anyhow", "async-stream", @@ -5186,7 +5186,7 @@ dependencies = [ [[package]] name = "goose-provider-types" -version = "0.1.0-alpha.0" +version = "0.1.0-alpha.5" dependencies = [ "anyhow", "async-stream", @@ -5215,7 +5215,7 @@ dependencies = [ [[package]] name = "goose-providers" -version = "0.1.0-alpha.0" +version = "0.1.0-alpha.5" dependencies = [ "anyhow", "async-stream", @@ -5247,14 +5247,16 @@ dependencies = [ [[package]] name = "goose-sdk" -version = "0.1.0-alpha.0" +version = "0.1.0-alpha.5" dependencies = [ "agent-client-protocol", "agent-client-protocol-schema", "anyhow", + "base64 0.22.1", "futures", "goose-providers", "goose-sdk-types", + "rmcp", "serde_json", "thiserror 2.0.18", "tokio", @@ -5264,7 +5266,7 @@ dependencies = [ [[package]] name = "goose-sdk-types" -version = "0.1.0-alpha.0" +version = "0.1.0-alpha.5" dependencies = [ "agent-client-protocol", "agent-client-protocol-schema", diff --git a/crates/goose-cli/src/session/mod.rs b/crates/goose-cli/src/session/mod.rs index 1bea927e8..c626aa695 100644 --- a/crates/goose-cli/src/session/mod.rs +++ b/crates/goose-cli/src/session/mod.rs @@ -2584,6 +2584,7 @@ mod tests { serde_json::json!(["output-128k-2025-02-19"]), )])), reasoning: Some(false), + request_headers: None, }; let switched = diff --git a/crates/goose-download-manager/Cargo.toml b/crates/goose-download-manager/Cargo.toml index c7ae701c3..39b2c29b1 100644 --- a/crates/goose-download-manager/Cargo.toml +++ b/crates/goose-download-manager/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "goose-download-manager" -version = "0.1.0-alpha.0" +version = "0.1.0-alpha.5" edition.workspace = true rust-version.workspace = true authors.workspace = true diff --git a/crates/goose-local-inference/Cargo.toml b/crates/goose-local-inference/Cargo.toml index 1d8befeb6..a4ecd640a 100644 --- a/crates/goose-local-inference/Cargo.toml +++ b/crates/goose-local-inference/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "goose-local-inference" -version = "0.1.0-alpha.0" +version = "0.1.0-alpha.5" edition.workspace = true rust-version.workspace = true authors.workspace = true @@ -27,9 +27,9 @@ etcetera = { workspace = true } encoding_rs = { version = "0.8.35", default-features = false } fs2 = { workspace = true } futures = { workspace = true } -goose-download-manager = { version = "0.1.0-alpha.0", path = "../goose-download-manager" } -goose-provider-types = { version = "0.1.0-alpha.0", path = "../goose-provider-types", default-features = false } -goose-sdk-types = { version = "0.1.0-alpha.0", path = "../goose-sdk-types", default-features = false } +goose-download-manager = { version = "0.1.0-alpha.5", path = "../goose-download-manager" } +goose-provider-types = { version = "0.1.0-alpha.5", path = "../goose-provider-types", default-features = false } +goose-sdk-types = { version = "0.1.0-alpha.5", path = "../goose-sdk-types", default-features = false } hf-hub = { version = "1.0.0-rc.1", default-features = false } include_dir = { workspace = true } llama-cpp-2 = { workspace = true } diff --git a/crates/goose-provider-types/Cargo.toml b/crates/goose-provider-types/Cargo.toml index bef8070b3..76b8b3821 100644 --- a/crates/goose-provider-types/Cargo.toml +++ b/crates/goose-provider-types/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "goose-provider-types" -version = "0.1.0-alpha.0" +version = "0.1.0-alpha.5" edition.workspace = true rust-version.workspace = true authors.workspace = true diff --git a/crates/goose-provider-types/src/base.rs b/crates/goose-provider-types/src/base.rs index 91e3c5e12..fd2e78a99 100644 --- a/crates/goose-provider-types/src/base.rs +++ b/crates/goose-provider-types/src/base.rs @@ -358,13 +358,21 @@ pub async fn collect_stream( } } - match final_message { - Some(msg) => { - let usage = final_usage + match (final_message, final_usage) { + (Some(msg), usage) => { + let usage = usage .unwrap_or_else(|| ProviderUsage::new("unknown".to_string(), Usage::default())); Ok((msg, usage)) } - None => Err(ProviderError::ExecutionError( + (None, Some(usage)) => Ok(( + Message::new( + rmcp::model::Role::Assistant, + chrono::Utc::now().timestamp(), + Vec::new(), + ), + usage, + )), + (None, None) => Err(ProviderError::ExecutionError( "Stream yielded no message".to_string(), )), } @@ -662,6 +670,22 @@ mod tests { assert_eq!(usage.model, "unknown"); } + #[tokio::test] + async fn test_collect_stream_usage_only_yields_empty_message() { + let usage = ProviderUsage::new("claude-sonnet-4".to_string(), Usage::default()); + let stream = futures::stream::once(async move { Ok((None, Some(usage))) }); + let (msg, usage) = collect_stream(Box::pin(stream)).await.unwrap(); + assert!(msg.content.is_empty()); + assert_eq!(usage.model, "claude-sonnet-4"); + } + + #[tokio::test] + async fn test_collect_stream_no_message_no_usage_errors() { + let stream = futures::stream::empty(); + let result = collect_stream(Box::pin(stream)).await; + assert!(matches!(result, Err(ProviderError::ExecutionError(_)))); + } + #[tokio::test] async fn test_collect_stream_preserves_text_audience_boundaries() { use futures::stream; diff --git a/crates/goose-provider-types/src/formats/databricks.rs b/crates/goose-provider-types/src/formats/databricks.rs index cfffabaa8..a3e0f7fbb 100644 --- a/crates/goose-provider-types/src/formats/databricks.rs +++ b/crates/goose-provider-types/src/formats/databricks.rs @@ -1071,6 +1071,7 @@ mod tests { toolshim_model: None, request_params: None, reasoning: None, + request_headers: None, }; let request = create_request(&model_config, "system", &[], &[], &ImageFormat::OpenAi)?; let obj = request.as_object().unwrap(); @@ -1105,6 +1106,7 @@ mod tests { toolshim_model: None, request_params: Some(params), reasoning: None, + request_headers: None, }; let request = create_request(&model_config, "system", &[], &[], &ImageFormat::OpenAi)?; assert_eq!(request["reasoning_effort"], "high"); @@ -1124,6 +1126,7 @@ mod tests { toolshim_model: None, request_params: Some(params), reasoning: None, + request_headers: None, }; let request = create_request(&model_config, "system", &[], &[], &ImageFormat::OpenAi)?; assert_eq!(request["reasoning_effort"], "low"); @@ -1144,6 +1147,7 @@ mod tests { toolshim_model: None, request_params: Some(params), reasoning: None, + request_headers: None, }; let request = create_request(&model_config, "system", &[], &[], &ImageFormat::OpenAi)?; assert_eq!(request["reasoning_effort"], "high"); @@ -1162,6 +1166,7 @@ mod tests { toolshim_model: None, request_params: None, reasoning: None, + request_headers: None, }; let request = create_request(&model_config, "system", &[], &[], &ImageFormat::OpenAi)?; assert_eq!(request["model"], "o3"); @@ -1180,6 +1185,7 @@ mod tests { toolshim_model: None, request_params: None, reasoning: None, + request_headers: None, }; let request = create_request(&model_config, "system", &[], &[], &ImageFormat::OpenAi)?; assert_eq!(request["model"], "o3"); @@ -1198,6 +1204,7 @@ mod tests { toolshim_model: None, request_params: None, reasoning: None, + request_headers: None, }; let request = create_request(&model_config, "system", &[], &[], &ImageFormat::OpenAi)?; assert_eq!(request["model"], "databricks-gpt-5.4"); @@ -1687,6 +1694,7 @@ mod tests { toolshim_model: None, request_params: None, reasoning: None, + request_headers: None, }; let messages = vec![ @@ -1739,6 +1747,7 @@ mod tests { toolshim_model: None, request_params: None, reasoning: None, + request_headers: None, }; let messages = vec![Message::user().with_text("Hello")]; diff --git a/crates/goose-provider-types/src/formats/openai_responses.rs b/crates/goose-provider-types/src/formats/openai_responses.rs index f8e0aadd0..6f555a1aa 100644 --- a/crates/goose-provider-types/src/formats/openai_responses.rs +++ b/crates/goose-provider-types/src/formats/openai_responses.rs @@ -1330,6 +1330,7 @@ mod tests { toolshim_model: None, request_params: None, reasoning: None, + request_headers: None, }; let messages = vec![ @@ -1421,6 +1422,7 @@ mod tests { toolshim_model: None, request_params: None, reasoning: None, + request_headers: None, }; let tool = Tool::new( @@ -1464,6 +1466,7 @@ mod tests { toolshim_model: None, request_params: None, reasoning: None, + request_headers: None, }; let result = @@ -1544,6 +1547,7 @@ mod tests { toolshim_model: None, request_params: None, reasoning: None, + request_headers: None, }; let result = @@ -1594,6 +1598,7 @@ mod tests { toolshim_model: None, request_params: None, reasoning: None, + request_headers: None, }; let result = create_responses_request(&model_config, "You are helpful.", &[], &[]).unwrap(); @@ -1619,6 +1624,7 @@ mod tests { serde_json::json!(true), )])), reasoning: None, + request_headers: None, }; let result = create_responses_request(&model_config, "", &[], &[]).unwrap(); @@ -1643,6 +1649,7 @@ mod tests { toolshim_model: None, request_params: None, reasoning: None, + request_headers: None, }; let result = @@ -1689,6 +1696,7 @@ mod tests { toolshim_model: None, request_params: None, reasoning: None, + request_headers: None, }; let result = create_responses_request(&model_config, "", &messages, &[]).unwrap(); @@ -1725,6 +1733,7 @@ mod tests { toolshim_model: None, request_params: None, reasoning: None, + request_headers: None, }; let result = create_responses_request(&model_config, "", &messages, &[]).unwrap(); @@ -1756,6 +1765,7 @@ mod tests { toolshim_model: None, request_params: None, reasoning: None, + request_headers: None, }; let result = create_responses_request(&model_config, "", &messages, &[]).unwrap(); @@ -1786,6 +1796,7 @@ mod tests { toolshim_model: None, request_params: None, reasoning: None, + request_headers: None, }; let result = create_responses_request(&model_config, "", &messages, &[]).unwrap(); @@ -1820,6 +1831,7 @@ mod tests { toolshim_model: None, request_params: None, reasoning: None, + request_headers: None, }; let result = create_responses_request(&model_config, "", &messages, &[]).unwrap(); @@ -1856,6 +1868,7 @@ mod tests { toolshim_model: None, request_params: None, reasoning: None, + request_headers: None, }; let result = create_responses_request(&model_config, "", &messages, &[]).unwrap(); @@ -1881,6 +1894,7 @@ mod tests { toolshim_model: None, request_params: None, reasoning: None, + request_headers: None, }; let result = create_responses_request(&model_config, "", &messages, &[]).unwrap(); @@ -1912,6 +1926,7 @@ mod tests { toolshim_model: None, request_params: None, reasoning: None, + request_headers: None, }; let result = create_responses_request(&model_config, "", &messages, &[]).unwrap(); @@ -1943,6 +1958,7 @@ mod tests { toolshim_model: None, request_params: None, reasoning: None, + request_headers: None, }; let result = create_responses_request(&model_config, "", &messages, &[]).unwrap(); @@ -2085,6 +2101,7 @@ mod tests { toolshim_model: None, request_params: None, reasoning: None, + request_headers: None, }; let result = create_responses_request(&model_config, "", &messages, &[]).unwrap(); @@ -2125,6 +2142,7 @@ mod tests { toolshim_model: None, request_params: None, reasoning: None, + request_headers: None, }; let result = create_responses_request(&model_config, "", &messages, &[]).unwrap(); @@ -2162,6 +2180,7 @@ mod tests { toolshim_model: None, request_params: None, reasoning: None, + request_headers: None, }; let result = create_responses_request(&model_config, "", &messages, &[]).unwrap(); @@ -2199,6 +2218,7 @@ mod tests { toolshim_model: None, request_params: None, reasoning: None, + request_headers: None, }; let result = create_responses_request(&model_config, "", &messages, &[]).unwrap(); diff --git a/crates/goose-provider-types/src/model.rs b/crates/goose-provider-types/src/model.rs index cc5bdb6a2..17367ab13 100644 --- a/crates/goose-provider-types/src/model.rs +++ b/crates/goose-provider-types/src/model.rs @@ -33,6 +33,10 @@ pub struct ModelConfig { pub request_params: Option>, #[serde(default, skip_serializing_if = "Option::is_none")] pub reasoning: Option, + /// Per-request HTTP headers attached to outgoing provider calls. + /// Never serialized into request bodies. + #[serde(skip)] + pub request_headers: Option>, } impl<'de> Deserialize<'de> for ModelConfig { @@ -64,6 +68,7 @@ impl<'de> Deserialize<'de> for ModelConfig { toolshim_model: raw.toolshim_model, request_params: raw.request_params, reasoning: raw.reasoning, + request_headers: None, }; config.normalize_effort_suffix(); Ok(config) @@ -81,6 +86,7 @@ impl ModelConfig { toolshim_model: None, request_params: None, reasoning: None, + request_headers: None, }; config.normalize_effort_suffix(); config @@ -162,6 +168,11 @@ impl ModelConfig { self } + pub fn with_request_headers(mut self, headers: Option>) -> Self { + self.request_headers = headers; + self + } + pub fn with_merged_request_params(mut self, params: HashMap) -> Self { match self.request_params.as_mut() { Some(existing) => { @@ -300,6 +311,25 @@ impl ModelConfig { mod tests { use super::*; + #[test] + fn request_headers_never_serialize_into_bodies() { + let config = ModelConfig::new("test-model").with_request_headers(Some(HashMap::from([( + "queue_threshold".to_string(), + "500".to_string(), + )]))); + + let serialized = serde_json::to_value(&config).unwrap(); + assert!(serialized.get("request_headers").is_none()); + assert_eq!( + config + .request_headers + .as_ref() + .unwrap() + .get("queue_threshold"), + Some(&"500".to_string()) + ); + } + mod thinking_effort_tests { use super::*; diff --git a/crates/goose-providers/Cargo.toml b/crates/goose-providers/Cargo.toml index 73d86b82a..29862de30 100644 --- a/crates/goose-providers/Cargo.toml +++ b/crates/goose-providers/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "goose-providers" -version = "0.1.0-alpha.0" +version = "0.1.0-alpha.5" edition.workspace = true rust-version.workspace = true authors.workspace = true @@ -35,8 +35,8 @@ anyhow = { workspace = true } async-stream = { workspace = true } chrono = { workspace = true } futures = { workspace = true } -goose-provider-types = { version = "0.1.0-alpha.0", path = "../goose-provider-types", default-features = false } -goose-local-inference = { version = "0.1.0-alpha.0", path = "../goose-local-inference", default-features = false, optional = true } +goose-provider-types = { version = "0.1.0-alpha.5", path = "../goose-provider-types", default-features = false } +goose-local-inference = { version = "0.1.0-alpha.5", path = "../goose-local-inference", default-features = false, optional = true } reqwest = { workspace = true } rmcp = { workspace = true, features = ["server", "macros"] } serde = { workspace = true } diff --git a/crates/goose-providers/src/anthropic.rs b/crates/goose-providers/src/anthropic.rs index 33755a686..04a520616 100644 --- a/crates/goose-providers/src/anthropic.rs +++ b/crates/goose-providers/src/anthropic.rs @@ -282,7 +282,10 @@ impl Provider for AnthropicProvider { let response = self .with_retry(|| async { - let request = self.api_client.request("v1/messages"); + let request = self + .api_client + .request("v1/messages") + .model_headers(model_config)?; let resp = request.response_post(&payload).await?; handle_status(resp).await }) diff --git a/crates/goose-providers/src/api_client.rs b/crates/goose-providers/src/api_client.rs index 34a6d7f02..57a44fc8c 100644 --- a/crates/goose-providers/src/api_client.rs +++ b/crates/goose-providers/src/api_client.rs @@ -408,6 +408,17 @@ impl<'a> ApiRequestBuilder<'a> { self } + /// Apply per-request headers from a model config, overriding any static + /// client headers on key collision. + pub fn model_headers(self, model_config: &crate::model::ModelConfig) -> Result { + match &model_config.request_headers { + Some(headers) => headers + .iter() + .try_fold(self, |builder, (key, value)| builder.header(key, value)), + None => Ok(self), + } + } + pub async fn api_post(self, payload: &Value) -> Result { let response = self.response_post(payload).await?; ApiResponse::from_response(response).await @@ -598,6 +609,67 @@ ShGoCNbfNS+COlPMRAujyDlATZcLs9p4tA== mod tests { use super::*; + #[test] + fn test_model_headers_applied_and_override_static_headers() { + let runtime = tokio::runtime::Runtime::new().unwrap(); + runtime.block_on(async { + let client = ApiClient::new_with_tls( + "http://localhost:8080".to_string(), + AuthMethod::NoAuth, + None, + ) + .unwrap() + .with_header("x-static", "static-value") + .unwrap() + .with_header("queue_threshold", "1000") + .unwrap(); + + let model_config = crate::model::ModelConfig::new("test-model").with_request_headers( + Some(std::collections::HashMap::from([ + ("queue_threshold".to_string(), "500".to_string()), + ("Idempotency-Key".to_string(), "abc-123".to_string()), + ])), + ); + + let request = client + .request("/test") + .model_headers(&model_config) + .unwrap() + .send_request(|url, client| client.get(url)) + .await + .unwrap(); + + let headers = request.build().unwrap().headers().clone(); + let get = |name: &str| { + headers + .get(name) + .and_then(|value| value.to_str().ok()) + .map(str::to_string) + }; + assert_eq!(get("queue_threshold"), Some("500".to_string())); + assert_eq!(get("Idempotency-Key"), Some("abc-123".to_string())); + }); + } + + #[test] + fn test_model_headers_rejects_invalid_header_name() { + let client = ApiClient::new_with_tls( + "http://localhost:8080".to_string(), + AuthMethod::NoAuth, + None, + ) + .unwrap(); + + let model_config = crate::model::ModelConfig::new("test-model").with_request_headers(Some( + std::collections::HashMap::from([("bad header name".to_string(), "value".to_string())]), + )); + + assert!(client + .request("/test") + .model_headers(&model_config) + .is_err()); + } + #[test] fn test_request_builder_decorator() { let runtime = tokio::runtime::Runtime::new().unwrap(); diff --git a/crates/goose-providers/src/databricks.rs b/crates/goose-providers/src/databricks.rs index b857b9137..0775635f3 100644 --- a/crates/goose-providers/src/databricks.rs +++ b/crates/goose-providers/src/databricks.rs @@ -595,7 +595,12 @@ impl Provider for DatabricksProvider { let response = self .with_retry(|| async { let payload_clone = payload.clone(); - let resp = self.api_client.response_post(&path, &payload_clone).await?; + let resp = self + .api_client + .request(&path) + .model_headers(model_config)? + .response_post(&payload_clone) + .await?; handle_status(resp).await }) .await @@ -655,7 +660,12 @@ impl Provider for DatabricksProvider { let mut log = start_log(model_config, &payload)?; let response = self .with_retry(|| async { - let resp = self.api_client.response_post(&path, &payload).await?; + let resp = self + .api_client + .request(&path) + .model_headers(model_config)? + .response_post(&payload) + .await?; if !resp.status().is_success() { let status = resp.status(); let url = sanitize_url(resp.url().as_str()); @@ -672,7 +682,12 @@ impl Provider for DatabricksProvider { Err(e) if e.to_string().contains("stream_options") => { payload.as_object_mut().unwrap().remove("stream_options"); self.with_retry(|| async { - let resp = self.api_client.response_post(&path, &payload).await?; + let resp = self + .api_client + .request(&path) + .model_headers(model_config)? + .response_post(&payload) + .await?; if !resp.status().is_success() { let status = resp.status(); let url = sanitize_url(resp.url().as_str()); diff --git a/crates/goose-providers/src/databricks_v2.rs b/crates/goose-providers/src/databricks_v2.rs index af3014059..8a2c1f8c4 100644 --- a/crates/goose-providers/src/databricks_v2.rs +++ b/crates/goose-providers/src/databricks_v2.rs @@ -205,7 +205,9 @@ impl DatabricksV2Provider { .with_retry(|| async { let resp = self .api_client - .response_post("ai-gateway/openai/v1/responses", &payload) + .request("ai-gateway/openai/v1/responses") + .model_headers(model_config)? + .response_post(&payload) .await?; handle_status(resp).await }) @@ -241,7 +243,9 @@ impl DatabricksV2Provider { .with_retry(|| async { let resp = self .api_client - .response_post("ai-gateway/mlflow/v1/chat/completions", &payload) + .request("ai-gateway/mlflow/v1/chat/completions") + .model_headers(model_config)? + .response_post(&payload) .await?; handle_status(resp).await }) @@ -275,7 +279,9 @@ impl DatabricksV2Provider { .with_retry(|| async { let resp = self .api_client - .response_post("ai-gateway/anthropic/v1/messages", &payload) + .request("ai-gateway/anthropic/v1/messages") + .model_headers(model_config)? + .response_post(&payload) .await?; handle_status(resp).await }) diff --git a/crates/goose-providers/src/google.rs b/crates/goose-providers/src/google.rs index 934a488cc..83e29353f 100644 --- a/crates/goose-providers/src/google.rs +++ b/crates/goose-providers/src/google.rs @@ -93,11 +93,19 @@ impl GoogleProvider { async fn post_stream( &self, - model_name: &str, + model_config: &ModelConfig, payload: &Value, ) -> Result { - let path = format!("v1beta/models/{}:streamGenerateContent?alt=sse", model_name); - let response = self.api_client.response_post(&path, payload).await?; + let path = format!( + "v1beta/models/{}:streamGenerateContent?alt=sse", + model_config.model_name + ); + let response = self + .api_client + .request(&path) + .model_headers(model_config)? + .response_post(payload) + .await?; handle_status(response).await } } @@ -181,7 +189,7 @@ impl Provider for GoogleProvider { let mut log = start_log(model_config, &payload)?; let response = self - .with_retry(|| async { self.post_stream(&model_config.model_name, &payload).await }) + .with_retry(|| async { self.post_stream(model_config, &payload).await }) .await .inspect_err(|e| { let _ = log.error(e); diff --git a/crates/goose-providers/src/ollama.rs b/crates/goose-providers/src/ollama.rs index f6ac65fcc..2a4cdbab8 100644 --- a/crates/goose-providers/src/ollama.rs +++ b/crates/goose-providers/src/ollama.rs @@ -426,7 +426,9 @@ impl Provider for OllamaProvider { .with_retry(|| async { let resp = self .api_client - .response_post("v1/chat/completions", &payload) + .request("v1/chat/completions") + .model_headers(model_config)? + .response_post(&payload) .await?; handle_status(resp).await }) diff --git a/crates/goose-providers/src/openai.rs b/crates/goose-providers/src/openai.rs index ac6dd4f0f..6b5234108 100644 --- a/crates/goose-providers/src/openai.rs +++ b/crates/goose-providers/src/openai.rs @@ -656,14 +656,13 @@ impl Provider for OpenAiProvider { let payload_clone = payload.clone(); let resp = self .api_client - .response_post( - &Self::map_base_path( - &self.base_path, - "responses", - OPEN_AI_DEFAULT_RESPONSES_PATH, - ), - &payload_clone, - ) + .request(&Self::map_base_path( + &self.base_path, + "responses", + OPEN_AI_DEFAULT_RESPONSES_PATH, + )) + .model_headers(model_config)? + .response_post(&payload_clone) .await?; handle_status(resp).await }) @@ -721,7 +720,9 @@ impl Provider for OpenAiProvider { .with_retry(|| async { let resp = self .api_client - .response_post(&self.base_path, &payload) + .request(&self.base_path) + .model_headers(model_config)? + .response_post(&payload) .await?; handle_status(resp).await }) diff --git a/crates/goose-providers/src/openai_compatible.rs b/crates/goose-providers/src/openai_compatible.rs index 78f96aed9..b9e4648a4 100644 --- a/crates/goose-providers/src/openai_compatible.rs +++ b/crates/goose-providers/src/openai_compatible.rs @@ -123,7 +123,9 @@ impl Provider for OpenAiCompatibleProvider { .with_retry(|| async { let resp = self .api_client - .response_post(&completions_path, &payload) + .request(&completions_path) + .model_headers(model_config)? + .response_post(&payload) .await?; handle_status(resp).await }) diff --git a/crates/goose-providers/src/snowflake.rs b/crates/goose-providers/src/snowflake.rs index 1ae0537d2..d30cb26d9 100644 --- a/crates/goose-providers/src/snowflake.rs +++ b/crates/goose-providers/src/snowflake.rs @@ -91,10 +91,16 @@ impl SnowflakeProvider { }) } - async fn post(&self, payload: &Value) -> Result { + async fn post( + &self, + model_config: &ModelConfig, + payload: &Value, + ) -> Result { let response = self .api_client - .response_post("api/v2/cortex/inference:complete", payload) + .request("api/v2/cortex/inference:complete") + .model_headers(model_config)? + .response_post(payload) .await?; let status = response.status(); @@ -317,7 +323,7 @@ impl Provider for SnowflakeProvider { let response = self .with_retry(|| async { let payload_clone = payload.clone(); - self.post(&payload_clone).await + self.post(model_config, &payload_clone).await }) .await?; diff --git a/crates/goose-sdk-types/Cargo.toml b/crates/goose-sdk-types/Cargo.toml index afda65dd8..95bcf5fd6 100644 --- a/crates/goose-sdk-types/Cargo.toml +++ b/crates/goose-sdk-types/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "goose-sdk-types" -version = "0.1.0-alpha.0" +version = "0.1.0-alpha.5" edition.workspace = true rust-version.workspace = true authors.workspace = true diff --git a/crates/goose-sdk/Cargo.toml b/crates/goose-sdk/Cargo.toml index 48cdfb7ca..d185719ca 100644 --- a/crates/goose-sdk/Cargo.toml +++ b/crates/goose-sdk/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "goose-sdk" -version = "0.1.0-alpha.0" +version = "0.1.0-alpha.5" edition.workspace = true rust-version.workspace = true authors.workspace = true @@ -27,20 +27,24 @@ uniffi = [ "dep:futures", "dep:serde_json", "dep:tokio", + "dep:rmcp", + "dep:base64", ] [dependencies] -goose-sdk-types = { version = "0.1.0-alpha.0", path = "../goose-sdk-types" } +goose-sdk-types = { version = "0.1.0-alpha.5", path = "../goose-sdk-types" } agent-client-protocol = { workspace = true, features = ["unstable"] } agent-client-protocol-schema = { workspace = true } uniffi = { version = "0.32", features = ["cli"], optional = true } thiserror = { version = "2", optional = true } -goose-providers = { version = "0.1.0-alpha.0", path = "../goose-providers", features = ["rustls-tls"], optional = true } +goose-providers = { version = "0.1.0-alpha.5", path = "../goose-providers", features = ["rustls-tls"], optional = true } futures = { workspace = true, optional = true } serde_json = { workspace = true, optional = true } tokio = { workspace = true, features = ["rt-multi-thread", "sync"], optional = true } anyhow = { workspace = true, optional = true } +rmcp = { workspace = true, optional = true } +base64 = { workspace = true, optional = true } [dev-dependencies] tokio = { workspace = true, features = ["rt-multi-thread", "macros", "process", "io-std", "io-util"] } diff --git a/crates/goose-sdk/README.md b/crates/goose-sdk/README.md index 1649f8eb5..0ec1418c5 100644 --- a/crates/goose-sdk/README.md +++ b/crates/goose-sdk/README.md @@ -4,9 +4,9 @@ The bindings layer for Goose. It houses the shared types used for both ACP and SDK access, and exposes a cross-language version of the Goose API. With `--features uniffi` the crate compiles to native bindings for Python and -Kotlin (namespace `goose` / `io.github.aaif_goose`). The UniFFI surface currently lets -callers construct declarative providers from JSON and stream provider -completions. +Kotlin (namespace `goose` / `io.github.aaif_goose`). The UniFFI surface lets +callers construct providers, stream provider completions, perform non-streaming +completion, and pass rich message/tool content across the FFI boundary. ```bash just python # build bindings + run examples/uniffi/provider.py diff --git a/crates/goose-sdk/examples/uniffi/kotlin/src/main/kotlin/Main.kt b/crates/goose-sdk/examples/uniffi/kotlin/src/main/kotlin/Main.kt index 7b2401d71..0bb9785f0 100644 --- a/crates/goose-sdk/examples/uniffi/kotlin/src/main/kotlin/Main.kt +++ b/crates/goose-sdk/examples/uniffi/kotlin/src/main/kotlin/Main.kt @@ -1,6 +1,8 @@ +import io.github.aaif_goose.MessageContent import io.github.aaif_goose.MessageRole import io.github.aaif_goose.ProviderMessage import io.github.aaif_goose.ProviderModelConfig +import io.github.aaif_goose.StreamChunk import io.github.aaif_goose.streamFlow import io.github.aaif_goose.providers.openai.defaultModel import io.github.aaif_goose.providers.openai.provider as openAiProvider @@ -17,7 +19,11 @@ fun main() = runBlocking { val messages = listOf( ProviderMessage( role = MessageRole.USER, - text = "What is the capital of France? Answer in one sentence.", + content = listOf( + MessageContent.Text( + text = "What is the capital of France? Answer in one sentence.", + ), + ), ), ) @@ -28,8 +34,12 @@ fun main() = runBlocking { messages, ) .collect { chunk -> - chunk.text?.let { print(it) } - chunk.usageJson?.let { println("\nusage: $it") } + when (chunk) { + is StreamChunk.TextChunk -> print(chunk.text) + is StreamChunk.EndChunk -> chunk.usage?.let { println("\nusage: $it") } + is StreamChunk.ErrorChunk -> System.err.println("\nerror: ${chunk.error.message}") + is StreamChunk.ToolChunk -> Unit + } } println() } diff --git a/crates/goose-sdk/examples/uniffi/provider.py b/crates/goose-sdk/examples/uniffi/provider.py index 87f9886b9..016973555 100755 --- a/crates/goose-sdk/examples/uniffi/provider.py +++ b/crates/goose-sdk/examples/uniffi/provider.py @@ -1,7 +1,6 @@ #!/usr/bin/env -S uv run --script """Goose SDK demo: build a declarative provider and stream a completion.""" import asyncio -import json import sys from pathlib import Path @@ -9,9 +8,11 @@ HERE = Path(__file__).resolve().parent sys.path.insert(0, str(HERE.parent.parent / "generated")) from goose import ( # noqa: E402 + MessageContent, MessageRole, ProviderMessage, ProviderModelConfig, + StreamChunk, declarative_provider_from_json, ) @@ -19,19 +20,26 @@ from goose import ( # noqa: E402 async def main() -> None: provider = declarative_provider_from_json((HERE.parent / "deepseek.json").read_text()) model = ProviderModelConfig(model_name="deepseek-v4-flash") - messages = [ProviderMessage(role=MessageRole.USER, text="what is the capital of France?")] + messages = [ + ProviderMessage( + role=MessageRole.USER, + content=[MessageContent.Text(text="what is the capital of France?")], + ) + ] stream = await provider.stream( model, "You are a knowledgable geography expert", messages, + [], ) - while chunk := await stream.next(): - if chunk.text: + while chunk := await stream.next_chunk(): + if isinstance(chunk, StreamChunk.TextChunk): print(chunk.text, end="") - if chunk.usage_json: - usage = json.loads(chunk.usage_json) - print(f"\nusage: {usage}") + elif isinstance(chunk, StreamChunk.EndChunk) and chunk.usage: + print(f"\nusage: {chunk.usage}") + elif isinstance(chunk, StreamChunk.ErrorChunk): + print(f"\nerror: {chunk.error.message}", file=sys.stderr) print() diff --git a/crates/goose-sdk/justfile b/crates/goose-sdk/justfile index fe49ac8d3..404d0b6db 100644 --- a/crates/goose-sdk/justfile +++ b/crates/goose-sdk/justfile @@ -77,9 +77,15 @@ maven-bindings profile="release": esac; \ crates/goose-sdk/scripts/prepare-maven-package.sh "$lib_path" +maven-clean: + rm -rf {{maven_dir}}/build {{maven_kotlin_dir}} {{maven_resources_dir}} + maven-package: (maven-bindings "release") cd {{maven_dir}} && ./gradlew --no-daemon publishToMavenLocal +maven-package-clean: maven-clean (maven-bindings "release") + cd {{maven_dir}} && ./gradlew --no-daemon publishToMavenLocal + maven-publish: (maven-bindings "release") cd {{maven_dir}} && ./gradlew --no-daemon publishAndReleaseToMavenCentral diff --git a/crates/goose-sdk/maven/README.md b/crates/goose-sdk/maven/README.md index 0f3a91312..80d51d17c 100644 --- a/crates/goose-sdk/maven/README.md +++ b/crates/goose-sdk/maven/README.md @@ -5,7 +5,10 @@ as the Maven artifact `io.github.aaif-goose:gdk`. The artifact version is read from `crates/goose-sdk/Cargo.toml`, so it stays in lockstep with the Rust crate version. The jar includes the generated Kotlin API -and native libraries under JNA platform resource directories. +and native libraries under JNA platform resource directories. Packaging supports +`darwin-aarch64`, `darwin-x86-64`, `linux-x86-64`, `linux-aarch64`, and +`win32-x86-64` resource prefixes; CI is responsible for assembling every native +library into the final published jar. Build locally from the repository root: diff --git a/crates/goose-sdk/maven/build.gradle.kts b/crates/goose-sdk/maven/build.gradle.kts index 97448236f..ce16ed8f9 100644 --- a/crates/goose-sdk/maven/build.gradle.kts +++ b/crates/goose-sdk/maven/build.gradle.kts @@ -42,10 +42,6 @@ tasks.jar { } } -tasks.withType().configureEach { - dependsOn(tasks.named("plainJavadocJar")) -} - mavenPublishing { publishToMavenCentral(automaticRelease = true) if (providers.gradleProperty("signingInMemoryKey").isPresent) { diff --git a/crates/goose-sdk/maven/src/support/kotlin/io/github/aaif_goose/ProviderFlow.kt b/crates/goose-sdk/maven/src/support/kotlin/io/github/aaif_goose/ProviderFlow.kt index d52b61929..2b683c87c 100644 --- a/crates/goose-sdk/maven/src/support/kotlin/io/github/aaif_goose/ProviderFlow.kt +++ b/crates/goose-sdk/maven/src/support/kotlin/io/github/aaif_goose/ProviderFlow.kt @@ -7,9 +7,16 @@ public fun Provider.streamFlow( model: ProviderModelConfig, system: String, messages: List, -): Flow = flow { - val stream = stream(model, system, messages) + tools: List = emptyList(), +): Flow = flow { + val stream = stream(model, system, messages, tools) while (true) { - emit(stream.next() ?: break) + emit(stream.nextChunk() ?: break) } } + +public suspend fun Provider.complete( + model: ProviderModelConfig, + system: String, + messages: List, +): ProviderCompletion = complete(model, system, messages, emptyList()) diff --git a/crates/goose-sdk/maven/src/support/kotlin/io/github/aaif_goose/providers/anthropic/Anthropic.kt b/crates/goose-sdk/maven/src/support/kotlin/io/github/aaif_goose/providers/anthropic/Anthropic.kt new file mode 100644 index 000000000..fadd95d2d --- /dev/null +++ b/crates/goose-sdk/maven/src/support/kotlin/io/github/aaif_goose/providers/anthropic/Anthropic.kt @@ -0,0 +1,9 @@ +package io.github.aaif_goose.providers.anthropic + +public fun provider( + apiKey: String, + baseUrl: String? = null, + betaHeaders: List = emptyList(), +): io.github.aaif_goose.Provider = io.github.aaif_goose.anthropicProvider(apiKey, baseUrl, betaHeaders) + +public fun defaultModel(): String = io.github.aaif_goose.anthropicDefaultModel() diff --git a/crates/goose-sdk/maven/src/support/kotlin/io/github/aaif_goose/providers/groq/Groq.kt b/crates/goose-sdk/maven/src/support/kotlin/io/github/aaif_goose/providers/groq/Groq.kt new file mode 100644 index 000000000..d52dfca5a --- /dev/null +++ b/crates/goose-sdk/maven/src/support/kotlin/io/github/aaif_goose/providers/groq/Groq.kt @@ -0,0 +1,5 @@ +package io.github.aaif_goose.providers.groq + +public fun provider(apiKey: String): io.github.aaif_goose.Provider = io.github.aaif_goose.groqProvider(apiKey) + +public fun defaultModel(): String = io.github.aaif_goose.groqDefaultModel() diff --git a/crates/goose-sdk/python/pyproject.toml b/crates/goose-sdk/python/pyproject.toml index c12f5b9b3..4a97beda7 100644 --- a/crates/goose-sdk/python/pyproject.toml +++ b/crates/goose-sdk/python/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "goose-sdk" -version = "0.1.0a0" +version = "0.1.0a5" description = "Python bindings for the Goose SDK" readme = "README.md" requires-python = ">=3.9" diff --git a/crates/goose-sdk/scripts/maven-resource-prefix.sh b/crates/goose-sdk/scripts/maven-resource-prefix.sh index fa8ece254..ef3ea1a14 100755 --- a/crates/goose-sdk/scripts/maven-resource-prefix.sh +++ b/crates/goose-sdk/scripts/maven-resource-prefix.sh @@ -5,6 +5,7 @@ case "$(uname -s)-$(uname -m)" in Darwin-arm64) echo "darwin-aarch64" ;; Darwin-x86_64) echo "darwin-x86-64" ;; Linux-x86_64) echo "linux-x86-64" ;; + Linux-aarch64|Linux-arm64) echo "linux-aarch64" ;; MINGW64_NT-*-x86_64|MSYS_NT-*-x86_64|CYGWIN_NT-*-x86_64) echo "win32-x86-64" ;; *) echo "unsupported platform: $(uname -s)-$(uname -m)" >&2; exit 1 ;; esac diff --git a/crates/goose-sdk/scripts/prepare-maven-package.sh b/crates/goose-sdk/scripts/prepare-maven-package.sh index f99f316a2..cae816a5f 100755 --- a/crates/goose-sdk/scripts/prepare-maven-package.sh +++ b/crates/goose-sdk/scripts/prepare-maven-package.sh @@ -10,7 +10,7 @@ native_lib="$1" resource_prefix="${2:-$(crates/goose-sdk/scripts/maven-resource-prefix.sh)}" case "$resource_prefix" in - darwin-aarch64|linux-x86-64|win32-x86-64) ;; + darwin-aarch64|darwin-x86-64|linux-x86-64|linux-aarch64|win32-x86-64) ;; *) echo "unsupported JNA resource prefix: $resource_prefix" >&2; exit 1 ;; esac diff --git a/crates/goose-sdk/src/bindings.rs b/crates/goose-sdk/src/bindings.rs index f6728344d..4852a3263 100644 --- a/crates/goose-sdk/src/bindings.rs +++ b/crates/goose-sdk/src/bindings.rs @@ -1,77 +1,369 @@ //! In-process uniffi bindings for the Goose SDK. //! -//! This is the API surface exposed to Python and Kotlin. It currently focuses -//! on declarative providers: consumers can construct a provider from JSON and -//! stream completions from it. +//! This is the API surface exposed to Python and Kotlin. It focuses on native +//! Goose providers and mirrors the provider message/tool/streaming model closely +//! enough for Kotlin agent frameworks to avoid JSON-only shims for common paths. -use std::{future::Future, sync::Arc, sync::OnceLock}; +use std::{collections::HashMap, future::Future, sync::Arc, sync::OnceLock, time::Duration}; +use base64::Engine as _; use futures::StreamExt; use goose_providers::{ + anthropic::AnthropicProviderBuilder, api_client::{ApiClient, AuthMethod}, base::{MessageStream, Provider as GooseProvider}, - conversation::message::Message, + conversation::{ + message::{Message, MessageContent as GooseMessageContent}, + token_usage::ProviderUsage, + }, databricks::DatabricksProvider as GooseDatabricksProvider, databricks_auth::DatabricksAuth, - declarative::EnvKeyResolver, + databricks_v2::DatabricksV2Provider as GooseDatabricksV2Provider, + declarative::{DeclarativeProviderConfig, EnvKeyResolver}, model::ModelConfig, openai::OpenAiProviderBuilder, + utils::sanitize_unicode_tags, }; +use rmcp::model::{ + CallToolRequestParams, CallToolResult, Content, ErrorCode, ErrorData, Role, Tool, +}; +use serde_json::Value; -/// Errors surfaced across the uniffi boundary. #[derive(Debug, thiserror::Error, uniffi::Error)] pub enum GooseError { - #[error("{0}")] - Generic(String), + #[error("Rate limit exceeded{retry_after_suffix}")] + RateLimited { + retry_after_ms: Option, + retry_after_suffix: String, + }, + #[error("Output token limit exceeded: {details}")] + OutputTokenLimitExceeded { details: String }, + #[error("Context length exceeded: {details}")] + ContextLengthExceeded { details: String }, + #[error("Authentication error: {details}")] + Authentication { details: String }, + #[error("Timeout: {details}")] + Timeout { details: String }, + #[error("Provider unavailable: {details}")] + ProviderUnavailable { details: String }, + #[error("{details}")] + Generic { details: String }, +} + +impl GooseError { + fn generic(error: impl ToString) -> Self { + Self::Generic { + details: error.to_string(), + } + } } impl From for GooseError { fn from(error: anyhow::Error) -> Self { - Self::Generic(error.to_string()) + Self::generic(error) } } impl From for GooseError { fn from(error: goose_providers::errors::ProviderError) -> Self { - Self::Generic(error.to_string()) + match error { + goose_providers::errors::ProviderError::Authentication(message) => { + Self::Authentication { details: message } + } + goose_providers::errors::ProviderError::ContextLengthExceeded(message) => { + Self::ContextLengthExceeded { details: message } + } + goose_providers::errors::ProviderError::RateLimitExceeded { retry_delay, .. } => { + let retry_after_ms = retry_delay.map(|delay| delay.as_millis() as u64); + let retry_after_suffix = retry_after_ms + .map(|ms| format!("; retry after {ms}ms")) + .unwrap_or_default(); + Self::RateLimited { + retry_after_ms, + retry_after_suffix, + } + } + goose_providers::errors::ProviderError::ServerError(message) + | goose_providers::errors::ProviderError::EndpointNotFound(message) + | goose_providers::errors::ProviderError::CreditsExhausted { + details: message, .. + } => Self::ProviderUnavailable { details: message }, + goose_providers::errors::ProviderError::NetworkError(message) + if is_timeout(&message) => + { + Self::Timeout { details: message } + } + goose_providers::errors::ProviderError::RequestFailed(message) + if is_timeout(&message) => + { + Self::Timeout { details: message } + } + goose_providers::errors::ProviderError::ExecutionError(message) + if is_output_token_limit(&message) => + { + Self::OutputTokenLimitExceeded { details: message } + } + other => Self::generic(other), + } } } impl From for GooseError { fn from(error: serde_json::Error) -> Self { - Self::Generic(error.to_string()) + Self::generic(error) } } +fn is_timeout(message: &str) -> bool { + let message = message.to_ascii_lowercase(); + message.contains("timed out") || message.contains("timeout") +} + +fn is_output_token_limit(message: &str) -> bool { + let message = message.to_ascii_lowercase(); + message.contains("output token") + || message.contains("max_tokens") + || message.contains("max tokens") +} + +/// Receives provider request logs as JSONL records. +/// +/// `start` returns an identifier that is passed to `write` for every record in +/// that request, allowing callers to keep concurrent request logs separate. +#[uniffi::export(callback_interface)] +pub trait RequestLogger: Send + Sync { + fn start(&self) -> Result; + fn write(&self, request_id: u64, record: String) -> Result<(), GooseError>; +} + +struct RequestLoggerAdapter { + logger: Arc, +} + +impl goose_providers::request_log::RequestLogger for RequestLoggerAdapter { + fn start( + &self, + ) -> Result< + Box, + Box, + > { + Ok(Box::new(RequestLogHandleAdapter { + request_id: self.logger.start()?, + logger: Arc::clone(&self.logger), + })) + } +} + +struct RequestLogHandleAdapter { + request_id: u64, + logger: Arc, +} + +impl goose_providers::request_log::RequestLogHandle for RequestLogHandleAdapter { + fn write(&mut self, record: &str) -> Result<(), Box> { + self.logger.write(self.request_id, record.to_string())?; + Ok(()) + } +} + +/// Installs the process-wide provider request logger. +/// +/// A logger can only be installed once for the lifetime of the process. +#[uniffi::export] +pub fn install_request_logger(logger: Box) -> Result<(), GooseError> { + goose_providers::request_log::install_logger(RequestLoggerAdapter { + logger: Arc::from(logger), + }) + .map_err(|error| GooseError::generic(error.to_string())) +} + /// A text message passed to a provider. #[derive(Debug, Clone, uniffi::Record)] pub struct ProviderMessage { pub role: MessageRole, - pub text: String, + pub content: Vec, } -/// Supported message roles for provider requests and streamed responses. #[derive(Debug, Clone, uniffi::Enum)] pub enum MessageRole { User, Assistant, + Tool, +} + +#[derive(Debug, Clone, uniffi::Enum)] +pub enum MessageContent { + Text { + text: String, + }, + Image { + mime_type: String, + data: Vec, + }, + ToolRequest { + id: String, + name: String, + arguments_json: String, + }, + ToolResult { + id: String, + success: bool, + content_json: String, + }, + Thinking { + thinking: String, + signature: String, + }, + RedactedThinking { + data: String, + }, } impl ProviderMessage { - fn to_goose_message(&self) -> Message { - match self.role { - MessageRole::User => Message::user().with_text(&self.text), - MessageRole::Assistant => Message::assistant().with_text(&self.text), + fn to_goose_message(&self) -> Result, GooseError> { + let role = match self.role { + MessageRole::User | MessageRole::Tool => Role::User, + MessageRole::Assistant => Role::Assistant, + }; + let mut message = Message::new(role, chrono_now(), Vec::new()); + for content in &self.content { + message = message.with_content(content.to_goose_content()?); + } + Ok(Some(message)) + } +} + +impl MessageContent { + fn to_goose_content(&self) -> Result { + match self { + MessageContent::Text { text } => { + Ok(GooseMessageContent::text(sanitize_unicode_tags(text))) + } + MessageContent::Image { mime_type, data } => Ok(GooseMessageContent::image( + base64::engine::general_purpose::STANDARD.encode(data), + mime_type.clone(), + )), + MessageContent::ToolRequest { + id, + name, + arguments_json, + } => { + let arguments = parse_json_object(arguments_json)?; + Ok(GooseMessageContent::tool_request( + id.clone(), + Ok(CallToolRequestParams::new(name.clone()).with_arguments(arguments)), + )) + } + MessageContent::ToolResult { + id, + success, + content_json, + } => { + let value: Value = serde_json::from_str(content_json)?; + let tool_result = if *success { + Ok(call_tool_result(value, false)) + } else { + Err(ErrorData::new( + ErrorCode::INTERNAL_ERROR, + value.to_string(), + None, + )) + }; + Ok(GooseMessageContent::tool_response(id.clone(), tool_result)) + } + MessageContent::Thinking { + thinking, + signature, + } => Ok(GooseMessageContent::thinking( + thinking.clone(), + signature.clone(), + )), + MessageContent::RedactedThinking { data } => { + Ok(GooseMessageContent::redacted_thinking(data.clone())) + } } } } -/// Model selection and optional generation settings for a provider request. +fn chrono_now() -> i64 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|duration| duration.as_secs() as i64) + .unwrap_or_default() +} + +fn call_tool_result(value: Value, is_error: bool) -> CallToolResult { + let content = match value { + Value::Array(items) => items.into_iter().map(value_to_content).collect(), + other => vec![value_to_content(other)], + }; + + if is_error { + CallToolResult::error(content) + } else { + CallToolResult::success(content) + } +} + +fn value_to_content(value: Value) -> Content { + match value { + Value::String(text) => Content::text(text), + Value::Object(object) => match object.get("type").and_then(|value| value.as_str()) { + Some("text") => object + .get("text") + .and_then(|value| value.as_str().map(str::to_owned)) + .map(Content::text) + .unwrap_or_else(|| Content::text(Value::Object(object).to_string())), + Some("image") => { + let mime_type = object + .get("mimeType") + .or_else(|| object.get("mime_type")) + .and_then(|value| value.as_str().map(str::to_owned)) + .unwrap_or_else(|| "image/png".to_string()); + let data = object + .get("data") + .and_then(|value| value.as_str().map(str::to_owned)) + .unwrap_or_default(); + Content::image(data, mime_type) + } + _ => Content::text(Value::Object(object).to_string()), + }, + other => Content::text(other.to_string()), + } +} + +#[derive(Debug, Clone, uniffi::Record)] +pub struct ProviderTool { + pub name: String, + pub description: String, + pub input_schema_json: String, + #[uniffi(default = None)] + pub annotations_json: Option, +} + +impl ProviderTool { + fn to_goose_tool(&self) -> Result { + let schema = parse_json_object(&self.input_schema_json)?; + let mut tool = Tool::new(self.name.clone(), self.description.clone(), schema); + if let Some(annotations_json) = &self.annotations_json { + tool.annotations = Some(serde_json::from_str(annotations_json)?); + } + Ok(tool) + } +} + +fn parse_json_object(json: &str) -> Result, GooseError> { + match serde_json::from_str(json)? { + Value::Object(object) => Ok(object), + _ => Err(GooseError::generic("expected a JSON object")), + } +} + #[derive(Debug, Clone, uniffi::Record)] pub struct ProviderModelConfig { pub model_name: String, #[uniffi(default = None)] - pub context_limit: Option, + pub context_limit: Option, #[uniffi(default = None)] pub temperature: Option, #[uniffi(default = None)] @@ -80,41 +372,182 @@ pub struct ProviderModelConfig { pub toolshim: bool, #[uniffi(default = None)] pub toolshim_model: Option, - /// Provider-specific request parameters as a JSON object string. #[uniffi(default = None)] pub request_params_json: Option, #[uniffi(default = None)] + pub provider_params_json: Option, + #[uniffi(default = None)] pub reasoning: Option, + #[uniffi(default = None)] + pub timeout_ms: Option, + /// Per-request HTTP headers attached to the outgoing provider call. + /// These override any static headers configured on the provider. + #[uniffi(default = None)] + pub request_headers: Option>, } impl ProviderModelConfig { fn to_goose_model_config(&self) -> Result { let mut config = ModelConfig::new(&self.model_name) - .with_context_limit(self.context_limit.map(|limit| limit as usize)) + .with_context_limit(self.context_limit.map(|limit| limit.max(0) as usize)) .with_temperature(self.temperature) .with_max_tokens(self.max_tokens) .with_toolshim(self.toolshim) .with_toolshim_model(self.toolshim_model.clone()); - if let Some(request_params_json) = &self.request_params_json { - let request_params = serde_json::from_str(request_params_json)?; - config = config.with_merged_request_params(request_params); + let mut request_params = serde_json::Map::new(); + merge_params(&mut request_params, self.request_params_json.as_ref())?; + merge_params(&mut request_params, self.provider_params_json.as_ref())?; + if !request_params.is_empty() { + config = config.with_merged_request_params(request_params.into_iter().collect()); } + config = config.with_request_headers(self.request_headers.clone()); config.reasoning = self.reasoning; Ok(config) } } -/// One item yielded by a provider stream. +fn merge_params( + target: &mut serde_json::Map, + params_json: Option<&String>, +) -> Result<(), GooseError> { + if let Some(params_json) = params_json { + for (key, value) in parse_json_object(params_json)? { + target.insert(key, value); + } + } + Ok(()) +} + #[derive(Debug, Clone, uniffi::Record)] -pub struct ProviderStreamChunk { - /// The concatenated text content in this message chunk, if one was emitted. - pub text: Option, - /// Full Goose message JSON for callers that need non-text content such as tool requests. - pub message_json: Option, - /// Provider usage JSON when the provider emits usage metadata. - pub usage_json: Option, +pub struct Usage { + pub input_tokens: Option, + pub output_tokens: Option, + pub total_tokens: Option, + pub cache_read_input_tokens: Option, + pub cache_creation_input_tokens: Option, + pub reasoning_tokens: Option, + pub model: String, + pub provider_metadata_json: Option, +} + +impl Usage { + fn from_provider_usage(usage: &ProviderUsage) -> Result { + Ok(Self { + input_tokens: usage.usage.input_tokens, + output_tokens: usage.usage.output_tokens, + total_tokens: usage.usage.total_tokens, + cache_read_input_tokens: usage.usage.cache_read_input_tokens, + cache_creation_input_tokens: usage.usage.cache_write_input_tokens, + reasoning_tokens: None, + model: usage.model.clone(), + provider_metadata_json: Some(serde_json::to_string(usage)?), + }) + } +} + +#[derive(Debug, Clone, uniffi::Enum)] +pub enum StreamChunk { + TextChunk { + text: String, + }, + ToolChunk { + id: String, + name: String, + arguments_json: String, + }, + ThinkingChunk { + thinking: String, + signature: String, + }, + RedactedThinkingChunk { + data: String, + }, + EndChunk { + usage: Option, + }, + ErrorChunk { + error: GooseStreamError, + }, +} + +#[derive(Debug, Clone, uniffi::Record)] +pub struct GooseStreamError { + pub kind: GooseStreamErrorKind, + pub message: String, + pub retry_after_ms: Option, +} + +#[derive(Debug, Clone, uniffi::Enum)] +pub enum GooseStreamErrorKind { + RateLimited, + OutputTokenLimitExceeded, + ContextLengthExceeded, + Authentication, + Timeout, + ProviderUnavailable, + Generic, +} + +impl From for GooseStreamError { + fn from(error: GooseError) -> Self { + match error { + GooseError::RateLimited { + retry_after_ms, + retry_after_suffix, + } => Self { + kind: GooseStreamErrorKind::RateLimited, + message: format!("Rate limit exceeded{retry_after_suffix}"), + retry_after_ms, + }, + GooseError::OutputTokenLimitExceeded { details } => Self { + kind: GooseStreamErrorKind::OutputTokenLimitExceeded, + message: details, + retry_after_ms: None, + }, + GooseError::ContextLengthExceeded { details } => Self { + kind: GooseStreamErrorKind::ContextLengthExceeded, + message: details, + retry_after_ms: None, + }, + GooseError::Authentication { details } => Self { + kind: GooseStreamErrorKind::Authentication, + message: details, + retry_after_ms: None, + }, + GooseError::Timeout { details } => Self { + kind: GooseStreamErrorKind::Timeout, + message: details, + retry_after_ms: None, + }, + GooseError::ProviderUnavailable { details } => Self { + kind: GooseStreamErrorKind::ProviderUnavailable, + message: details, + retry_after_ms: None, + }, + GooseError::Generic { details } => Self { + kind: GooseStreamErrorKind::Generic, + message: details, + retry_after_ms: None, + }, + } + } +} + +#[derive(Debug, Clone, uniffi::Record)] +pub struct ProviderCompletion { + pub message_json: String, + pub usage: Option, +} + +#[derive(Debug, Clone, uniffi::Enum)] +pub enum Feature { + Tools, + Streaming, + Images, + JsonSchema, + Reasoning, } static RUNTIME: OnceLock = OnceLock::new(); @@ -127,26 +560,44 @@ fn runtime() -> Result<&'static tokio::runtime::Runtime, GooseError> { let runtime = tokio::runtime::Builder::new_multi_thread() .enable_all() .build() - .map_err(|error| GooseError::Generic(error.to_string()))?; + .map_err(GooseError::generic)?; - let _ = RUNTIME.set(runtime); + if let Err(unused) = RUNTIME.set(runtime) { + // Lost the init race; dropping a runtime inside an async context + // panics, so shut it down without blocking. + unused.shutdown_background(); + } Ok(RUNTIME.get().expect("runtime was initialized")) } +struct AbortOnDrop { + handle: tokio::task::JoinHandle, +} + +impl Drop for AbortOnDrop { + fn drop(&mut self) { + self.handle.abort(); + } +} + async fn run_on_runtime( future: impl Future + Send + 'static, ) -> Result where T: Send + 'static, { - let (sender, receiver) = tokio::sync::oneshot::channel(); - runtime()?.spawn(async move { - let _ = sender.send(future.await); - }); - - receiver - .await - .map_err(|_| GooseError::Generic("runtime task was cancelled".to_string())) + let mut task = AbortOnDrop { + handle: runtime()?.spawn(future), + }; + (&mut task.handle).await.map_err(|error| { + if error.is_cancelled() { + GooseError::Timeout { + details: "runtime task was cancelled".to_string(), + } + } else { + GooseError::generic(error) + } + }) } struct ProviderHandle { @@ -169,24 +620,85 @@ impl ProviderHandle { model: ProviderModelConfig, system: String, messages: Vec, + tools: Vec, ) -> Result, GooseError> { + let timeout_ms = model.timeout_ms; let model = model.to_goose_model_config()?; - let messages = messages - .iter() - .map(ProviderMessage::to_goose_message) - .collect::>(); + let messages = convert_messages(messages)?; + let tools = convert_tools(tools)?; let provider = Arc::clone(&self.provider); - let stream = - run_on_runtime(async move { provider.stream(&model, &system, &messages, &[]).await }) - .await??; + let stream = run_provider_future(timeout_ms, async move { + provider.stream(&model, &system, &messages, &tools).await + }) + .await??; Ok(Arc::new(ProviderStream { - stream: Arc::new(tokio::sync::Mutex::new(stream)), + state: Arc::new(tokio::sync::Mutex::new(ProviderStreamState { + stream, + pending: Vec::new(), + final_usage: None, + ended: false, + })), + timeout_ms, })) } + + async fn complete( + &self, + model: ProviderModelConfig, + system: String, + messages: Vec, + tools: Vec, + ) -> Result { + let timeout_ms = model.timeout_ms; + let model = model.to_goose_model_config()?; + let messages = convert_messages(messages)?; + let tools = convert_tools(tools)?; + let provider = Arc::clone(&self.provider); + let (message, usage) = run_provider_future(timeout_ms, async move { + provider.complete(&model, &system, &messages, &tools).await + }) + .await??; + + Ok(ProviderCompletion { + message_json: serde_json::to_string(&message)?, + usage: Some(Usage::from_provider_usage(&usage)?), + }) + } +} + +async fn run_provider_future( + timeout_ms: Option, + future: impl Future + Send + 'static, +) -> Result +where + T: Send + 'static, +{ + run_on_runtime(async move { + if let Some(timeout_ms) = timeout_ms { + tokio::time::timeout(Duration::from_millis(timeout_ms), future) + .await + .map_err(|_| GooseError::Timeout { + details: format!("request timed out after {timeout_ms}ms"), + }) + } else { + Ok(future.await) + } + }) + .await? +} + +fn convert_messages(messages: Vec) -> Result, GooseError> { + messages + .iter() + .filter_map(|message| message.to_goose_message().transpose()) + .collect() +} + +fn convert_tools(tools: Vec) -> Result, GooseError> { + tools.iter().map(ProviderTool::to_goose_tool).collect() } -/// A Goose provider backed by one of Goose's native provider implementations. #[derive(uniffi::Object)] pub struct Provider { handle: ProviderHandle, @@ -206,15 +718,42 @@ impl Provider { self.handle.name() } - /// Start a streaming completion request. Tools are not yet exposed over the - /// uniffi boundary, so this calls providers with an empty tool list. + pub fn supported_features(&self) -> Vec { + let name = self.name(); + let mut features = vec![Feature::Streaming, Feature::Tools, Feature::JsonSchema]; + if matches!( + name.as_str(), + "openai" | "anthropic" | "databricks" | "databricks_v2" | "groq" + ) { + features.push(Feature::Images); + } + if matches!( + name.as_str(), + "openai" | "anthropic" | "databricks" | "databricks_v2" + ) { + features.push(Feature::Reasoning); + } + features + } + pub async fn stream( &self, model: ProviderModelConfig, system: String, messages: Vec, + tools: Vec, ) -> Result, GooseError> { - self.handle.stream(model, system, messages).await + self.handle.stream(model, system, messages, tools).await + } + + pub async fn complete( + &self, + model: ProviderModelConfig, + system: String, + messages: Vec, + tools: Vec, + ) -> Result { + self.handle.complete(model, system, messages, tools).await } } @@ -240,6 +779,83 @@ pub fn openai_provider(api_key: String) -> Result, GooseError> { Ok(Provider::new(Box::new(provider))) } +#[uniffi::export] +pub fn anthropic_default_model() -> String { + goose_providers::anthropic::ANTHROPIC_DEFAULT_MODEL.to_string() +} + +#[uniffi::export] +pub fn anthropic_provider( + api_key: String, + base_url: Option, + beta_headers: Vec, +) -> Result, GooseError> { + let mut api_client = ApiClient::new_with_tls( + base_url.unwrap_or_else(|| "https://api.anthropic.com".to_string()), + AuthMethod::ApiKey { + header_name: "x-api-key".to_string(), + key: api_key, + }, + None, + )? + .with_header( + "anthropic-version", + goose_providers::anthropic::ANTHROPIC_API_VERSION, + )?; + + if !beta_headers.is_empty() { + api_client = api_client.with_header("anthropic-beta", &beta_headers.join(","))?; + } + + let provider = AnthropicProviderBuilder::new(api_client).build(); + Ok(Provider::new(Box::new(provider))) +} + +#[uniffi::export] +pub fn groq_default_model() -> String { + groq_config() + .models + .first() + .map(|model| model.name.clone()) + .unwrap_or_else(|| "moonshotai/kimi-k2-instruct-0905".to_string()) +} + +#[uniffi::export] +pub fn groq_provider(api_key: String) -> Result, GooseError> { + let json = goose_providers::groq::JSON.replace("${GROQ_API_KEY}", &api_key); + let provider = goose_providers::declarative::from_json( + &json, + None, + StaticKeyResolver { + key_name: "GROQ_API_KEY".to_string(), + key: api_key, + }, + )?; + Ok(Provider::new(provider)) +} + +fn groq_config() -> DeclarativeProviderConfig { + goose_providers::declarative::deserialize_provider_config(goose_providers::groq::JSON) + .expect("bundled groq provider config is valid") +} + +struct StaticKeyResolver { + key_name: String, + key: String, +} + +impl goose_providers::declarative::KeyResolver for StaticKeyResolver { + type Error = std::env::VarError; + + fn resolve_key(&self, key: &str) -> std::result::Result { + if key == self.key_name { + Ok(self.key.clone()) + } else { + std::env::var(key) + } + } +} + #[uniffi::export] pub fn databricks_default_model() -> String { goose_providers::databricks::DATABRICKS_DEFAULT_MODEL.to_string() @@ -264,52 +880,167 @@ pub fn databricks_provider(host: String, token: String) -> Result, Ok(Provider::new(Box::new(provider))) } -/// An async iterator over provider stream chunks. +#[uniffi::export] +pub fn databricks_v2_default_model() -> String { + goose_providers::databricks_v2::DATABRICKS_V2_DEFAULT_MODEL.to_string() +} + +#[uniffi::export] +pub fn databricks_v2_provider(host: String, token: String) -> Result, GooseError> { + let retry_config = GooseDatabricksV2Provider::load_retry_config(|key| std::env::var(key).ok()); + let provider = GooseDatabricksV2Provider::new( + host, + DatabricksAuth::token(token), + retry_config, + None, + None, + None, + None, + None, + )?; + + Ok(Provider::new(Box::new(provider))) +} + #[derive(uniffi::Object)] pub struct ProviderStream { - stream: Arc>, + state: Arc>, + timeout_ms: Option, +} + +struct ProviderStreamState { + stream: MessageStream, + pending: Vec, + final_usage: Option, + ended: bool, } #[uniffi::export] impl ProviderStream { - /// Return the next stream chunk, or `None` when the stream is exhausted. - pub async fn next(&self) -> Result, GooseError> { - let stream = Arc::clone(&self.stream); + pub async fn next_chunk(&self) -> Result, GooseError> { + let state = Arc::clone(&self.state); + let timeout_ms = self.timeout_ms; run_on_runtime(async move { - let mut stream = stream.lock().await; - let Some((message, usage)) = stream.next().await.transpose()? else { - return Ok(None); - }; + let mut state = state.lock().await; + loop { + if let Some(chunk) = state.pending.pop() { + return Ok(Some(chunk)); + } - let text = message.as_ref().map(Message::as_concat_text); - let message_json = message.as_ref().map(serde_json::to_string).transpose()?; - let usage_json = usage.as_ref().map(serde_json::to_string).transpose()?; + if state.ended { + return Ok(None); + } - Ok(Some(ProviderStreamChunk { - text, - message_json, - usage_json, - })) + let next = if let Some(timeout_ms) = timeout_ms { + tokio::time::timeout(Duration::from_millis(timeout_ms), state.stream.next()) + .await + .map_err(|_| GooseError::Timeout { + details: format!("request timed out after {timeout_ms}ms"), + })? + } else { + state.stream.next().await + }; + + match next { + Some(Ok((message, usage))) => { + if let Some(usage) = usage { + state.final_usage = Some(Usage::from_provider_usage(&usage)?); + } + let Some(message) = message else { + continue; + }; + let mut chunks = message_to_chunks(message); + if chunks.is_empty() { + continue; + } + chunks.reverse(); + let first = chunks.pop(); + state.pending = chunks; + return Ok(first); + } + Some(Err(error)) => { + state.ended = true; + return Ok(Some(StreamChunk::ErrorChunk { + error: GooseStreamError::from(GooseError::from(error)), + })); + } + None => { + state.ended = true; + return Ok(Some(StreamChunk::EndChunk { + usage: state.final_usage.clone(), + })); + } + } + } }) .await? } } +fn message_to_chunks(message: Message) -> Vec { + message + .content + .into_iter() + .filter_map(|content| match content { + GooseMessageContent::Text(text) if !text.text.is_empty() => { + Some(StreamChunk::TextChunk { + text: text.text.clone(), + }) + } + GooseMessageContent::ToolRequest(request) => match request.tool_call { + Ok(tool_call) => Some(StreamChunk::ToolChunk { + id: request.id, + name: tool_call.name.to_string(), + arguments_json: serde_json::to_string(&tool_call.arguments.unwrap_or_default()) + .unwrap_or_else(|_| "{}".to_string()), + }), + Err(error) => Some(StreamChunk::ErrorChunk { + error: GooseStreamError { + kind: GooseStreamErrorKind::Generic, + message: error.to_string(), + retry_after_ms: None, + }, + }), + }, + GooseMessageContent::Thinking(thinking) => Some(StreamChunk::ThinkingChunk { + thinking: thinking.thinking, + signature: thinking.signature, + }), + GooseMessageContent::RedactedThinking(redacted) => { + Some(StreamChunk::RedactedThinkingChunk { + data: redacted.data, + }) + } + _ => None, + }) + .collect() +} + #[cfg(test)] mod tests { use super::*; - #[test] - fn model_config_rejects_invalid_request_params_json() { - let config = ProviderModelConfig { + fn base_model_config() -> ProviderModelConfig { + ProviderModelConfig { model_name: "test".to_string(), context_limit: None, temperature: None, max_tokens: None, toolshim: false, toolshim_model: None, - request_params_json: Some("not json".to_string()), + request_params_json: None, + provider_params_json: None, reasoning: None, + timeout_ms: None, + request_headers: None, + } + } + + #[test] + fn model_config_rejects_invalid_request_params_json() { + let config = ProviderModelConfig { + request_params_json: Some("not json".to_string()), + ..base_model_config() }; assert!(config.to_goose_model_config().is_err()); @@ -319,10 +1050,48 @@ mod tests { fn provider_message_converts_user_text() { let message = ProviderMessage { role: MessageRole::User, - text: "what is the capital of France?".to_string(), + content: vec![MessageContent::Text { + text: "what is the capital of France?".to_string(), + }], } - .to_goose_message(); + .to_goose_message() + .unwrap() + .unwrap(); assert_eq!(message.as_concat_text(), "what is the capital of France?"); } + + #[test] + fn tool_config_converts_to_rmcp_tool() { + let tool = ProviderTool { + name: "lookup".to_string(), + description: "Lookup a value".to_string(), + input_schema_json: r#"{"type":"object","properties":{"key":{"type":"string"}}}"# + .to_string(), + annotations_json: None, + } + .to_goose_tool() + .unwrap(); + + assert_eq!(tool.name.as_ref(), "lookup"); + assert_eq!(tool.input_schema["type"], "object"); + } + + #[test] + fn tool_result_content_converts() { + let content = MessageContent::ToolResult { + id: "call_1".to_string(), + success: true, + content_json: r#"{"type":"text","text":"done"}"#.to_string(), + } + .to_goose_content() + .unwrap(); + + let GooseMessageContent::ToolResponse(response) = content else { + panic!("expected tool response"); + }; + let result = response.tool_result.unwrap(); + assert_eq!(result.is_error, Some(false)); + assert_eq!(result.content[0].as_text().unwrap().text, "done"); + } } diff --git a/crates/goose/src/context_mgmt/mod.rs b/crates/goose/src/context_mgmt/mod.rs index 1f0d074e3..d4e7fd7f8 100644 --- a/crates/goose/src/context_mgmt/mod.rs +++ b/crates/goose/src/context_mgmt/mod.rs @@ -710,6 +710,7 @@ mod tests { toolshim_model: None, request_params: None, reasoning: None, + request_headers: None, }, max_tool_responses: None, } diff --git a/crates/goose/src/model_config.rs b/crates/goose/src/model_config.rs index 145e6ac77..abd1584b9 100644 --- a/crates/goose/src/model_config.rs +++ b/crates/goose/src/model_config.rs @@ -98,6 +98,7 @@ pub async fn get_fast_model( match fast_model_name { Some(name) if name != model_config.model_name => { model_config_from_user_config(provider_name, name) + .map(|config| config.with_request_headers(model_config.request_headers.clone())) } _ => Ok(model_config.clone()), } @@ -179,6 +180,7 @@ fn base_model_config_from_user_config(model_name: &str) -> Result { toolshim_model: get_goose_toolshim_model(config)?, request_params: None, reasoning: None, + request_headers: None, }; model.normalize_effort_suffix(); Ok(model) diff --git a/crates/goose/src/providers/bedrock.rs b/crates/goose/src/providers/bedrock.rs index ca1759168..d6fee4cec 100644 --- a/crates/goose/src/providers/bedrock.rs +++ b/crates/goose/src/providers/bedrock.rs @@ -941,6 +941,7 @@ mod tests { toolshim_model: None, request_params: None, reasoning: None, + request_headers: None, }, ) } diff --git a/crates/goose/src/providers/gcpvertexai.rs b/crates/goose/src/providers/gcpvertexai.rs index 1230d41dc..6a2ea2ebc 100644 --- a/crates/goose/src/providers/gcpvertexai.rs +++ b/crates/goose/src/providers/gcpvertexai.rs @@ -279,6 +279,7 @@ impl GcpVertexAIProvider { async fn send_request_with_retry( &self, + model: &ModelConfig, url: Url, payload: &Value, ) -> Result { @@ -314,12 +315,18 @@ impl GcpVertexAIProvider { } }; - let request = self + let mut request = self .client .post(url.clone()) .json(payload) .header("Authorization", auth_header); + if let Some(headers) = &model.request_headers { + for (key, value) in headers { + request = request.header(key, value); + } + } + let response = (self.request_builder)(request) .map_err(|e| ProviderError::ExecutionError(e.to_string()))? .send() @@ -400,7 +407,7 @@ impl GcpVertexAIProvider { .build_request_url(model, context.provider(), location, true) .map_err(|e| ProviderError::RequestFailed(e.to_string()))?; - self.send_request_with_retry(url, payload).await + self.send_request_with_retry(model, url, payload).await } async fn post_stream( diff --git a/crates/goose/src/providers/githubcopilot.rs b/crates/goose/src/providers/githubcopilot.rs index b69c9c048..cc218215f 100644 --- a/crates/goose/src/providers/githubcopilot.rs +++ b/crates/goose/src/providers/githubcopilot.rs @@ -261,6 +261,7 @@ impl GithubCopilotProvider { async fn post( &self, + model_config: &ModelConfig, path: &str, is_user_initiated: bool, payload: &mut Value, @@ -279,7 +280,9 @@ impl GithubCopilotProvider { .with_headers(headers)?; api_client - .response_post(path, payload) + .request(path) + .model_headers(model_config)? + .response_post(payload) .await .map_err(|e| e.into()) } @@ -413,6 +416,7 @@ impl GithubCopilotProvider { let mut payload_clone = payload.clone(); let resp = self .post( + model_config, "responses", is_user_initiated, &mut payload_clone, @@ -459,6 +463,7 @@ impl GithubCopilotProvider { let mut payload_clone = payload.clone(); let resp = self .post( + model_config, "chat/completions", is_user_initiated, &mut payload_clone, @@ -488,6 +493,7 @@ impl GithubCopilotProvider { .with_retry(|| async { let mut payload_clone = payload.clone(); self.post( + model_config, "chat/completions", is_user_initiated, &mut payload_clone, diff --git a/crates/goose/src/providers/litellm.rs b/crates/goose/src/providers/litellm.rs index 87b5ceb27..e9caf70f5 100644 --- a/crates/goose/src/providers/litellm.rs +++ b/crates/goose/src/providers/litellm.rs @@ -136,10 +136,16 @@ impl LiteLLMProvider { Ok(models) } - async fn post(&self, payload: &Value) -> Result { + async fn post( + &self, + model_config: &ModelConfig, + payload: &Value, + ) -> Result { let response = self .api_client - .response_post(&self.base_path, payload) + .request(&self.base_path) + .model_headers(model_config)? + .response_post(payload) .await?; handle_response_openai_compat(response).await } @@ -248,7 +254,7 @@ impl Provider for LiteLLMProvider { let response = self .with_retry(|| async { let payload_clone = payload.clone(); - self.post(&payload_clone).await + self.post(model_config, &payload_clone).await }) .await?; diff --git a/crates/goose/src/providers/nanogpt.rs b/crates/goose/src/providers/nanogpt.rs index 56829b472..df7fcb57f 100644 --- a/crates/goose/src/providers/nanogpt.rs +++ b/crates/goose/src/providers/nanogpt.rs @@ -196,7 +196,9 @@ impl Provider for NanoGptProvider { .with_retry(|| async { let resp = self .api_client - .response_post("chat/completions", &payload) + .request("chat/completions") + .model_headers(model_config)? + .response_post(&payload) .await?; handle_status(resp).await }) diff --git a/crates/goose/src/providers/openrouter.rs b/crates/goose/src/providers/openrouter.rs index 92249bd10..d9d5088da 100644 --- a/crates/goose/src/providers/openrouter.rs +++ b/crates/goose/src/providers/openrouter.rs @@ -335,7 +335,9 @@ impl Provider for OpenRouterProvider { .with_retry(|| async { let resp = self .api_client - .response_post("api/v1/chat/completions", &payload) + .request("api/v1/chat/completions") + .model_headers(model_config)? + .response_post(&payload) .await?; handle_status(resp).await }) @@ -369,6 +371,7 @@ mod tests { toolshim_model: None, request_params: None, reasoning: None, + request_headers: None, } } diff --git a/crates/goose/src/providers/tetrate.rs b/crates/goose/src/providers/tetrate.rs index 3dddae360..2bb64231d 100644 --- a/crates/goose/src/providers/tetrate.rs +++ b/crates/goose/src/providers/tetrate.rs @@ -152,7 +152,9 @@ impl Provider for TetrateProvider { .with_retry(|| async { let resp = self .api_client - .response_post("v1/chat/completions", &payload) + .request("v1/chat/completions") + .model_headers(model_config)? + .response_post(&payload) .await?; let resp = handle_status(resp) .await