Merge branch 'main' into feat/storybook-ci-tests

Author: Aaron1011

.github/workflows/general.yml

@@ -734,8 +734,8 @@ jobs:
     # the main 'validate' job)
     runs-on: ${{ matrix.replicated && 'namespace-profile-tensorzero-8x16;ephemeral-storage.size-multiplier=2' || 'namespace-profile-tensorzero-8x16' }}
     continue-on-error: ${{ matrix.clickhouse_version.allow_failure }}
-    # This needs to pull from Docker Hub, so skip for external PR CI
-    if: ${{ (github.event.pull_request.head.repo.full_name == github.repository) || inputs.is_merge_group }}
+    # This needs to pull from Docker Hub, so only run in the merge queue, or when running on a PR from the main repository
+    if: ${{ github.event_name == 'merge_group' || (github.event.pull_request.head.repo.full_name == github.repository) }}
     strategy:
       matrix:
         # Only include replicated: true when running in merge queue
@@ -1185,7 +1185,11 @@ jobs:
       ]
     runs-on: ubuntu-latest
     steps:
+      - name: Print all job results
+        run: |
+          echo "'needs': ${{ toJson(needs) }}"
+          echo "github.event_name: ${{ github.event_name }}"
       # When running in the merge queue, jobs should never be skipped.
       # In PR CI, some jobs may be intentionally skipped (e.g. due to running from a fork, or to save money)
-      - if: ${{ contains(needs.*.result, 'failure') || contains(needs.*.result, 'cancelled' || (github.event_name == 'merge_group' && contains(needs.*.result, 'skipped'))) }}
+      - if: ${{ contains(needs.*.result, 'failure') || contains(needs.*.result, 'cancelled') || (github.event_name == 'merge_group' && contains(needs.*.result, 'skipped')) }}
         run: exit 1

clients/rust/src/lib.rs

@@ -71,6 +71,7 @@ pub use tensorzero_core::endpoints::inference::{
     ChatCompletionInferenceParams, InferenceOutput, InferenceParams, InferenceResponse,
     InferenceResponseChunk, InferenceStream,
 };
+pub use tensorzero_core::endpoints::internal::config::GetConfigResponse;
 pub use tensorzero_core::endpoints::object_storage::ObjectResponse;
 pub use tensorzero_core::endpoints::stored_inferences::v1::types::{
     GetInferencesRequest, GetInferencesResponse, ListInferencesRequest,
@@ -467,6 +468,24 @@ pub trait ClientExt {
 
     fn get_config(&self) -> Result<Arc<Config>, TensorZeroError>;
 
+    /// Gets a config snapshot by hash, or the live config if no hash is provided.
+    ///
+    /// # Arguments
+    ///
+    /// * `hash` - Optional hash of the config snapshot to retrieve. If `None`, returns the live config.
+    ///
+    /// # Returns
+    ///
+    /// A `GetConfigResponse` containing the config snapshot.
+    ///
+    /// # Errors
+    ///
+    /// Returns a `TensorZeroError` if the request fails or the config snapshot is not found.
+    async fn get_config_snapshot(
+        &self,
+        hash: Option<&str>,
+    ) -> Result<GetConfigResponse, TensorZeroError>;
+
     #[cfg(any(feature = "e2e_tests", feature = "pyo3"))]
     fn get_app_state_data(&self) -> Option<&tensorzero_core::utils::gateway::AppStateData>;
 }
@@ -1323,6 +1342,60 @@ impl ClientExt for Client {
         }
     }
 
