From a25f9e94d3e50b8a6302dc95ce880b328d451156 Mon Sep 17 00:00:00 2001
From: Jude Edwards <judeedwards@squareup.com>
Date: Tue, 23 Jun 2026 13:19:14 -0700
Subject: [PATCH] fix(cost): resolve databricks_v2 pricing and surface cost in
 standard usage update (#9925)

Co-authored-by: Douwe M Osinga <douwe@sidewalklabs.com>
---
 .../src/canonical/name_builder.rs             | 53 ++++++++++++++++++-
 crates/goose/src/acp/server.rs                | 10 +++-
 2 files changed, 60 insertions(+), 3 deletions(-)

diff --git a/crates/goose-providers/src/canonical/name_builder.rs b/crates/goose-providers/src/canonical/name_builder.rs
index 640820101..2d377bcf4 100644
--- a/crates/goose-providers/src/canonical/name_builder.rs
+++ b/crates/goose-providers/src/canonical/name_builder.rs
@@ -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())
+        );
+    }
 }
diff --git a/crates/goose/src/acp/server.rs b/crates/goose/src/acp/server.rs
index 4b9b0c6af..5b038cf03 100644
--- a/crates/goose/src/acp/server.rs
+++ b/crates/goose/src/acp/server.rs
@@ -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
+        },
     })
 }