tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 34 Packages Projects Releases Wiki Activity

Add Azure AI Foundry multi-LLM provider (#10622)

Browse Source 
Co-authored-by: Jack Amadeo <jackamadeo@squareup.com>
This commit is contained in:
jonathan MERCIER
2026-07-30 17:14:26 +02:00
committed by GitHub
parent 866d3e4e47
commit a6420ea34b
25 changed files with 2048 additions and 174 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-provider-types/src/canonical.rs
+14 -7
View File
@@ -65,10 +65,10 @@ pub fn recommended_models_from_registry(provider: &str) -> Vec<String> {
.collect()
}
/// Providers that run models locally — their cost is always zero regardless
/// of what the canonical registry says for the underlying model architecture.
fn is_local_provider(provider: &str) -> bool {
matches!(provider, "ollama" | "local")
/// Catalog pricing is not valid for local inference or Azure Foundry deployments.
/// Azure billing depends on deployment region, SKU, offer, and contract.
fn should_clear_catalog_pricing(provider: &str) -> bool {
matches!(provider, "ollama" | "local" | "azure_foundry")
}
pub fn maybe_get_canonical_model(provider: &str, model: &str) -> Option<CanonicalModel> {
@@ -81,9 +81,7 @@ pub fn maybe_get_canonical_model(provider: &str, model: &str) -> Option<Canonica
return None;
};
// Local providers run models on the user's own hardware — zero out cloud
// pricing so every consumer (CLI, server, etc.) sees the correct cost.
if is_local_provider(provider) {
if should_clear_catalog_pricing(provider) {
canonical.cost = Pricing::default();
}
@@ -108,6 +106,15 @@ mod tests {
);
}
#[test]
fn azure_foundry_models_retain_limits_without_catalog_pricing() {
let canonical = maybe_get_canonical_model("azure_foundry", "gpt-5")
.expect("gpt-5 should resolve through the Azure catalog");
assert_eq!(canonical.limit.context, 400_000);
assert_eq!(canonical.cost.input, None);
assert_eq!(canonical.cost.output, None);
}
#[test]
fn cloud_provider_retains_cost() {
let canonical = maybe_get_canonical_model("anthropic", "claude-sonnet-4-5-20250929")
crates/goose-provider-types/src/canonical/name_builder.rs
+6 -2
View File
@@ -37,7 +37,7 @@ pub fn canonical_name(provider: &str, model: &str) -> String {
fn is_meta_provider(provider: &str) -> bool {
matches!(
provider,
"databricks" | "databricks_v2" | "tetrate" | "bedrock" | "azure"
"databricks" | "databricks_v2" | "tetrate" | "bedrock" | "azure" | "azure_foundry"
)
}
@@ -45,7 +45,7 @@ pub fn map_provider_name(provider: &str) -> &str {
match provider {
// Goose provider names that differ from models.dev names
"xai" => "x-ai",
"azure_openai" => "azure",
"azure_openai" | "azure_foundry" => "azure",
"aws_bedrock" => "amazon-bedrock",
"gcp_vertex_ai" => "google-vertex",
"gemini_oauth" => "google",
@@ -420,6 +420,10 @@ mod tests {
map_to_canonical_model("azure", "gpt-4o", r),
Some("openai/gpt-4o".to_string())
);
assert_eq!(
map_to_canonical_model("azure_foundry", "gpt-4o", r),
Some("openai/gpt-4o".to_string())
);
// === OpenAI O-series ===
assert_eq!(
crates/goose-provider-types/src/formats/anthropic.rs
+21 -1
View File
@@ -776,6 +776,26 @@ pub fn create_request(
messages: &[Message],
tools: &[Tool],
options: AnthropicFormatOptions,
) -> Result<Value> {
create_request_for_model(
provider_name,
model_config,
&model_config.model_name,
system,
messages,
tools,
options,
)
}
pub fn create_request_for_model(
provider_name: &str,
model_config: &ModelConfig,
wire_model_name: &str,
system: &str,
messages: &[Message],
tools: &[Tool],
options: AnthropicFormatOptions,
) -> Result<Value> {
let options = options.for_model(model_config);
let anthropic_messages = format_messages_with_options(messages, options);
@@ -788,7 +808,7 @@ pub fn create_request(
let max_tokens = model_config.max_output_tokens();
let mut payload = json!({
"model": model_config.model_name,
"model": wire_model_name,
"messages": anthropic_messages,
"max_tokens": max_tokens,
});
crates/goose-provider-types/src/formats/openai.rs
+28 -2
View File
@@ -1401,6 +1401,32 @@ pub fn create_request_with_options(
image_format: &ImageFormat,
for_streaming: bool,
format_options: OpenAiFormatOptions,
) -> anyhow::Result<Value, Error> {
let (wire_model_name, _) = extract_reasoning_effort(&model_config.model_name);
create_request_for_model_with_options(
model_config,
&wire_model_name,
&model_config.model_name,
system,
messages,
tools,
image_format,
for_streaming,
format_options,
)
}
#[allow(clippy::too_many_arguments)]
pub fn create_request_for_model_with_options(
model_config: &ModelConfig,
wire_model_name: &str,
capability_model_name: &str,
system: &str,
messages: &[Message],
tools: &[Tool],
image_format: &ImageFormat,
for_streaming: bool,
format_options: OpenAiFormatOptions,
) -> anyhow::Result<Value, Error> {
if model_config.model_name.starts_with("o1-mini") {
return Err(anyhow!(
@@ -1408,7 +1434,7 @@ pub fn create_request_with_options(
));
}
let (model_name, legacy_reasoning_effort) = extract_reasoning_effort(&model_config.model_name);
let (model_name, legacy_reasoning_effort) = extract_reasoning_effort(capability_model_name);
let is_reasoning_model = is_openai_responses_model(&model_name);
let supports_xai_effort = supports_xai_reasoning_effort(&model_name);
let reasoning_effort = if is_reasoning_model {
@@ -1439,7 +1465,7 @@ pub fn create_request_with_options(
messages_array.extend(messages_spec);
let mut payload = json!({
"model": model_name,
"model": wire_model_name,
"messages": messages_array
});
crates/goose-provider-types/src/formats/openai_responses.rs
+28 -8
View File
@@ -388,6 +388,7 @@ fn add_message_items(input_items: &mut Vec<Value>, messages: &[Message]) {
MessageContentBlock::ToolRequest(request) if message.role == Role::Assistant => {
if !text_items.is_empty() {
input_items.push(json!({
"type": "message",
"role": role,
"content": text_items
}));
@@ -435,6 +436,7 @@ fn add_message_items(input_items: &mut Vec<Value>, messages: &[Message]) {
MessageContentBlock::ToolResponse(response) => {
if !text_items.is_empty() {
input_items.push(json!({
"type": "message",
"role": role,
"content": text_items
}));
@@ -520,6 +522,7 @@ fn add_message_items(input_items: &mut Vec<Value>, messages: &[Message]) {
MessageContentBlock::FrontendToolRequest(request) => {
if !text_items.is_empty() {
input_items.push(json!({
"type": "message",
"role": role,
"content": text_items
}));
@@ -559,6 +562,7 @@ fn add_message_items(input_items: &mut Vec<Value>, messages: &[Message]) {
if !text_items.is_empty() {
input_items.push(json!({
"type": "message",
"role": role,
"content": text_items
}));
@@ -579,11 +583,31 @@ pub fn create_responses_request(
system: &str,
messages: &[Message],
tools: &[Tool],
) -> anyhow::Result<Value, Error> {
let (wire_model_name, _) = extract_reasoning_effort(&model_config.model_name);
create_responses_request_for_model(
model_config,
&wire_model_name,
&model_config.model_name,
system,
messages,
tools,
)
}
pub fn create_responses_request_for_model(
model_config: &ModelConfig,
wire_model_name: &str,
capability_model_name: &str,
system: &str,
messages: &[Message],
tools: &[Tool],
) -> anyhow::Result<Value, Error> {
let mut input_items = Vec::new();
if !system.is_empty() {
input_items.push(json!({
"type": "message",
"role": "system",
"content": [{
"type": "input_text",
@@ -594,7 +618,7 @@ pub fn create_responses_request(
add_message_items(&mut input_items, messages);
let (model_name, legacy_reasoning_effort) = extract_reasoning_effort(&model_config.model_name);
let (model_name, legacy_reasoning_effort) = extract_reasoning_effort(capability_model_name);
// All models routed here are responses-capable; temperature is rejected
// by the API for reasoning models regardless of whether an explicit
// effort suffix was provided.
@@ -639,7 +663,7 @@ pub fn create_responses_request(
));
}
let mut payload = json!({
"model": model_name,
"model": wire_model_name,
"input": input_items,
"store": store,
});
@@ -1358,16 +1382,12 @@ mod tests {
let types: Vec<&str> = input
.iter()
.map(|item| {
item.get("type")
.and_then(|v| v.as_str())
.unwrap_or_else(|| item["role"].as_str().unwrap())
})
.map(|item| item["type"].as_str().unwrap())
.collect();
assert_eq!(
types,
vec!["assistant", "function_call", "assistant", "function_call"]
vec!["message", "function_call", "message", "function_call"]
);
}
crates/goose-providers/src/anthropic.rs
+61 -40
View File
@@ -16,7 +16,8 @@ use tokio_util::io::StreamReader;
use super::api_client::ApiClient;
use super::base::{ConfigKey, MessageStream, ModelInfo, Provider, ProviderMetadata};
use super::formats::anthropic::{
create_request, response_to_streaming_message, AnthropicFormatOptions, ANTHROPIC_PROVIDER_NAME,
create_request_for_model, response_to_streaming_message, AnthropicFormatOptions,
ANTHROPIC_PROVIDER_NAME,
};
use super::openai_compatible::handle_status;
use super::openai_compatible::map_http_error_to_provider_error;
@@ -155,6 +156,54 @@ impl AnthropicProviderBuilder {
}
impl AnthropicProvider {
pub async fn stream_for_model(
&self,
model_config: &ModelConfig,
wire_model: &str,
system: &str,
messages: &[Message],
tools: &[Tool],
) -> Result<MessageStream, ProviderError> {
let mut payload = create_request_for_model(
ANTHROPIC_PROVIDER_NAME,
model_config,
wire_model,
system,
messages,
tools,
self.format_options,
)?;
payload["stream"] = Value::Bool(true);
let mut log = start_log(model_config, &payload)?;
let response = self
.with_retry(|| async {
handle_status(
self.api_client
.request("v1/messages")
.model_headers(model_config)?
.response_post(&payload)
.await?,
)
.await
})
.await
.inspect_err(|e| {
let _ = log.error(e);
})?;
let stream = response.bytes_stream().map_err(io::Error::other);
Ok(Box::pin(try_stream! {
let reader = StreamReader::new(stream);
let framed = tokio_util::codec::FramedRead::new(reader, tokio_util::codec::LinesCodec::new()).map_err(anyhow::Error::from);
let messages = response_to_streaming_message(framed);
pin!(messages);
while let Some(message) = futures::StreamExt::next(&mut messages).await {
let (message, usage) = message.map_err(ProviderError::from_stream_error)?;
log.write(&message, usage.as_ref().map(|f| f.usage).as_ref())?;
yield (message, usage);
}
}))
}
async fn fetch_models_from_api(&self) -> Result<Vec<String>, ProviderError> {
let response = self.api_client.request("v1/models").api_get().await?;
@@ -233,6 +282,13 @@ impl Provider for AnthropicProvider {
&self.name
}
async fn refresh_credentials(&self) -> Result<(), ProviderError> {
self.api_client
.refresh_credentials()
.await
.map_err(|error| ProviderError::Authentication(error.to_string()))
}
fn skip_canonical_filtering(&self) -> bool {
self.skip_canonical_filtering
}
@@ -266,49 +322,14 @@ impl Provider for AnthropicProvider {
messages: &[Message],
tools: &[Tool],
) -> Result<MessageStream, ProviderError> {
let mut payload = create_request(
ANTHROPIC_PROVIDER_NAME,
self.stream_for_model(
model_config,
&model_config.model_name,
system,
messages,
tools,
self.format_options,
)?;
payload
.as_object_mut()
.unwrap()
.insert("stream".to_string(), Value::Bool(true));
let mut log = start_log(model_config, &payload)?;
let response = self
.with_retry(|| async {
let request = self
.api_client
.request("v1/messages")
.model_headers(model_config)?;
let resp = request.response_post(&payload).await?;
handle_status(resp).await
})
.await
.inspect_err(|e| {
let _ = log.error(e);
})?;
let stream = response.bytes_stream().map_err(io::Error::other);
Ok(Box::pin(try_stream! {
let stream_reader = StreamReader::new(stream);
let framed = tokio_util::codec::FramedRead::new(stream_reader, tokio_util::codec::LinesCodec::new()).map_err(anyhow::Error::from);
let message_stream = response_to_streaming_message(framed);
pin!(message_stream);
while let Some(message) = futures::StreamExt::next(&mut message_stream).await {
let (message, usage) = message.map_err(ProviderError::from_stream_error)?;
log.write(&message, usage.as_ref().map(|f| f.usage).as_ref())?;
yield (message, usage);
}
}))
)
.await
}
}
crates/goose-providers/src/api_client.rs
+11
View File
@@ -195,6 +195,10 @@ fn convert_key_to_pkcs8_pem(key_pem_str: &str) -> Result<String> {
#[async_trait]
pub trait AuthProvider: Send + Sync {
async fn get_auth_header(&self) -> Result<(String, String)>;
async fn refresh_credentials(&self) -> Result<()> {
anyhow::bail!("credential refresh not supported")
}
}
pub struct ApiResponse {
@@ -356,6 +360,13 @@ impl ApiClient {
}
}
pub async fn refresh_credentials(&self) -> Result<()> {
match &self.auth {
AuthMethod::Custom(provider) => provider.refresh_credentials().await,
_ => anyhow::bail!("credential refresh not supported"),
}
}
pub async fn api_post(&self, path: &str, payload: &Value) -> Result<ApiResponse> {
self.request(path).api_post(payload).await
}
crates/goose-providers/src/azure_foundry.rs
+1175
View File
File diff suppressed because it is too large Load Diff
crates/goose-providers/src/lib.rs
+1
View File
@@ -1,5 +1,6 @@
pub mod anthropic;
pub mod api_client;
pub mod azure_foundry;
pub mod databricks;
pub mod databricks_auth;
pub mod databricks_v2;
crates/goose-providers/src/openai.rs
+94 -56
View File
@@ -11,7 +11,8 @@ use crate::formats::openai::{
create_request_with_options, get_cost, get_usage, response_to_message, OpenAiFormatOptions,
};
use crate::formats::openai_responses::{
create_responses_request, get_responses_usage, responses_api_to_message, ResponsesApiResponse,
create_responses_request_for_model, get_responses_usage, responses_api_to_message,
ResponsesApiResponse,
};
use crate::images::ImageFormat;
use crate::openai_compatible::{
@@ -272,6 +273,80 @@ impl OpenAiProviderBuilder {
}
impl OpenAiProvider {
pub async fn stream_for_model(
&self,
model_config: &ModelConfig,
wire_model: &str,
capability_model: &str,
system: &str,
messages: &[Message],
tools: &[Tool],
) -> Result<MessageStream, ProviderError> {
let mut payload = create_responses_request_for_model(
model_config,
wire_model,
capability_model,
system,
messages,
tools,
)?;
payload["stream"] = serde_json::Value::Bool(self.supports_streaming);
self.stream_responses_payload(model_config, payload).await
}
async fn stream_responses_payload(
&self,
model_config: &ModelConfig,
payload: serde_json::Value,
) -> Result<MessageStream, ProviderError> {
let mut log = start_log(model_config, &payload)?;
let response = self
.with_retry(|| async {
handle_status(
self.api_client
.request(&Self::map_base_path(
&self.base_path,
"responses",
OPEN_AI_DEFAULT_RESPONSES_PATH,
))
.model_headers(model_config)?
.response_post(&payload)
.await?,
)
.await
})
.await
.inspect_err(|e| {
let _ = log.error(e);
})?;
if self.supports_streaming {
stream_responses_compat(response, log)
} else {
let json: serde_json::Value = response.json().await.map_err(|e| {
ProviderError::RequestFailed(format!("Failed to parse JSON: {}", e))
})?;
let parsed: ResponsesApiResponse =
serde_json::from_value(json.clone()).map_err(|e| {
ProviderError::ExecutionError(format!(
"Failed to parse responses API response: {}",
e
))
})?;
let message = responses_api_to_message(&parsed)?;
let usage_data = get_responses_usage(&parsed);
let usage_json = json.get("usage").unwrap_or(&serde_json::Value::Null);
let mut usage = ProviderUsage::new(model_config.model_name.clone(), usage_data);
if let Some(cost) = get_cost(usage_json) {
usage = usage.with_cost(cost, CostSource::ProviderReported);
}
log.write(
&serde_json::to_value(&message).unwrap_or_default(),
Some(&usage_data),
)?;
Ok(super::base::stream_from_single_message(message, usage))
}
}
#[doc(hidden)]
pub fn new(api_client: ApiClient) -> Self {
Self {
@@ -584,6 +659,13 @@ impl Provider for OpenAiProvider {
&self.name
}
async fn refresh_credentials(&self) -> Result<(), ProviderError> {
self.api_client
.refresh_credentials()
.await
.map_err(|error| ProviderError::Authentication(error.to_string()))
}
fn skip_canonical_filtering(&self) -> bool {
self.skip_canonical_filtering
}
@@ -655,61 +737,17 @@ impl Provider for OpenAiProvider {
tools: &[Tool],
) -> Result<MessageStream, ProviderError> {
if self.should_use_responses_api_for_provider(&model_config.model_name) {
let mut payload = create_responses_request(model_config, system, messages, tools)?;
payload["stream"] = serde_json::Value::Bool(self.supports_streaming);
let mut log = start_log(model_config, &payload)?;
let response = self
.with_retry(|| async {
let payload_clone = payload.clone();
let resp = self
.api_client
.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
})
.await
.inspect_err(|e| {
let _ = log.error(e);
})?;
if self.supports_streaming {
stream_responses_compat(response, log)
} else {
let json: serde_json::Value = response.json().await.map_err(|e| {
ProviderError::RequestFailed(format!("Failed to parse JSON: {}", e))
})?;
let responses_api_response: ResponsesApiResponse =
serde_json::from_value(json.clone()).map_err(|e| {
ProviderError::ExecutionError(format!(
"Failed to parse responses API response: {}",
e
))
})?;
let message = responses_api_to_message(&responses_api_response)?;
let usage_data = get_responses_usage(&responses_api_response);
let usage_json = json.get("usage").unwrap_or(&serde_json::Value::Null);
let mut usage = ProviderUsage::new(model_config.model_name.clone(), usage_data);
if let Some(cost) = get_cost(usage_json) {
usage = usage.with_cost(cost, CostSource::ProviderReported);
}
log.write(
&serde_json::to_value(&message).unwrap_or_default(),
Some(&usage_data),
)?;
Ok(super::base::stream_from_single_message(message, usage))
}
let (wire_model, _) =
crate::formats::openai::extract_reasoning_effort(&model_config.model_name);
self.stream_for_model(
model_config,
&wire_model,
&model_config.model_name,
system,
messages,
tools,
)
.await
} else {
let payload = create_request_with_options(
model_config,
crates/goose-providers/src/openai_compatible.rs
+103 -44
View File
@@ -18,7 +18,8 @@ use super::retry::ProviderRetry;
use crate::conversation::message::Message;
use crate::errors::ProviderError;
use crate::formats::openai::{
create_request, get_cost, get_usage, response_to_message, response_to_streaming_message,
create_request, create_request_for_model_with_options, get_cost, get_usage,
response_to_message, response_to_streaming_message, OpenAiFormatOptions,
};
use crate::formats::openai_responses::responses_api_to_streaming_message;
use crate::model::ModelConfig;
@@ -49,6 +50,99 @@ impl OpenAiCompatibleProvider {
self
}
#[allow(clippy::too_many_arguments)]
fn build_request_for_model(
&self,
model_config: &ModelConfig,
wire_model: &str,
capability_model: &str,
system: &str,
messages: &[Message],
tools: &[Tool],
for_streaming: bool,
) -> Result<Value, ProviderError> {
create_request_for_model_with_options(
model_config,
wire_model,
capability_model,
system,
messages,
tools,
&ImageFormat::OpenAi,
for_streaming,
OpenAiFormatOptions {
preserve_thinking_context: true,
},
)
.map_err(|e| ProviderError::RequestFailed(format!("Failed to create request: {}", e)))
}
pub async fn stream_for_model(
&self,
model_config: &ModelConfig,
wire_model: &str,
capability_model: &str,
system: &str,
messages: &[Message],
tools: &[Tool],
) -> Result<MessageStream, ProviderError> {
let payload = self.build_request_for_model(
model_config,
wire_model,
capability_model,
system,
messages,
tools,
self.supports_streaming,
)?;
self.stream_payload(model_config, payload).await
}
async fn stream_payload(
&self,
model_config: &ModelConfig,
payload: Value,
) -> Result<MessageStream, ProviderError> {
let mut log = start_log(model_config, &payload)?;
let path = format!("{}chat/completions", self.completions_prefix);
let response = self
.with_retry(|| async {
handle_status(
self.api_client
.request(&path)
.model_headers(model_config)?
.response_post(&payload)
.await?,
)
.await
})
.await
.inspect_err(|e| {
let _ = log.error(e);
})?;
if self.supports_streaming {
stream_openai_compat(response, log)
} else {
let json = response.json().await.map_err(|e| {
ProviderError::RequestFailed(format!("Failed to parse JSON: {}", e))
})?;
let message = response_to_message(&json).map_err(|e| {
ProviderError::RequestFailed(format!("Failed to parse message: {}", e))
})?;
let usage_json = json.get("usage").unwrap_or(&Value::Null);
let usage_data = get_usage(usage_json);
let mut usage = ProviderUsage::new(model_config.model_name.clone(), usage_data);
if let Some(cost) = get_cost(usage_json) {
usage = usage.with_cost(cost, CostSource::ProviderReported);
}
log.write(
&serde_json::to_value(&message).unwrap_or_default(),
Some(&usage.usage),
)?;
Ok(stream_from_single_message(message, usage))
}
}
fn build_request(
&self,
model_config: &ModelConfig,
@@ -75,6 +169,13 @@ impl Provider for OpenAiCompatibleProvider {
&self.name
}
async fn refresh_credentials(&self) -> Result<(), ProviderError> {
self.api_client
.refresh_credentials()
.await
.map_err(|error| ProviderError::Authentication(error.to_string()))
}
async fn fetch_supported_models(&self) -> Result<Vec<String>, ProviderError> {
let response = self
.api_client
@@ -116,49 +217,7 @@ impl Provider for OpenAiCompatibleProvider {
tools,
self.supports_streaming,
)?;
let mut log = start_log(model_config, &payload)?;
let completions_path = format!("{}chat/completions", self.completions_prefix);
let response = self
.with_retry(|| async {
let resp = self
.api_client
.request(&completions_path)
.model_headers(model_config)?
.response_post(&payload)
.await?;
handle_status(resp).await
})
.await
.inspect_err(|e| {
let _ = log.error(e);
})?;
if self.supports_streaming {
stream_openai_compat(response, log)
} else {
let json: serde_json::Value = response.json().await.map_err(|e| {
ProviderError::RequestFailed(format!("Failed to parse JSON: {}", e))
})?;
let message = response_to_message(&json).map_err(|e| {
ProviderError::RequestFailed(format!("Failed to parse message: {}", e))
})?;
let usage_json = json.get("usage").unwrap_or(&serde_json::Value::Null);
let usage_data = get_usage(usage_json);
let mut usage = ProviderUsage::new(model_config.model_name.clone(), usage_data);
if let Some(cost) = get_cost(usage_json) {
usage = usage.with_cost(cost, CostSource::ProviderReported);
}
log.write(
&serde_json::to_value(&message).unwrap_or_default(),
Some(&usage.usage),
)?;
Ok(stream_from_single_message(message, usage))
}
self.stream_payload(model_config, payload).await
}
}
crates/goose/src/model_config.rs
+43 -6
View File
@@ -14,7 +14,7 @@ pub fn model_config_from_user_config(
provider_name: &str,
model_name: impl AsRef<str>,
) -> Result<ModelConfig> {
let model = base_model_config_from_user_config(model_name.as_ref())?;
let model = base_model_config_from_user_config(provider_name, model_name.as_ref())?;
materialize_model_config(provider_name, model)
}
@@ -26,18 +26,26 @@ pub fn model_config_from_user_config_with_session_settings(
context_limit: Option<usize>,
) -> Result<ModelConfig> {
let config = Config::global();
let model = base_model_config_from_user_config(model_name.as_ref())?;
let model = base_model_config_from_user_config(provider_name, model_name.as_ref())?;
let model = materialize_model_config_inner(model, provider_name, false)?
.with_context_limit(context_limit)
.with_inherited_session_settings_from(previous, request_params)
.with_default_thinking_effort(config.get_goose_thinking_effort());
Ok(model.with_canonical_limits(provider_name))
Ok(apply_canonical_limits(provider_name, model))
}
pub fn materialize_model_config(provider_name: &str, model: ModelConfig) -> Result<ModelConfig> {
let model = materialize_model_config_inner(model, provider_name, true)?;
Ok(model.with_canonical_limits(provider_name))
Ok(apply_canonical_limits(provider_name, model))
}
fn apply_canonical_limits(provider_name: &str, model: ModelConfig) -> ModelConfig {
if provider_name == goose_providers::azure_foundry::AZURE_FOUNDRY_PROVIDER_NAME {
model
} else {
model.with_canonical_limits(provider_name)
}
}
fn materialize_model_config_inner(
@@ -169,7 +177,10 @@ fn apply_openai_request_params(mut model: ModelConfig) -> ModelConfig {
model
}
fn base_model_config_from_user_config(model_name: &str) -> Result<ModelConfig> {
fn base_model_config_from_user_config(
provider_name: &str,
model_name: &str,
) -> Result<ModelConfig> {
let config = Config::global();
let mut model = ModelConfig {
model_name: model_name.to_string(),
@@ -182,7 +193,9 @@ fn base_model_config_from_user_config(model_name: &str) -> Result<ModelConfig> {
reasoning: None,
request_headers: None,
};
model.normalize_effort_suffix();
if provider_name != goose_providers::azure_foundry::AZURE_FOUNDRY_PROVIDER_NAME {
model.normalize_effort_suffix();
}
Ok(model)
}
@@ -247,3 +260,27 @@ fn parse_yaml_bool_config(key: &str, value: serde_yaml::Value) -> Result<bool> {
}
}
}
#[cfg(test)]
mod azure_foundry_tests {
use super::*;
#[test]
fn deployment_name_survives_thinking_effort_changes() {
let config = base_model_config_from_user_config("azure_foundry", "gpt-5-high")
.unwrap()
.with_thinking_effort(ThinkingEffort::Off);
assert_eq!(config.model_name, "gpt-5-high");
assert_eq!(config.context_limit, None);
assert_eq!(config.thinking_effort(), Some(ThinkingEffort::Off));
}
#[test]
fn none_suffixed_deployment_name_is_preserved() {
let config = base_model_config_from_user_config("azure_foundry", "gpt-5-none").unwrap();
assert_eq!(config.model_name, "gpt-5-none");
assert_eq!(config.thinking_effort(), None);
}
}
crates/goose/src/providers/azure_foundry_def.rs
+173
View File
@@ -0,0 +1,173 @@
use std::sync::Arc;
use anyhow::Result;
use async_trait::async_trait;
use futures::future::BoxFuture;
use goose_providers::api_client::{AuthMethod, AuthProvider, TlsConfig};
use goose_providers::azure_foundry::{endpoint_kind, AzureFoundryProvider, EndpointKind};
use goose_providers::base::{ProviderDescriptor, ProviderMetadata};
use crate::config::{Config, ExtensionConfig};
use crate::providers::azureauth::{AzureAuth, AzureCredentials};
use crate::providers::base::ProviderDef;
const AZURE_PROJECT_ENTRA_RESOURCE: &str = "https://ai.azure.com";
const AZURE_MAAS_ENTRA_RESOURCE: &str = "https://ml.azure.com";
enum AuthHeader {
ApiKey,
Bearer,
}
struct AzureFoundryAuthProvider {
auth: Arc<AzureAuth>,
header: AuthHeader,
}
#[async_trait]
impl AuthProvider for AzureFoundryAuthProvider {
async fn get_auth_header(&self) -> Result<(String, String)> {
let token = self.auth.get_token().await?;
match &self.header {
AuthHeader::ApiKey => Ok(("api-key".to_string(), token.token_value)),
AuthHeader::Bearer => Ok((
"Authorization".to_string(),
format!("Bearer {}", token.token_value),
)),
}
}
async fn refresh_credentials(&self) -> Result<()> {
self.auth.invalidate_token().await;
Ok(())
}
}
pub struct AzureFoundryProviderDef;
impl ProviderDescriptor for AzureFoundryProviderDef {
fn metadata() -> ProviderMetadata {
AzureFoundryProvider::metadata()
}
}
impl ProviderDef for AzureFoundryProviderDef {
type Provider = AzureFoundryProvider;
fn from_env(
_extensions: Vec<ExtensionConfig>,
tls_config: Option<TlsConfig>,
) -> BoxFuture<'static, Result<Self::Provider>> {
Box::pin(from_env(tls_config))
}
}
pub async fn from_env(tls_config: Option<TlsConfig>) -> Result<AzureFoundryProvider> {
let config = Config::global();
let endpoint: String = config.get_param("AZURE_FOUNDRY_ENDPOINT")?;
let api_version = config.get_param("AZURE_FOUNDRY_API_VERSION").ok();
let maas_model = config
.get_param::<String>("AZURE_FOUNDRY_MODEL")
.ok()
.filter(|model| !model.trim().is_empty());
let api_key = config
.get_secret::<String>("AZURE_FOUNDRY_API_KEY")
.ok()
.filter(|key| !key.is_empty());
let ad_token = config
.get_secret::<String>("AZURE_FOUNDRY_AD_TOKEN")
.ok()
.filter(|token| !token.is_empty());
let endpoint_kind = endpoint_kind(&endpoint);
let resource = if endpoint_kind == EndpointKind::Maas {
AZURE_MAAS_ENTRA_RESOURCE
} else {
AZURE_PROJECT_ENTRA_RESOURCE
};
let auth = Arc::new(AzureAuth::new_with_resource(
api_key,
ad_token,
resource.to_string(),
)?);
let auth_method = |header| {
AuthMethod::Custom(Box::new(AzureFoundryAuthProvider {
auth: Arc::clone(&auth),
header,
}))
};
let anthropic_auth = match auth.credential_type() {
AzureCredentials::ApiKey(key) => AuthMethod::ApiKey {
header_name: "x-api-key".to_string(),
key: key.clone(),
},
_ => auth_method(AuthHeader::Bearer),
};
let api_key_auth_header = || match auth.credential_type() {
AzureCredentials::ApiKey(_) => AuthHeader::ApiKey,
_ => AuthHeader::Bearer,
};
let chat_auth_header = match auth.credential_type() {
AzureCredentials::ApiKey(_) if endpoint_kind == EndpointKind::Maas => AuthHeader::Bearer,
_ => api_key_auth_header(),
};
AzureFoundryProvider::create(
endpoint,
api_version,
maas_model,
auth_method(chat_auth_header),
auth_method(api_key_auth_header()),
anthropic_auth,
auth_method(api_key_auth_header()),
tls_config,
Some(crate::session_context::session_id_request_builder()),
)
}
#[cfg(test)]
mod tests {
use super::*;
async fn header(
api_key: Option<&str>,
ad_token: Option<&str>,
header: AuthHeader,
) -> (String, String) {
let auth = Arc::new(
AzureAuth::new_with_resource(
api_key.map(str::to_string),
ad_token.map(str::to_string),
AZURE_PROJECT_ENTRA_RESOURCE.to_string(),
)
.unwrap(),
);
AzureFoundryAuthProvider { auth, header }
.get_auth_header()
.await
.unwrap()
}
#[tokio::test]
async fn project_api_key_uses_api_key_header() {
assert_eq!(
header(Some("key"), None, AuthHeader::ApiKey).await,
("api-key".to_string(), "key".to_string())
);
}
#[tokio::test]
async fn maas_api_key_uses_bearer_header() {
assert_eq!(
header(Some("key"), None, AuthHeader::Bearer).await,
("Authorization".to_string(), "Bearer key".to_string())
);
}
#[tokio::test]
async fn entra_token_uses_bearer_header() {
assert_eq!(
header(None, Some("token"), AuthHeader::Bearer).await,
("Authorization".to_string(), "Bearer token".to_string())
);
}
}
crates/goose/src/providers/azureauth.rs
+19 -6
View File
@@ -59,6 +59,7 @@ struct TokenResponse {
#[derive(Debug)]
pub struct AzureAuth {
credentials: AzureCredentials,
resource: String,
cached_token: Arc<RwLock<Option<CachedToken>>>,
}
@@ -73,6 +74,18 @@ impl AzureAuth {
/// # Returns
/// * `Result<Self, AuthError>` - A new AzureAuth instance or an error if initialization fails
pub fn new(api_key: Option<String>, ad_token: Option<String>) -> Result<Self, AuthError> {
Self::new_with_resource(
api_key,
ad_token,
"https://cognitiveservices.azure.com".to_string(),
)
}
pub fn new_with_resource(
api_key: Option<String>,
ad_token: Option<String>,
resource: String,
) -> Result<Self, AuthError> {
let credentials = match (ad_token, api_key) {
(Some(token), _) => AzureCredentials::BearerToken(token),
(None, Some(key)) => AzureCredentials::ApiKey(key),
@@ -81,6 +94,7 @@ impl AzureAuth {
Ok(Self {
credentials,
resource,
cached_token: Arc::new(RwLock::new(None)),
})
}
@@ -90,6 +104,10 @@ impl AzureAuth {
&self.credentials
}
pub async fn invalidate_token(&self) {
*self.cached_token.write().await = None;
}
/// Retrieves a valid authentication token.
///
/// This method implements an efficient token management strategy:
@@ -137,12 +155,7 @@ impl AzureAuth {
let az = if cfg!(windows) { "az.cmd" } else { "az" };
let output = tokio::process::Command::new(az)
.args([
"account",
"get-access-token",
"--resource",
"https://cognitiveservices.azure.com",
])
.args(["account", "get-access-token", "--resource", &self.resource])
.set_no_window()
.output()
.await
crates/goose/src/providers/init.rs
+5
View File
@@ -37,6 +37,7 @@ use super::{
};
use crate::config::ExtensionConfig;
use crate::providers::anthropic_def::AnthropicProviderDef;
use crate::providers::azure_foundry_def::AzureFoundryProviderDef;
use crate::providers::base::ProviderType;
use crate::providers::databricks_def::{self, DatabricksProviderDef};
use crate::providers::databricks_v2_def::{self, DatabricksV2ProviderDef};
@@ -69,6 +70,10 @@ async fn init_registry() -> RwLock<ProviderRegistry> {
);
registry.register::<AvianProvider>(false);
registry.register::<AzureProvider>(false);
registry.register_with_inventory::<AzureFoundryProviderDef>(
true,
Some(registrations::azure_foundry_inventory()),
);
#[cfg(feature = "aws-providers")]
registry.register::<BedrockProvider>(false);
#[cfg(feature = "local-inference")]
crates/goose/src/providers/inventory/registrations.rs
+71
View File
@@ -20,6 +20,7 @@ use crate::providers::ollama::OLLAMA_PROVIDER_NAME;
use crate::providers::openai::{OPEN_AI_DEFAULT_BASE_PATH, OPEN_AI_PROVIDER_NAME};
use crate::providers::pi_acp::{PI_ACP_BINARY, PI_ACP_PROVIDER_NAME};
use crate::providers::xai_oauth::TokenCache as XaiOAuthTokenCache;
use goose_providers::azure_foundry::{endpoint_kind, EndpointKind, AZURE_FOUNDRY_PROVIDER_NAME};
pub fn openai_inventory() -> InventoryRegistration {
InventoryRegistration::new(true, || {
@@ -67,6 +68,52 @@ pub fn openai_inventory() -> InventoryRegistration {
})
}
pub fn azure_foundry_inventory() -> InventoryRegistration {
InventoryRegistration::new(true, || {
let config = Config::global();
let mut identity =
InventoryIdentityInput::new(AZURE_FOUNDRY_PROVIDER_NAME, AZURE_FOUNDRY_PROVIDER_NAME);
if let Ok(endpoint) = config.get_param::<String>("AZURE_FOUNDRY_ENDPOINT") {
identity = identity.with_public("endpoint", endpoint);
}
if let Ok(api_version) = config.get_param::<String>("AZURE_FOUNDRY_API_VERSION") {
identity = identity.with_public("api_version", api_version);
}
if let Ok(model) = config.get_param::<String>("AZURE_FOUNDRY_MODEL") {
identity = identity.with_public("model", model);
}
if let Some(api_key) = config_secret_value(config, "AZURE_FOUNDRY_API_KEY") {
identity = identity.with_secret("api_key", api_key);
}
if let Some(ad_token) = config_secret_value(config, "AZURE_FOUNDRY_AD_TOKEN") {
identity = identity.with_secret("ad_token", ad_token);
}
Ok(identity)
})
.with_configured(|| azure_foundry_configured(Config::global()))
}
fn azure_foundry_configured(config: &Config) -> bool {
azure_foundry_configured_values(
config
.get_param::<String>("AZURE_FOUNDRY_ENDPOINT")
.ok()
.as_deref(),
config
.get_param::<String>("AZURE_FOUNDRY_MODEL")
.ok()
.as_deref(),
)
}
fn azure_foundry_configured_values(endpoint: Option<&str>, model: Option<&str>) -> bool {
let Some(endpoint) = endpoint.filter(|endpoint| !endpoint.trim().is_empty()) else {
return false;
};
endpoint_kind(endpoint) != EndpointKind::Maas
|| model.is_some_and(|model| !model.trim().is_empty())
}
pub fn anthropic_inventory() -> InventoryRegistration {
InventoryRegistration::new(true, || {
let config = Config::global();
@@ -222,6 +269,30 @@ mod tests {
use crate::config::paths::Paths;
use chrono::Utc;
#[test]
fn azure_foundry_maas_requires_a_model_to_be_configured() {
assert!(!azure_foundry_configured_values(
Some("https://deployment.models.ai.azure.com"),
None,
));
assert!(!azure_foundry_configured_values(
Some("https://deployment.models.ai.azure.com"),
Some(" "),
));
assert!(azure_foundry_configured_values(
Some("https://deployment.models.ai.azure.com"),
Some("Phi-4"),
));
assert!(azure_foundry_configured_values(
Some("https://hub.services.ai.azure.com/api/projects/project"),
None,
));
assert!(azure_foundry_configured_values(
Some("https://hub.services.ai.azure.com"),
None,
));
}
#[test]
#[serial_test::serial]
fn gemini_oauth_inventory_configured_uses_token_cache() {
crates/goose/src/providers/mod.rs
+1
View File
@@ -9,6 +9,7 @@ pub mod api_client {
}
pub mod avian;
pub mod azure;
pub mod azure_foundry_def;
pub mod azureauth;
pub mod base;
#[cfg(feature = "aws-providers")]
crates/goose/src/session/session_manager.rs
+81 -2
View File
@@ -715,6 +715,27 @@ impl Session {
}
}
fn deserialize_session_model_config(
provider_name: Option<&str>,
json: &str,
) -> Option<ModelConfig> {
let mut model_config: ModelConfig = serde_json::from_str(json).ok()?;
// TODO: Remove this workaround once ModelConfig guarantees deserialize(serialize(config)) == config.
if provider_name == Some(goose_providers::azure_foundry::AZURE_FOUNDRY_PROVIDER_NAME) {
#[derive(Deserialize)]
struct AzurePersistedFields {
model_name: String,
#[serde(default)]
request_params: Option<HashMap<String, serde_json::Value>>,
}
let persisted: AzurePersistedFields = serde_json::from_str(json).ok()?;
model_config.model_name = persisted.model_name;
model_config.request_params = persisted.request_params;
}
Some(model_config)
}
impl sqlx::FromRow<'_, sqlx::sqlite::SqliteRow> for Session {
fn from_row(row: &sqlx::sqlite::SqliteRow) -> Result<Self, sqlx::Error> {
use sqlx::Row;
@@ -726,8 +747,11 @@ impl sqlx::FromRow<'_, sqlx::sqlite::SqliteRow> for Session {
let user_recipe_values =
user_recipe_values_json.and_then(|json| serde_json::from_str(&json).ok());
let provider_name: Option<String> = row.try_get("provider_name").ok().flatten();
let model_config_json: Option<String> = row.try_get("model_config_json").ok().flatten();
let model_config = model_config_json.and_then(|json| serde_json::from_str(&json).ok());
let model_config = model_config_json
.as_deref()
.and_then(|json| deserialize_session_model_config(provider_name.as_deref(), json));
let name: String = {
let name_val: String = row.try_get("name").unwrap_or_default();
@@ -788,7 +812,7 @@ impl sqlx::FromRow<'_, sqlx::sqlite::SqliteRow> for Session {
conversation: None,
message_count: row.try_get("message_count").unwrap_or(0) as usize,
last_message_at,
provider_name: row.try_get("provider_name").ok().flatten(),
provider_name,
model_config,
goose_mode: row
.try_get::<String, _>("goose_mode")
@@ -2587,6 +2611,61 @@ mod tests {
const NUM_CONCURRENT_SESSIONS: i32 = 10;
const GENERATED_SESSION_NAME: &str = "Generated session name";
#[test]
fn azure_session_model_config_preserves_suffixed_deployment_id() {
let json = serde_json::to_string(&ModelConfig {
model_name: "gpt-5-high".to_string(),
context_limit: None,
temperature: None,
max_tokens: None,
toolshim: false,
toolshim_model: None,
request_params: None,
reasoning: None,
request_headers: None,
})
.unwrap();
let config = deserialize_session_model_config(
Some(goose_providers::azure_foundry::AZURE_FOUNDRY_PROVIDER_NAME),
&json,
)
.unwrap();
assert_eq!(config.model_name, "gpt-5-high");
assert_eq!(config.thinking_effort(), None);
}
#[test]
fn azure_session_model_config_preserves_explicit_thinking_effort() {
let config = deserialize_session_model_config(
Some(goose_providers::azure_foundry::AZURE_FOUNDRY_PROVIDER_NAME),
r#"{"model_name":"gpt-5-high","context_limit":null,"temperature":null,"max_tokens":null,"toolshim":false,"toolshim_model":null,"request_params":{"thinking_effort":"low"}}"#,
)
.unwrap();
assert_eq!(config.model_name, "gpt-5-high");
assert_eq!(
config.thinking_effort(),
Some(goose_providers::thinking::ThinkingEffort::Low)
);
}
#[test]
fn non_azure_session_model_config_keeps_suffix_normalization() {
let config = deserialize_session_model_config(
Some(goose_providers::openai::OPEN_AI_PROVIDER_NAME),
r#"{"model_name":"gpt-5-high","context_limit":null,"temperature":null,"max_tokens":null,"toolshim":false,"toolshim_model":null}"#,
)
.unwrap();
assert_eq!(config.model_name, "gpt-5");
assert_eq!(
config.thinking_effort(),
Some(goose_providers::thinking::ThinkingEffort::High)
);
}
struct NamingTestProvider;
#[async_trait::async_trait]