+    async fn get_config_snapshot(
+        &self,
+        hash: Option<&str>,
+    ) -> Result<GetConfigResponse, TensorZeroError> {
+        match self.mode() {
+            ClientMode::HTTPGateway(client) => {
+                let endpoint = match hash {
+                    Some(h) => format!("internal/config/{h}"),
+                    None => "internal/config".to_string(),
+                };
+                let url = client
+                    .base_url
+                    .join(&endpoint)
+                    .map_err(|e| TensorZeroError::Other {
+                        source: Error::new(ErrorDetails::InvalidBaseUrl {
+                            message: format!(
+                                "Failed to join base URL with /{endpoint} endpoint: {e}"
+                            ),
+                        })
+                        .into(),
+                    })?;
+                let builder = client.http_client.get(url);
+                Ok(client.send_and_parse_http_response(builder).await?.0)
+            }
+            ClientMode::EmbeddedGateway { gateway, timeout } => {
+                with_embedded_timeout(*timeout, async {
+                    use tensorzero_core::db::ConfigQueries;
+                    let snapshot_hash = match hash {
+                        Some(h) => h.parse().map_err(|_| {
+                            err_to_http(Error::new(ErrorDetails::ConfigSnapshotNotFound {
+                                snapshot_hash: h.to_string(),
+                            }))
+                        })?,
+                        None => gateway.handle.app_state.config.hash.clone(),
+                    };
+                    let snapshot = gateway
+                        .handle
+                        .app_state
+                        .clickhouse_connection_info
+                        .get_config_snapshot(snapshot_hash)
+                        .await
+                        .map_err(err_to_http)?;
+                    Ok(GetConfigResponse {
+                        hash: snapshot.hash.to_string(),
+                        config: snapshot.config,
+                        extra_templates: snapshot.extra_templates,
+                        tags: snapshot.tags,
+                    })
+                })
+                .await
+            }
+        }
+    }
+
     async fn get_variant_sampling_probabilities(
         &self,
         function_name: &str,

gateway/src/routes/internal.rs

@@ -99,4 +99,13 @@ pub fn build_internal_non_otel_enabled_routes() -> Router<AppStateData> {
             "/internal/models/latency",
             get(endpoints::internal::models::get_model_latency_handler),
         )
+        // Config snapshot endpoints
+        .route(
+            "/internal/config",
+            get(endpoints::internal::config::get_live_config_handler),
+        )
+        .route(
+            "/internal/config/{hash}",
+            get(endpoints::internal::config::get_config_by_hash_handler),
+        )
 }

tensorzero-core/src/config/snapshot.rs

@@ -76,6 +76,10 @@ pub struct ConfigSnapshot {
     /// uses the compiled templates in `Config.templates`, not these raw strings.
     pub extra_templates: HashMap<String, String>,
 
+    /// User-defined tags for categorizing or labeling this config snapshot.
+    /// Tags are metadata and do not affect the config hash.
+    pub tags: HashMap<String, String>,
+
     __private: (),
 }
 
@@ -105,7 +109,6 @@ impl SnapshotHash {
     }
 }
 
-#[cfg(any(test, feature = "e2e_tests"))]
 impl std::str::FromStr for SnapshotHash {
     type Err = std::convert::Infallible;
 
@@ -137,11 +140,13 @@ impl ConfigSnapshot {
                     message: format!("Failed to serialize stored config: {e}"),
                 })
             })?);
+
         let hash = ConfigSnapshot::hash(&stored_config_toml, &extra_templates)?;
         Ok(Self {
             config: stored_config,
             hash,
             extra_templates,
+            tags: HashMap::new(),
             __private: (),
         })
     }
@@ -169,6 +174,7 @@ impl ConfigSnapshot {
             config: StoredConfig::default(),
             hash: SnapshotHash::new_test(),
             extra_templates: HashMap::new(),
+            tags: HashMap::new(),
             __private: (),
         }
     }
@@ -177,9 +183,14 @@ impl ConfigSnapshot {
     ///
     /// This is used when loading a previously stored config snapshot from ClickHouse.
     /// The hash is recomputed from the config and templates to ensure consistency.
+    ///
+    /// Note: We deserialize as `StoredConfig` (not `UninitializedConfig`) to support
+    /// backward compatibility with historical snapshots that may contain deprecated
+    /// fields like `timeouts`.
     pub fn from_stored(
         config_toml: &str,
         extra_templates: HashMap<String, String>,
+        tags: HashMap<String, String>,
         original_hash: &SnapshotHash,
     ) -> Result<Self, Error> {
         let table: toml::Table = config_toml.parse().map_err(|e| {
@@ -189,19 +200,24 @@ impl ConfigSnapshot {
         })?;
 
         let sorted_table = prepare_table_for_snapshot(table);
-        let config = UninitializedConfig::try_from(sorted_table.clone())?;
-        let hash = ConfigSnapshot::hash(&sorted_table, &extra_templates)?;
-        if hash != *original_hash {
-            return Err(Error::new(ErrorDetails::ConfigSnapshotHashMismatch {
-                expected: original_hash.clone(),
-                actual: hash.clone(),
-            }));
-        }
 
+        // Deserialize as StoredConfig to accept deprecated fields (e.g., `timeouts`)
+        let stored_config: StoredConfig =
+            serde_path_to_error::deserialize(sorted_table).map_err(|e| {
+                let path = e.path().clone();
+                Error::new(ErrorDetails::Config {
+                    message: format!("{}: {}", path, e.into_inner().message()),
+                })
+            })?;
+
+        // Use the original hash from the database rather than recomputing it.
+        // Recomputing can produce different hashes due to floating-point serialization
+        // differences (e.g., 0.2 vs 0.20000000298023224) even when the config is identical.
         Ok(Self {
-            config: config.into(),
-            hash,
+            config: stored_config,
+            hash: original_hash.clone(),
             extra_templates,
+            tags,
             __private: (),
         })
     }

tensorzero-core/src/db/clickhouse/migration_manager/migrations/migration_0045.rs

@@ -0,0 +1,44 @@
+use super::check_column_exists;
+use crate::db::clickhouse::ClickHouseConnectionInfo;
+use crate::db::clickhouse::migration_manager::migration_trait::Migration;
+use crate::error::Error;
+use async_trait::async_trait;
+
+const MIGRATION_ID: &str = "0045";
+
+/// This migration adds a `tags` column to the `ConfigSnapshot` table.
+pub struct Migration0045<'a> {
+    pub clickhouse: &'a ClickHouseConnectionInfo,
+}
+
+#[async_trait]
+impl Migration for Migration0045<'_> {
+    async fn can_apply(&self) -> Result<(), Error> {
+        Ok(())
+    }
+
+    async fn should_apply(&self) -> Result<bool, Error> {
+        Ok(!check_column_exists(self.clickhouse, "ConfigSnapshot", "tags", MIGRATION_ID).await?)
+    }
+
+    async fn apply(&self, _clean_start: bool) -> Result<(), Error> {
+        let on_cluster_name = self.clickhouse.get_on_cluster_name();
+
+        self.clickhouse
+            .run_query_synchronous_no_params(format!(
+                "ALTER TABLE ConfigSnapshot{on_cluster_name} ADD COLUMN IF NOT EXISTS tags Map(String, String) DEFAULT map()"
+            ))
+            .await?;
+
+        Ok(())
+    }
+
+    fn rollback_instructions(&self) -> String {
+        let on_cluster_name = self.clickhouse.get_on_cluster_name();
+        format!("ALTER TABLE ConfigSnapshot{on_cluster_name} DROP COLUMN tags;")
+    }
+
+    async fn has_succeeded(&self) -> Result<bool, Error> {
+        Ok(check_column_exists(self.clickhouse, "ConfigSnapshot", "tags", MIGRATION_ID).await?)
+    }
+}

