Merge branch 'tensorzero:main' into main

Author: hongantran3804

.cargo/config.toml

@@ -38,11 +38,26 @@ test-e2e-no-creds = [
     "e2e_tests",
     "--profile",
     "clickhouse",
+    ]
+test-rate-limit-load = [
+    "run",
+    "--release",
+    "--package",
+    "rate-limit-load-test",
+    "--",
+]
+test-feedback-load = [
+    "run",
+    "--release",
+    "--package",
+    "feedback-load-test",
+    "--",
 ]
 
 
 build-e2e = "build --bin gateway --features e2e_tests"
 run-e2e = "run --bin gateway --features e2e_tests -- --config-file tensorzero-core/tests/e2e/tensorzero.toml"
+migrate-postgres = "run --bin gateway --features e2e_tests -- --run-postgres-migrations"
 watch-e2e = "watch -x run-e2e"
 
 tsbuild = [
@@ -52,4 +67,8 @@ tsbuild = [
     "tensorzero-core",
     "-p",
     "tensorzero",
+    "-p",
+    "tensorzero-node",
+    "-p",
+    "tensorzero-auth",
 ] # Export Typescript bindings for TensorZero

.claude/commands/clippy.md

@@ -0,0 +1 @@
+Run `cargo clippy --all-targets --no-default-features -- -D warnings` and fix all issues. Confirm that it is clean if necessary. Then run `cargo fmt`.

.config/nextest.toml

@@ -26,14 +26,14 @@ slow-timeout = { period = "20s", terminate-after = 3 }
 # We also define cargo aliases for common use-cases (e.g. `cargo test-e2e`)
 [profile.e2e]
 retries = { backoff = "exponential", count = 4, delay = "5s", jitter = true, max-delay = "60s" }
-default-filter = "not (test(batch)) and not (test(test_dummy_only) | test(clickhouse))"
+default-filter = "not (test(batch)) and not (test(test_dummy_only) | test(clickhouse) | test(db::))"
 junit.path = "junit.xml"
 
 [profile.batch]
 default-filter = 'test(batch)'
 
 [profile.clickhouse]
-default-filter = 'test(test_dummy_only) | test(clickhouse)'
+default-filter = 'test(test_dummy_only) | test(clickhouse) | test(db::)'
 junit.path = "junit.xml"
 
 [profile.optimization]
@@ -70,6 +70,10 @@ clickhouse_migration = { max-threads = 1 }
 [[profile.default.overrides]]
 filter = 'test(test_otel_export)'
 slow-timeout = { period = "60s", terminate-after = 2 }
+# We sleep long enough that these tests should pass 100% of the time.
+# Any flakiness indicates a potential bug that could cause us to lose spans,
+# so we don't want to retry.
+retries = 0
 
 [[profile.default.overrides]]
 filter = 'binary(e2e) and test(providers::aws_bedrock::)'
@@ -97,6 +101,16 @@ test-group = 'clickhouse_migration'
 filter = 'binary(e2e) and test(providers::vllm::)'
 slow-timeout = { period = "60s", terminate-after = 2 }
 
+[[profile.default.overrides]]
+# Web search is very slow
+filter = 'test(websearch)'
+slow-timeout = { period = "120s", terminate-after = 2 }
+
+# Image inference seems to be slow on GCP Vertex Gemini, so we give it a longer timeout
+[[profile.default.overrides]]
+filter = 'test(providers::gcp_vertex_gemini::test_image_inference)'
+slow-timeout = { period = "60s", terminate-after = 2 }
+
 [[profile.default.overrides]]
 filter = 'binary(e2e) and test(providers::sglang::)'
 # The model we run on SGLang often fails to emit valid tool calls, so we need many retries
@@ -122,7 +136,7 @@ slow-timeout = { period = "900s", terminate-after = 1 }
 [[profile.default.overrides]]
 filter = 'test(test_concurrent_requests)'
 # This test creates and uses lots of `reqwest::Client`s, which can be slow
-slow-timeout = { period = "30s", terminate-after = 1 }
+slow-timeout = { period = "60s", terminate-after = 1 }
 
 [[profile.default.overrides]]
 # Settings for running unit tests

.dockerignore

@@ -9,6 +9,8 @@ target-rust-analyzer
 **/.env*
 **/.vscode
 **/.claude
+ui/fixtures/*.jsonl
+ui/fixtures/s3-fixtures
 /recipes
 /examples
 gcp_jwt_key.json

.github/dependabot.yml

@@ -33,7 +33,11 @@ updates:
       interval: "weekly"
     labels: []
     groups:
+      opentelemetry-dependencies:
+        patterns:
+          - "*opentelemetry*"
+          - "tonic"
       rust-dependencies:
         patterns:
           - "*"
-    open-pull-requests-limit: 1 # This does not affect the security update PR limit
+    open-pull-requests-limit: 2 # This does not affect the security update PR limit

.github/workflows/batch-test.yml

@@ -37,11 +37,14 @@ jobs:
     runs-on: ubuntu-latest
     if: github.repository == 'tensorzero/tensorzero'
 
-    timeout-minutes: 15
+    timeout-minutes: 60
 
     steps:
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683
 
+      - name: Cleanup disk space
+        run: ./ci/free-disk-space.sh
+
       - name: Install Rust toolchain
         run: |
           for attempt in 1 2 3; do

.github/workflows/build-gateway-container.yml

@@ -5,31 +5,40 @@ on:
 
 jobs:
   build-gateway-container:
-    runs-on: namespace-profile-tensorzero-2x8;overrides.cache-tag=build-gateway-cache
+    runs-on: ubuntu-latest
     if: github.repository == 'tensorzero/tensorzero'
+    permissions:
+      # Permission to checkout the repository
+      contents: read
+      # Permission to fetch GitHub OIDC token authentication
+      id-token: write
 
     steps:
       # TODO - investigate why using the Namespace checkout action causes
       # 'tensorzero_core::built_info::GIT_COMMIT_HASH_SHORT' to be `None`
       - name: Check out the repo
         uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683
 
+      # We allow the namespace builder setup to fail on Dependabot PRs and PRs from forks
+      # (where the oidc token is not available)
+
+      - name: Install Namespace CLI
+        uses: namespacelabs/nscloud-setup@d1c625762f7c926a54bd39252efff0705fd11c64
+        continue-on-error: ${{ github.event.pull_request.head.repo.full_name != github.repository || github.actor == 'dependabot[bot]' }}
+
+      - name: Configure Namespace-powered Buildx
+        uses: namespacelabs/nscloud-setup-buildx-action@91c2e6537780e3b092cb8476406be99a8f91bd5e
+        with:
+          wait-for-builder: true
+        continue-on-error: ${{ github.event.pull_request.head.repo.full_name != github.repository || github.actor == 'dependabot[bot]' }}
+
       - name: Build `gateway` container
         run: |
-          docker build --build-arg BUILDKIT_CONTEXT_KEEP_GIT_DIR=1 -f gateway/Dockerfile . -t tensorzero/gateway:sha-${{ github.sha }}
+          docker buildx build --build-arg BUILDKIT_CONTEXT_KEEP_GIT_DIR=1 -f gateway/Dockerfile . -t tensorzero/gateway:sha-${{ github.sha }}
 
       - name: Save `gateway` container
         run: docker save tensorzero/gateway:sha-${{ github.sha }} > gateway-container.tar
 
-      - name: Upload `gateway` container as an artifact to Namespace
-        uses: namespace-actions/upload-artifact@9a78c62e083914789d908952f9773e42744b9f68
-        with:
-          name: build-gateway-container
-          path: gateway-container.tar
-          retention-days: 1
-          if-no-files-found: error
-          overwrite: false
-
       - name: Upload `gateway` container as an artifact to Github
         uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02
         with:

.github/workflows/build-gateway-e2e-container.yml

@@ -0,0 +1,62 @@
+name: Build Gateway Container
+
+on:
+  workflow_call:
+    secrets:
+      DOCKERHUB_USERNAME:
+        required: true
+      DOCKERHUB_TOKEN:
+        required: true
+
+env:
+  TAG: tensorzero/gateway-e2e:sha-${{ github.sha }}
+  LATEST_TAG: tensorzero/gateway-e2e:latest
+
+jobs:
+  build-gateway-container:
+    runs-on: ubuntu-latest
+    if: github.repository == 'tensorzero/tensorzero'
+    permissions:
+      # Permission to checkout the repository
+      contents: read
+      # Permission to fetch GitHub OIDC token authentication
+      id-token: write
+
+    steps:
+      # TODO - investigate why using the Namespace checkout action causes
+      # 'tensorzero_core::built_info::GIT_COMMIT_HASH_SHORT' to be `None`
+      - name: Check out the repo
+        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683
+
+      - name: Login to DockerHub
+        uses: docker/login-action@74a5d142397b4f367a81961eba4e8cd7edddf772
+        with:
+          username: ${{ secrets.DOCKERHUB_USERNAME }}
+          password: ${{ secrets.DOCKERHUB_TOKEN }}
+
+      # We allow the namespace builder setup to fail on Dependabot PRs and PRs from forks
+      # (where the oidc token is not available)
+
+      - name: Install Namespace CLI
+        uses: namespacelabs/nscloud-setup@d1c625762f7c926a54bd39252efff0705fd11c64
+        continue-on-error: ${{ github.event.pull_request.head.repo.full_name != github.repository || github.actor == 'dependabot[bot]' }}
+
+      - name: Configure Namespace-powered Buildx
+        uses: namespacelabs/nscloud-setup-buildx-action@91c2e6537780e3b092cb8476406be99a8f91bd5e
+        with:
+          wait-for-builder: true
+        continue-on-error: ${{ github.event.pull_request.head.repo.full_name != github.repository || github.actor == 'dependabot[bot]' }}
+
+      - name: Pull latest image for caching (ignore errors if image doesn't exist)
+        run: docker pull $LATEST_TAG || true
+
+      - name: Build `gateway-e2e` container
+        run: docker build --load --build-arg BUILDKIT_CONTEXT_KEEP_GIT_DIR=1 --build-arg PROFILE=release --build-arg CARGO_BUILD_FLAGS="--features e2e_tests" --cache-from $LATEST_TAG -f gateway/Dockerfile . -t $TAG
+
+      - name: Tag latest image and push both tags
+        run: |
+          docker tag $TAG $LATEST_TAG
+          echo "Pushing $TAG"
+          docker push $TAG
+          echo "Pushing $LATEST_TAG"
+          docker push $LATEST_TAG

.github/workflows/build-ui-container.yml

@@ -5,28 +5,40 @@ on:
 
 jobs:
   build-ui-container:
-    runs-on: namespace-profile-tensorzero-2x8;overrides.cache-tag=build-ui-cache
+    runs-on: ubuntu-latest
     if: github.repository == 'tensorzero/tensorzero'
+    permissions:
+      # Permission to checkout the repository
+      contents: read
+      # Permission to fetch GitHub OIDC token authentication
+      id-token: write
 
     steps:
-      - uses: namespacelabs/nscloud-checkout-action@953fed31a6113cc2347ca69c9d823743c65bc84b
+      # TODO - investigate why using the Namespace checkout action causes
+      # 'tensorzero_core::built_info::GIT_COMMIT_HASH_SHORT' to be `None`
+      - name: Check out the repo
+        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683
+
+      # We allow the namespace builder setup to fail on Dependabot PRs and PRs from forks
+      # (where the oidc token is not available)
+
+      - name: Install Namespace CLI
+        uses: namespacelabs/nscloud-setup@d1c625762f7c926a54bd39252efff0705fd11c64
+        continue-on-error: ${{ github.event.pull_request.head.repo.full_name != github.repository || github.actor == 'dependabot[bot]' }}
+
+      - name: Configure Namespace-powered Buildx
+        uses: namespacelabs/nscloud-setup-buildx-action@91c2e6537780e3b092cb8476406be99a8f91bd5e
+        with:
+          wait-for-builder: true
+        continue-on-error: ${{ github.event.pull_request.head.repo.full_name != github.repository || github.actor == 'dependabot[bot]' }}
 
       - name: Build `ui` container
         run: |
-          docker build --build-arg BUILDKIT_CONTEXT_KEEP_GIT_DIR=1 -f ui/Dockerfile . -t tensorzero/ui:sha-${{ github.sha }}
+          docker buildx build --build-arg BUILDKIT_CONTEXT_KEEP_GIT_DIR=1 -f ui/Dockerfile . -t tensorzero/ui:sha-${{ github.sha }}
 
       - name: Save `ui` container
         run: docker save tensorzero/ui:sha-${{ github.sha }} > ui-container.tar
 
-      - name: Upload `ui` container as an artifact to Namespace
-        uses: namespace-actions/upload-artifact@9a78c62e083914789d908952f9773e42744b9f68
-        with:
-          name: build-ui-container
-          path: ui-container.tar
-          retention-days: 1
-          if-no-files-found: error
-          overwrite: false
-
       - name: Upload `ui` container as an artifact
         uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02
         with:

.github/workflows/buildkite-merge-queue.yml

@@ -5,6 +5,8 @@ on:
   # merge_group:
   workflow_dispatch:
 
+permissions: {}
+
 jobs:
   merge-queue-tests:
     runs-on: ubuntu-latest