From a5a497b3416252b8946f5d8a0215a4d2a808873e Mon Sep 17 00:00:00 2001 From: Lifei Zhou Date: Thu, 16 Apr 2026 14:32:19 +1000 Subject: [PATCH] pass globally unique conversation identifier as sessionId in databricks api call (#8576) --- crates/goose/src/instance_id.rs | 49 ++++++++++++++++++++++++ crates/goose/src/lib.rs | 1 + crates/goose/src/providers/databricks.rs | 40 ++++++++++++++++++- 3 files changed, 89 insertions(+), 1 deletion(-) create mode 100644 crates/goose/src/instance_id.rs diff --git a/crates/goose/src/instance_id.rs b/crates/goose/src/instance_id.rs new file mode 100644 index 00000000..e0013ec7 --- /dev/null +++ b/crates/goose/src/instance_id.rs @@ -0,0 +1,49 @@ +use crate::config::paths::Paths; +use once_cell::sync::Lazy; +use std::fs; +use uuid::Uuid; + +static INSTANCE_ID: Lazy = Lazy::new(load_or_create); + +fn file_path() -> std::path::PathBuf { + Paths::state_dir().join("instance_id") +} + +fn load_or_create() -> String { + let path = file_path(); + + if let Ok(id) = fs::read_to_string(&path) { + let id = id.trim().to_string(); + if !id.is_empty() { + return id; + } + } + + let id = Uuid::new_v4().to_string(); + + if let Some(parent) = path.parent() { + let _ = fs::create_dir_all(parent); + } + let _ = fs::write(&path, &id); + + id +} + +/// Returns a stable, globally unique identifier for this Goose installation. +/// The ID is generated once and persisted to disk, surviving restarts. +pub fn get_instance_id() -> &'static str { + &INSTANCE_ID +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_instance_id_is_stable() { + let id1 = get_instance_id(); + let id2 = get_instance_id(); + assert_eq!(id1, id2); + assert!(!id1.is_empty()); + } +} diff --git a/crates/goose/src/lib.rs b/crates/goose/src/lib.rs index 8e05ebb3..cd60b4de 100644 --- a/crates/goose/src/lib.rs +++ b/crates/goose/src/lib.rs @@ -18,6 +18,7 @@ pub mod execution; pub mod gateway; pub mod goose_apps; pub mod hints; +pub mod instance_id; pub mod logging; pub mod mcp_utils; pub mod model; diff --git a/crates/goose/src/providers/databricks.rs b/crates/goose/src/providers/databricks.rs index 652ff5ad..dcbe3d11 100644 --- a/crates/goose/src/providers/databricks.rs +++ b/crates/goose/src/providers/databricks.rs @@ -29,6 +29,7 @@ use super::retry::ProviderRetry; use super::utils::{ImageFormat, RequestLog}; use crate::config::ConfigError; use crate::conversation::message::Message; +use crate::instance_id::get_instance_id; use crate::model::ModelConfig; use crate::providers::retry::{ RetryConfig, DEFAULT_BACKOFF_MULTIPLIER, DEFAULT_INITIAL_RETRY_INTERVAL_MS, @@ -132,6 +133,8 @@ pub struct DatabricksProvider { name: String, #[serde(skip)] token_cache: Arc>>, + #[serde(skip)] + instance_id: Option, } impl DatabricksProvider { @@ -186,6 +189,7 @@ impl DatabricksProvider { fast_retry_config, name: DATABRICKS_PROVIDER_NAME.to_string(), token_cache, + instance_id: Self::resolve_instance_id(), }; provider.model = model.with_fast(DATABRICKS_DEFAULT_FAST_MODEL, DATABRICKS_PROVIDER_NAME)?; @@ -249,9 +253,21 @@ impl DatabricksProvider { fast_retry_config: RetryConfig::new(0, 0, 1.0, 0), name: DATABRICKS_PROVIDER_NAME.to_string(), token_cache, + instance_id: Self::resolve_instance_id(), }) } + fn resolve_instance_id() -> Option { + let enabled = crate::config::Config::global() + .get_param::("GOOSE_DATABRICKS_CLIENT_REQUEST_ID") + .unwrap_or(false); + if enabled { + Some(get_instance_id().to_string()) + } else { + None + } + } + fn is_responses_model(model_name: &str) -> bool { let normalized = model_name.to_ascii_lowercase(); normalized.contains("codex") @@ -267,16 +283,31 @@ impl DatabricksProvider { } } + fn build_client_request_id(&self, session_id: &str) -> Option { + self.instance_id.as_ref().map(|instance_id| { + json!({ + "sessionId": format!("{}_{}", instance_id, session_id), + }) + .to_string() + }) + } + async fn post( &self, session_id: Option<&str>, - payload: Value, + mut payload: Value, model_name: Option<&str>, ) -> Result { let is_embedding = payload.get("input").is_some() && payload.get("messages").is_none(); let model_to_use = model_name.unwrap_or(&self.model.model_name); let path = self.get_endpoint_path(model_to_use, is_embedding); + if let Some(session_id) = session_id { + if let Some(client_request_id) = self.build_client_request_id(session_id) { + payload["client_request_id"] = Value::String(client_request_id); + } + } + let response = self .api_client .response_post(session_id, &path, &payload) @@ -341,10 +372,14 @@ impl Provider for DatabricksProvider { tools: &[Tool], ) -> Result { let path = self.get_endpoint_path(&model_config.model_name, false); + let client_request_id = self.build_client_request_id(session_id); if Self::is_responses_model(&model_config.model_name) { let mut payload = create_responses_request(model_config, system, messages, tools)?; payload["stream"] = Value::Bool(true); + if let Some(ref client_request_id) = client_request_id { + payload["client_request_id"] = Value::String(client_request_id.clone()); + } let mut log = RequestLog::start(model_config, &payload)?; @@ -383,6 +418,9 @@ impl Provider for DatabricksProvider { .as_object_mut() .expect("payload should have model key") .remove("model"); + if let Some(client_request_id) = client_request_id { + payload["client_request_id"] = Value::String(client_request_id); + } payload .as_object_mut()