tensorzero-core/src/db/clickhouse/migration_manager/migrations/mod.rs

@@ -42,6 +42,7 @@ pub mod migration_0041;
 pub mod migration_0042;
 pub mod migration_0043;
 pub mod migration_0044;
+pub mod migration_0045;
 
 /// Returns true if the table exists, false if it does not
 /// Errors if the query fails

tensorzero-core/src/db/clickhouse/migration_manager/mod.rs

@@ -52,11 +52,12 @@ use migrations::migration_0041::Migration0041;
 use migrations::migration_0042::Migration0042;
 use migrations::migration_0043::Migration0043;
 use migrations::migration_0044::Migration0044;
+use migrations::migration_0045::Migration0045;
 use serde::{Deserialize, Serialize};
 
 /// This must match the number of migrations returned by `make_all_migrations` - the tests
 /// will panic if they don't match.
-pub const NUM_MIGRATIONS: usize = 38;
+pub const NUM_MIGRATIONS: usize = 39;
 pub fn get_run_migrations_command() -> String {
     let version = env!("CARGO_PKG_VERSION");
     format!(
@@ -126,6 +127,7 @@ pub fn make_all_migrations<'a>(
         Box::new(Migration0042 { clickhouse }),
         Box::new(Migration0043 { clickhouse }),
         Box::new(Migration0044 { clickhouse }),
+        Box::new(Migration0045 { clickhouse }),
     ];
     assert_eq!(
         migrations.len(),

tensorzero-core/src/db/clickhouse/mod.rs

@@ -340,11 +340,13 @@ impl ConfigQueries for ClickHouseConnectionInfo {
         struct ConfigSnapshotRow {
             config: String,
             extra_templates: HashMap<String, String>,
+            #[serde(default)]
+            tags: HashMap<String, String>,
         }
 
         let hash_str = snapshot_hash.to_string();
         let query = format!(
-            "SELECT config, extra_templates \
+            "SELECT config, extra_templates, tags \
              FROM ConfigSnapshot FINAL \
              WHERE hash = toUInt256('{hash_str}') \
              LIMIT 1 \
@@ -365,7 +367,7 @@ impl ConfigQueries for ClickHouseConnectionInfo {
             })
         })?;
 
-        ConfigSnapshot::from_stored(&row.config, row.extra_templates, &snapshot_hash)
+        ConfigSnapshot::from_stored(&row.config, row.extra_templates, row.tags, &snapshot_hash)
     }
 }
 

tensorzero-core/src/db/postgres/mod.rs

@@ -259,6 +259,6 @@ async fn get_applied_migrations(pool: &PgPool) -> Result<HashSet<i64>, sqlx::Err
     Ok(applied_migrations)
 }
 
-fn make_migrator() -> sqlx::migrate::Migrator {
+pub fn make_migrator() -> sqlx::migrate::Migrator {
     migrate!("src/db/postgres/migrations")
 }