fix(cost): resolve databricks_v2 pricing and surface cost in standard usage update (#9925)

Co-authored-by: Douwe M Osinga <douwe@sidewalklabs.com>
This commit is contained in:
Jude Edwards
2026-06-23 13:19:14 -07:00
committed by GitHub
parent 51ccfe5b7d
commit a25f9e94d3
2 changed files with 60 additions and 3 deletions
@@ -35,7 +35,10 @@ pub fn canonical_name(provider: &str, model: &str) -> String {
}
fn is_meta_provider(provider: &str) -> bool {
matches!(provider, "databricks" | "tetrate" | "bedrock" | "azure")
matches!(
provider,
"databricks" | "databricks_v2" | "tetrate" | "bedrock" | "azure"
)
}
pub fn map_provider_name(provider: &str) -> &str {
@@ -46,6 +49,7 @@ pub fn map_provider_name(provider: &str) -> &str {
"aws_bedrock" => "amazon-bedrock",
"gcp_vertex_ai" => "google-vertex",
"gemini_oauth" => "google",
"databricks_v2" => "databricks",
"zhipu" => "zhipuai",
"novita" => "novita-ai",
"opencode_go" => "opencode-go",
@@ -121,6 +125,21 @@ pub fn map_to_canonical_model(
}
}
// Fallback for meta-providers: some native aliases are keyed under the
// meta-provider itself (e.g. "databricks/databricks-gpt-oss-120b") and do
// not infer back to a first-party provider. Only try this after inference,
// so models that DO infer (e.g. databricks-claude-* -> anthropic/*) keep
// resolving to the richer first-party catalog entry.
if is_meta_provider(provider) {
if let Some(canonical) = registry.get(registry_provider, model) {
return Some(canonical.id.clone());
}
let normalized_model = strip_version_suffix(model);
if let Some(canonical) = registry.get(registry_provider, &normalized_model) {
return Some(canonical.id.clone());
}
}
None
}
@@ -537,4 +556,36 @@ mod tests {
Some("google-vertex/claude-haiku-4.5".to_string())
);
}
// Databricks-native open-weight ids are keyed under the meta-provider itself
// (e.g. "databricks/databricks-gpt-oss-120b") and do not infer back to another
// provider, so they must resolve via the direct meta-provider lookup. These
// particular ids are unversioned, so the assertions are not catalog-version brittle.
#[test]
fn test_databricks_native_open_weight_ids_resolve() {
let r = super::super::CanonicalModelRegistry::bundled().unwrap();
assert_eq!(
map_to_canonical_model("databricks_v2", "databricks-gpt-oss-120b", r),
Some("databricks/databricks-gpt-oss-120b".to_string())
);
assert_eq!(
map_to_canonical_model("databricks_v2", "databricks-gpt-oss-20b", r),
Some("databricks/databricks-gpt-oss-20b".to_string())
);
// Legacy provider name resolves identically.
assert_eq!(
map_to_canonical_model("databricks", "databricks-gpt-oss-120b", r),
Some("databricks/databricks-gpt-oss-120b".to_string())
);
// Regression guard: the meta-provider lookup must remain a *fallback*
// after inference. databricks-claude-* aliases infer back to the richer
// first-party "anthropic/*" entry (which carries thinking_mode used for
// adaptive thinking), not the metadata-poor "databricks/databricks-*" one.
assert_eq!(
map_to_canonical_model("databricks", "databricks-claude-opus-4-7", r),
Some("anthropic/claude-opus-4.7".to_string())
);
}
}
+8 -2
View File
@@ -42,7 +42,7 @@ use crate::utils::sanitize_unicode_tags;
use agent_client_protocol::schema::{
AgentCapabilities, Annotations, AuthMethod, AuthMethodAgent, AuthenticateRequest,
AuthenticateResponse, BlobResourceContents, CancelNotification, CloseSessionRequest,
CloseSessionResponse, ConfigOptionUpdate, Content, ContentBlock, ContentChunk,
CloseSessionResponse, ConfigOptionUpdate, Content, ContentBlock, ContentChunk, Cost,
CurrentModeUpdate, EmbeddedResource, EmbeddedResourceResource, FileSystemCapabilities,
ForkSessionRequest, ForkSessionResponse, ImageContent, Implementation, InitializeRequest,
InitializeResponse, ListSessionsRequest, ListSessionsResponse, LoadSessionRequest,
@@ -834,7 +834,13 @@ pub(super) fn build_usage_updates(session: &Session) -> Option<UsageUpdates> {
accumulated_cost: session.accumulated_cost,
}),
},
standard: UsageUpdate::new(used, ctx_limit),
standard: {
let mut standard = UsageUpdate::new(used, ctx_limit);
if let Some(amount) = session.accumulated_cost {
standard = standard.cost(Cost::new(amount, "USD"));
}
standard
},
})
}