chore: clean up stop signal redis fence (#7119)

Co-authored-by: Raunak Bhagat <r@rabh.io>
Co-authored-by: acaprau <48705707+acaprau@users.noreply.github.com>
Co-authored-by: Justin Tahara <105671973+justin-tahara@users.noreply.github.com>

.github/CODEOWNERS

@@ -1,3 +1,10 @@
 * @onyx-dot-app/onyx-core-team
 # Helm charts Owners
 /helm/ @justin-tahara
+
+# Web standards updates
+/web/STANDARDS.md @raunakab @Weves
+
+# Agent context files
+/CLAUDE.md.template @Weves
+/AGENTS.md.template @Weves

.github/actions/setup-python-and-install-dependencies/action.yml

@@ -7,12 +7,6 @@ inputs:
 runs:
   using: "composite"
   steps:
-    - name: Setup uv
-      uses: astral-sh/setup-uv@caf0cab7a618c569241d31dcd442f54681755d39 # ratchet:astral-sh/setup-uv@v3
-      # TODO: Enable caching once there is a uv.lock file checked in.
-      # with:
-      #   enable-cache: true
-
     - name: Compute requirements hash
       id: req-hash
       shell: bash
@@ -28,6 +22,8 @@ runs:
         done <<< "$REQUIREMENTS"
         echo "hash=$(echo "$hash" | sha256sum | cut -d' ' -f1)" >> "$GITHUB_OUTPUT"
 
+    # NOTE: This comes before Setup uv since clean-ups run in reverse chronological order
+    # such that Setup uv's prune-cache is able to prune the cache before we upload.
     - name: Cache uv cache directory
       uses: runs-on/cache@50350ad4242587b6c8c2baa2e740b1bc11285ff4 # ratchet:runs-on/cache@v4
       with:
@@ -36,6 +32,14 @@ runs:
         restore-keys: |
           ${{ runner.os }}-uv-
 
+    - name: Setup uv
+      uses: astral-sh/setup-uv@ed21f2f24f8dd64503750218de024bcf64c7250a # ratchet:astral-sh/setup-uv@v7
+      with:
+        version: "0.9.9"
+      # TODO: Enable caching once there is a uv.lock file checked in.
+      # with:
+      #   enable-cache: true
+
     - name: Setup Python
       uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # ratchet:actions/setup-python@v5
       with:

.github/dependabot.yml

@@ -5,7 +5,7 @@ updates:
     schedule:
       interval: "weekly"
     cooldown:
-      default-days: 4
+      default-days: 7
     open-pull-requests-limit: 3
     assignees:
       - "jmelahman"
@@ -16,7 +16,7 @@ updates:
     schedule:
       interval: "weekly"
     cooldown:
-      default-days: 4
+      default-days: 7
     open-pull-requests-limit: 3
     assignees:
       - "jmelahman"

.github/pull_request_template.md

@@ -1,10 +1,10 @@
 ## Description
 
-[Provide a brief description of the changes in this PR]
+<!--- Provide a brief description of the changes in this PR --->
 
 ## How Has This Been Tested?
 
-[Describe the tests you ran to verify your changes]
+<!--- Describe the tests you ran to verify your changes --->
 
 ## Additional Options
 

.github/workflows/check-lazy-imports.yml

@@ -1,33 +0,0 @@
-name: Check Lazy Imports
-concurrency:
-  group: Check-Lazy-Imports-${{ github.workflow }}-${{ github.head_ref || github.event.workflow_run.head_branch || github.run_id }}
-  cancel-in-progress: true
-
-on:
-  merge_group:
-  pull_request:
-    branches:
-      - main
-      - 'release/**'
-
-permissions:
-  contents: read
-
-jobs:
-  check-lazy-imports:
-    runs-on: ubuntu-latest
-    timeout-minutes: 45
-
-    steps:
-    - name: Checkout code
-      uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # ratchet:actions/checkout@v6
-      with:
-        persist-credentials: false
-
-    - name: Set up Python
-      uses: actions/setup-python@e797f83bcb11b83ae66e0230d6156d7c80228e7c # ratchet:actions/setup-python@v6
-      with:
-        python-version: '3.11'
-
-    - name: Check lazy imports
-      run: python3 backend/scripts/check_lazy_imports.py

.github/workflows/deployment.yml

@@ -6,8 +6,9 @@ on:
       - "*"
   workflow_dispatch:
 
-permissions:
-  contents: read
+# Set restrictive default permissions for all jobs. Jobs that need more permissions
+# should explicitly declare them.
+permissions: {}
 
 env:
   IS_DRY_RUN: ${{ github.event_name == 'workflow_dispatch' }}
@@ -20,6 +21,7 @@ jobs:
     runs-on: ubuntu-slim
     timeout-minutes: 90
     outputs:
+      build-desktop: ${{ steps.check.outputs.build-desktop }}
       build-web: ${{ steps.check.outputs.build-web }}
       build-web-cloud: ${{ steps.check.outputs.build-web-cloud }}
       build-backend: ${{ steps.check.outputs.build-backend }}
@@ -30,6 +32,7 @@ jobs:
       is-stable-standalone: ${{ steps.check.outputs.is-stable-standalone }}
       is-beta-standalone: ${{ steps.check.outputs.is-beta-standalone }}
       sanitized-tag: ${{ steps.check.outputs.sanitized-tag }}
+      short-sha: ${{ steps.check.outputs.short-sha }}
     steps:
       - name: Check which components to build and version info
         id: check
@@ -38,6 +41,7 @@ jobs:
           # Sanitize tag name by replacing slashes with hyphens (for Docker tag compatibility)
           SANITIZED_TAG=$(echo "$TAG" | tr '/' '-')
           IS_CLOUD=false
+          BUILD_DESKTOP=false
           BUILD_WEB=false
           BUILD_WEB_CLOUD=false
           BUILD_BACKEND=true
@@ -47,13 +51,6 @@ jobs:
           IS_STABLE_STANDALONE=false
           IS_BETA_STANDALONE=false
 
-          if [[ "$TAG" == *cloud* ]]; then
-            IS_CLOUD=true
-            BUILD_WEB_CLOUD=true
-          else
-            BUILD_WEB=true
-          fi
-
           # Version checks (for web - any stable version)
           if [[ "$TAG" =~ ^v[0-9]+\.[0-9]+\.[0-9]+$ ]]; then
             IS_STABLE=true
@@ -62,6 +59,17 @@ jobs:
             IS_BETA=true
           fi
 
+          if [[ "$TAG" == *cloud* ]]; then
+            IS_CLOUD=true
+            BUILD_WEB_CLOUD=true
+          else
+            BUILD_WEB=true
+            # Skip desktop builds on beta tags and nightly runs
+            if [[ "$IS_BETA" != "true" ]] && [[ "$TAG" != nightly* ]]; then
+              BUILD_DESKTOP=true
+            fi
+          fi
+
           # Version checks (for backend/model-server - stable version excluding cloud tags)
           if [[ "$TAG" =~ ^v[0-9]+\.[0-9]+\.[0-9]+$ ]] && [[ "$TAG" != *cloud* ]]; then
             IS_STABLE_STANDALONE=true
@@ -70,7 +78,9 @@ jobs:
             IS_BETA_STANDALONE=true
           fi
 
+          SHORT_SHA="${GITHUB_SHA::7}"
           {
+            echo "build-desktop=$BUILD_DESKTOP"
             echo "build-web=$BUILD_WEB"
             echo "build-web-cloud=$BUILD_WEB_CLOUD"
             echo "build-backend=$BUILD_BACKEND"
@@ -81,8 +91,181 @@ jobs:
             echo "is-stable-standalone=$IS_STABLE_STANDALONE"
             echo "is-beta-standalone=$IS_BETA_STANDALONE"
             echo "sanitized-tag=$SANITIZED_TAG"
+            echo "short-sha=$SHORT_SHA"
           } >> "$GITHUB_OUTPUT"
 
+  check-version-tag:
+    runs-on: ubuntu-slim
+    timeout-minutes: 10
+    if: ${{ !startsWith(github.ref_name, 'nightly-latest') && github.event_name != 'workflow_dispatch' }}
+    steps:
+      - name: Checkout
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6
+        with:
+          persist-credentials: false
+          fetch-depth: 0
+
+      - name: Setup uv
+        uses: astral-sh/setup-uv@ed21f2f24f8dd64503750218de024bcf64c7250a # ratchet:astral-sh/setup-uv@v7
+        with:
+          version: "0.9.9"
+          # NOTE: This isn't caching much and zizmor suggests this could be poisoned, so disable.
+          enable-cache: false
+
+      - name: Validate tag is versioned correctly
+        run: |
+          uv run --no-sync --with release-tag tag --check
+
+  notify-slack-on-tag-check-failure:
+    needs:
+      - check-version-tag
+    if: always() && needs.check-version-tag.result == 'failure' && github.event_name != 'workflow_dispatch'
+    runs-on: ubuntu-slim
+    timeout-minutes: 10
+    steps:
+      - name: Checkout
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6
+        with:
+          persist-credentials: false
+
+      - name: Send Slack notification
+        uses: ./.github/actions/slack-notify
+        with:
+          webhook-url: ${{ secrets.MONITOR_DEPLOYMENTS_WEBHOOK }}
+          failed-jobs: "• check-version-tag"
+          title: "🚨 Version Tag Check Failed"
+          ref-name: ${{ github.ref_name }}
+
+  build-desktop:
+    needs: determine-builds
+    if: needs.determine-builds.outputs.build-desktop == 'true'
+    permissions:
+      contents: write
+      actions: read
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - platform: "macos-latest" # Build a universal image for macOS.
+            args: "--target universal-apple-darwin"
+          - platform: "ubuntu-24.04"
+            args: "--bundles deb,rpm"
+          - platform: "ubuntu-24.04-arm" # Only available in public repos.
+            args: "--bundles deb,rpm"
+          - platform: "windows-latest"
+            args: ""
+
+    runs-on: ${{ matrix.platform }}
+    timeout-minutes: 90
+    steps:
+      - uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6.0.1
+        with:
+          # NOTE: persist-credentials is needed for tauri-action to create GitHub releases.
+          persist-credentials: true # zizmor: ignore[artipacked]
+
+      - name: install dependencies (ubuntu only)
+        if: startsWith(matrix.platform, 'ubuntu-')
+        run: |
+          sudo apt-get update
+          sudo apt-get install -y \
+            build-essential \
+            libglib2.0-dev \
+            libgirepository1.0-dev \
+            libgtk-3-dev \
+            libjavascriptcoregtk-4.1-dev \
+            libwebkit2gtk-4.1-dev \
+            libayatana-appindicator3-dev \
+            gobject-introspection \
+            pkg-config \
+            curl \
+            xdg-utils
+
+      - name: setup node
+        uses: actions/setup-node@395ad3262231945c25e8478fd5baf05154b1d79f # ratchet:actions/setup-node@v6.1.0
+        with:
+          node-version: 24
+          package-manager-cache: false
+
+      - name: install Rust stable
+        uses: dtolnay/rust-toolchain@6d9817901c499d6b02debbb57edb38d33daa680b # zizmor: ignore[impostor-commit]
+        with:
+          # Those targets are only used on macos runners so it's in an `if` to slightly speed up windows and linux builds.
+          targets: ${{ matrix.platform == 'macos-latest' && 'aarch64-apple-darwin,x86_64-apple-darwin' || '' }}
+
+      - name: install frontend dependencies
+        working-directory: ./desktop
+        run: npm install
+
+      - name: Inject version (Unix)
+        if: runner.os != 'Windows'
+        working-directory: ./desktop
+        env:
+          SHORT_SHA: ${{ needs.determine-builds.outputs.short-sha }}
+          EVENT_NAME: ${{ github.event_name }}
+        run: |
+          if [ "${EVENT_NAME}" == "workflow_dispatch" ]; then
+            VERSION="0.0.0-dev+${SHORT_SHA}"
+          else
+            VERSION="${GITHUB_REF_NAME#v}"
+          fi
+          echo "Injecting version: $VERSION"
+
+          # Update Cargo.toml
+          sed "s/^version = .*/version = \"$VERSION\"/" src-tauri/Cargo.toml > src-tauri/Cargo.toml.tmp
+          mv src-tauri/Cargo.toml.tmp src-tauri/Cargo.toml
+
+          # Update tauri.conf.json
+          jq --arg v "$VERSION" '.version = $v' src-tauri/tauri.conf.json > src-tauri/tauri.conf.json.tmp
+          mv src-tauri/tauri.conf.json.tmp src-tauri/tauri.conf.json
+
+          # Update package.json
+          jq --arg v "$VERSION" '.version = $v' package.json > package.json.tmp
+          mv package.json.tmp package.json
+
+          echo "Versions set to: $VERSION"
+
+      - name: Inject version (Windows)
+        if: runner.os == 'Windows'
+        working-directory: ./desktop
+        shell: pwsh
+        run: |
+          # Windows MSI requires numeric-only build metadata, so we skip the SHA suffix
+          if ("${{ github.event_name }}" -eq "workflow_dispatch") {
+            $VERSION = "0.0.0"
+          } else {
+            # Strip 'v' prefix and any pre-release suffix (e.g., -beta.13) for MSI compatibility
+            $VERSION = "$env:GITHUB_REF_NAME" -replace '^v', '' -replace '-.*$', ''
+          }
+          Write-Host "Injecting version: $VERSION"
+
+          # Update Cargo.toml
+          $cargo = Get-Content src-tauri/Cargo.toml -Raw
+          $cargo = $cargo -replace '(?m)^version = .*', "version = `"$VERSION`""
+          Set-Content src-tauri/Cargo.toml $cargo -NoNewline
+
+          # Update tauri.conf.json
+          $json = Get-Content src-tauri/tauri.conf.json | ConvertFrom-Json
+          $json.version = $VERSION
+          $json | ConvertTo-Json -Depth 100 | Set-Content src-tauri/tauri.conf.json
+
+          # Update package.json
+          $pkg = Get-Content package.json | ConvertFrom-Json
+          $pkg.version = $VERSION
+          $pkg | ConvertTo-Json -Depth 100 | Set-Content package.json
+
+          Write-Host "Versions set to: $VERSION"
+
+      - uses: tauri-apps/tauri-action@19b93bb55601e3e373a93cfb6eb4242e45f5af20 # ratchet:tauri-apps/tauri-action@action-v0.6.0
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        with:
+          tagName: ${{ github.event_name != 'workflow_dispatch' && 'v__VERSION__' || format('v0.0.0-dev+{0}', needs.determine-builds.outputs.short-sha) }}
+          releaseName: ${{ github.event_name != 'workflow_dispatch' && 'v__VERSION__' || format('v0.0.0-dev+{0}', needs.determine-builds.outputs.short-sha) }}
+          releaseBody: "See the assets to download this version and install."
+          releaseDraft: true
+          prerelease: false
+          args: ${{ matrix.args }}
+
   build-web-amd64:
     needs: determine-builds
     if: needs.determine-builds.outputs.build-web == 'true'
@@ -100,13 +283,13 @@ jobs:
       - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2
 
       - name: Checkout
-        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # ratchet:actions/checkout@v6
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6
         with:
           persist-credentials: false
 
       - name: Docker meta
         id: meta
-        uses: docker/metadata-action@318604b99e75e41977312d83839a89be02ca4893 # ratchet:docker/metadata-action@v5
+        uses: docker/metadata-action@c299e40c65443455700f0fdfc63efafe5b349051 # ratchet:docker/metadata-action@v5
         with:
           images: ${{ github.event_name == 'workflow_dispatch' && env.RUNS_ON_ECR_CACHE || env.REGISTRY_IMAGE }}
           flavor: |
@@ -158,13 +341,13 @@ jobs:
       - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2
 
       - name: Checkout
-        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # ratchet:actions/checkout@v6
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6
         with:
           persist-credentials: false
 
       - name: Docker meta
         id: meta
-        uses: docker/metadata-action@318604b99e75e41977312d83839a89be02ca4893 # ratchet:docker/metadata-action@v5
+        uses: docker/metadata-action@c299e40c65443455700f0fdfc63efafe5b349051 # ratchet:docker/metadata-action@v5
         with:
           images: ${{ github.event_name == 'workflow_dispatch' && env.RUNS_ON_ECR_CACHE || env.REGISTRY_IMAGE }}
           flavor: |
@@ -226,7 +409,7 @@ jobs:
 
       - name: Docker meta
         id: meta
-        uses: docker/metadata-action@318604b99e75e41977312d83839a89be02ca4893 # ratchet:docker/metadata-action@v5
+        uses: docker/metadata-action@c299e40c65443455700f0fdfc63efafe5b349051 # ratchet:docker/metadata-action@v5
         with:
           images: ${{ github.event_name == 'workflow_dispatch' && env.RUNS_ON_ECR_CACHE || env.REGISTRY_IMAGE }}
           flavor: |
@@ -266,13 +449,13 @@ jobs:
       - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2
 
       - name: Checkout
-        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # ratchet:actions/checkout@v6
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6
         with:
           persist-credentials: false
 
       - name: Docker meta
         id: meta
-        uses: docker/metadata-action@318604b99e75e41977312d83839a89be02ca4893 # ratchet:docker/metadata-action@v5
+        uses: docker/metadata-action@c299e40c65443455700f0fdfc63efafe5b349051 # ratchet:docker/metadata-action@v5
         with:
           images: ${{ github.event_name == 'workflow_dispatch' && env.RUNS_ON_ECR_CACHE || env.REGISTRY_IMAGE }}
           flavor: |
@@ -332,13 +515,13 @@ jobs:
       - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2
 
       - name: Checkout
-        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # ratchet:actions/checkout@v6
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6
         with:
           persist-credentials: false
 
       - name: Docker meta
         id: meta
-        uses: docker/metadata-action@318604b99e75e41977312d83839a89be02ca4893 # ratchet:docker/metadata-action@v5
+        uses: docker/metadata-action@c299e40c65443455700f0fdfc63efafe5b349051 # ratchet:docker/metadata-action@v5
         with:
           images: ${{ github.event_name == 'workflow_dispatch' && env.RUNS_ON_ECR_CACHE || env.REGISTRY_IMAGE }}
           flavor: |
@@ -408,7 +591,7 @@ jobs:
 
       - name: Docker meta
         id: meta
-        uses: docker/metadata-action@318604b99e75e41977312d83839a89be02ca4893 # ratchet:docker/metadata-action@v5
+        uses: docker/metadata-action@c299e40c65443455700f0fdfc63efafe5b349051 # ratchet:docker/metadata-action@v5
         with:
           images: ${{ github.event_name == 'workflow_dispatch' && env.RUNS_ON_ECR_CACHE || env.REGISTRY_IMAGE }}
           flavor: |
@@ -445,13 +628,13 @@ jobs:
       - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2
 
       - name: Checkout code
-        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # ratchet:actions/checkout@v6
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6
         with:
           persist-credentials: false
 
       - name: Docker meta
         id: meta
-        uses: docker/metadata-action@318604b99e75e41977312d83839a89be02ca4893 # ratchet:docker/metadata-action@v5
+        uses: docker/metadata-action@c299e40c65443455700f0fdfc63efafe5b349051 # ratchet:docker/metadata-action@v5
         with:
           images: ${{ github.event_name == 'workflow_dispatch' && env.RUNS_ON_ECR_CACHE || env.REGISTRY_IMAGE }}
           flavor: |
@@ -502,13 +685,13 @@ jobs:
       - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2
 
       - name: Checkout code
-        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # ratchet:actions/checkout@v6
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6
         with:
           persist-credentials: false
 
       - name: Docker meta
         id: meta
-        uses: docker/metadata-action@318604b99e75e41977312d83839a89be02ca4893 # ratchet:docker/metadata-action@v5
+        uses: docker/metadata-action@c299e40c65443455700f0fdfc63efafe5b349051 # ratchet:docker/metadata-action@v5
         with:
           images: ${{ github.event_name == 'workflow_dispatch' && env.RUNS_ON_ECR_CACHE || env.REGISTRY_IMAGE }}
           flavor: |
@@ -569,7 +752,7 @@ jobs:
 
       - name: Docker meta
         id: meta
-        uses: docker/metadata-action@318604b99e75e41977312d83839a89be02ca4893 # ratchet:docker/metadata-action@v5
+        uses: docker/metadata-action@c299e40c65443455700f0fdfc63efafe5b349051 # ratchet:docker/metadata-action@v5
         with:
           images: ${{ github.event_name == 'workflow_dispatch' && env.RUNS_ON_ECR_CACHE || env.REGISTRY_IMAGE }}
           flavor: |
@@ -610,13 +793,13 @@ jobs:
       - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2
 
       - name: Checkout code
-        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # ratchet:actions/checkout@v6
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6
         with:
           persist-credentials: false
 
       - name: Docker meta
         id: meta
-        uses: docker/metadata-action@318604b99e75e41977312d83839a89be02ca4893 # ratchet:docker/metadata-action@v5
+        uses: docker/metadata-action@c299e40c65443455700f0fdfc63efafe5b349051 # ratchet:docker/metadata-action@v5
         with:
           images: ${{ github.event_name == 'workflow_dispatch' && env.RUNS_ON_ECR_CACHE || env.REGISTRY_IMAGE }}
           flavor: |
@@ -674,13 +857,13 @@ jobs:
       - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2
 
       - name: Checkout code
-        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # ratchet:actions/checkout@v6
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6
         with:
           persist-credentials: false
 
       - name: Docker meta
         id: meta
-        uses: docker/metadata-action@318604b99e75e41977312d83839a89be02ca4893 # ratchet:docker/metadata-action@v5
+        uses: docker/metadata-action@c299e40c65443455700f0fdfc63efafe5b349051 # ratchet:docker/metadata-action@v5
         with:
           images: ${{ github.event_name == 'workflow_dispatch' && env.RUNS_ON_ECR_CACHE || env.REGISTRY_IMAGE }}
           flavor: |
@@ -747,7 +930,7 @@ jobs:
 
       - name: Docker meta
         id: meta
-        uses: docker/metadata-action@318604b99e75e41977312d83839a89be02ca4893 # ratchet:docker/metadata-action@v5
+        uses: docker/metadata-action@c299e40c65443455700f0fdfc63efafe5b349051 # ratchet:docker/metadata-action@v5
         with:
           images: ${{ github.event_name == 'workflow_dispatch' && env.RUNS_ON_ECR_CACHE || env.REGISTRY_IMAGE }}
           flavor: |
@@ -867,7 +1050,7 @@ jobs:
       - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2
 
       - name: Checkout
-        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # ratchet:actions/checkout@v6
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6
         with:
           persist-credentials: false
 
@@ -939,6 +1122,7 @@ jobs:
 
   notify-slack-on-failure:
     needs:
+      - build-desktop
       - build-web-amd64
       - build-web-arm64
       - merge-web
@@ -951,13 +1135,13 @@ jobs:
       - build-model-server-amd64
       - build-model-server-arm64
       - merge-model-server
-    if: always() && (needs.build-web-amd64.result == 'failure' || needs.build-web-arm64.result == 'failure' || needs.merge-web.result == 'failure' || needs.build-web-cloud-amd64.result == 'failure' || needs.build-web-cloud-arm64.result == 'failure' || needs.merge-web-cloud.result == 'failure' || needs.build-backend-amd64.result == 'failure' || needs.build-backend-arm64.result == 'failure' || needs.merge-backend.result == 'failure' || needs.build-model-server-amd64.result == 'failure' || needs.build-model-server-arm64.result == 'failure' || needs.merge-model-server.result == 'failure') && github.event_name != 'workflow_dispatch'
+    if: always() && (needs.build-desktop.result == 'failure' || needs.build-web-amd64.result == 'failure' || needs.build-web-arm64.result == 'failure' || needs.merge-web.result == 'failure' || needs.build-web-cloud-amd64.result == 'failure' || needs.build-web-cloud-arm64.result == 'failure' || needs.merge-web-cloud.result == 'failure' || needs.build-backend-amd64.result == 'failure' || needs.build-backend-arm64.result == 'failure' || needs.merge-backend.result == 'failure' || needs.build-model-server-amd64.result == 'failure' || needs.build-model-server-arm64.result == 'failure' || needs.merge-model-server.result == 'failure') && github.event_name != 'workflow_dispatch'
     # NOTE: Github-hosted runners have about 20s faster queue times and are preferred here.
     runs-on: ubuntu-slim
     timeout-minutes: 90
     steps:
       - name: Checkout
-        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # ratchet:actions/checkout@v6
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6
         with:
           persist-credentials: false
 
@@ -966,6 +1150,9 @@ jobs:
         shell: bash
         run: |
           FAILED_JOBS=""
+          if [ "${NEEDS_BUILD_DESKTOP_RESULT}" == "failure" ]; then
+            FAILED_JOBS="${FAILED_JOBS}• build-desktop\\n"
+          fi
           if [ "${NEEDS_BUILD_WEB_AMD64_RESULT}" == "failure" ]; then
             FAILED_JOBS="${FAILED_JOBS}• build-web-amd64\\n"
           fi
@@ -1006,6 +1193,7 @@ jobs:
           FAILED_JOBS=$(printf '%s' "$FAILED_JOBS" | sed 's/\\n$//')
           echo "jobs=$FAILED_JOBS" >> "$GITHUB_OUTPUT"
         env:
+          NEEDS_BUILD_DESKTOP_RESULT: ${{ needs.build-desktop.result }}
           NEEDS_BUILD_WEB_AMD64_RESULT: ${{ needs.build-web-amd64.result }}
           NEEDS_BUILD_WEB_ARM64_RESULT: ${{ needs.build-web-arm64.result }}
           NEEDS_MERGE_WEB_RESULT: ${{ needs.merge-web.result }}

.github/workflows/helm-chart-releases.yml

@@ -15,7 +15,7 @@ jobs:
     timeout-minutes: 45
     steps:
       - name: Checkout
-        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # ratchet:actions/checkout@v6
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6
         with:
           fetch-depth: 0
           persist-credentials: false
@@ -32,6 +32,7 @@ jobs:
           helm repo add cloudnative-pg https://cloudnative-pg.github.io/charts
           helm repo add ot-container-kit https://ot-container-kit.github.io/helm-charts
           helm repo add minio https://charts.min.io/
+          helm repo add code-interpreter https://onyx-dot-app.github.io/code-interpreter/
           helm repo update
 
       - name: Build chart dependencies

.github/workflows/merge-group.yml

@@ -0,0 +1,31 @@
+name: Merge Group-Specific
+
+on:
+  merge_group:
+
+permissions:
+  contents: read
+
+jobs:
+  # This job immediately succeeds to satisfy branch protection rules on merge_group events.
+  # There is a similarly named "required" job in pr-integration-tests.yml which runs the actual
+  # integration tests. That job runs on both pull_request and merge_group events, and this job
+  # exists solely to provide a fast-passing check with the same name for branch protection.
+  # The actual tests remain enforced on presubmit (pull_request events).
+  required:
+    runs-on: ubuntu-latest
+    timeout-minutes: 45
+    steps:
+      - name: Success
+        run: echo "Success"
+  # This job immediately succeeds to satisfy branch protection rules on merge_group events.
+  # There is a similarly named "playwright-required" job in pr-playwright-tests.yml which runs
+  # the actual playwright tests. That job runs on both pull_request and merge_group events, and
+  # this job exists solely to provide a fast-passing check with the same name for branch protection.
+  # The actual tests remain enforced on presubmit (pull_request events).
+  playwright-required:
+    runs-on: ubuntu-latest
+    timeout-minutes: 45
+    steps:
+      - name: Success
+        run: echo "Success"

.github/workflows/nightly-scan-licenses.yml

@@ -28,12 +28,12 @@ jobs:
 
     steps:
       - name: Checkout code
-        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # ratchet:actions/checkout@v6
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6
         with:
           persist-credentials: false
 
       - name: Set up Python
-        uses: actions/setup-python@e797f83bcb11b83ae66e0230d6156d7c80228e7c # ratchet:actions/setup-python@v6
+        uses: actions/setup-python@83679a892e2d95755f2dac6acb0bfd1e9ac5d548 # ratchet:actions/setup-python@v6
         with:
           python-version: '3.11'
           cache: 'pip'

.github/workflows/pr-database-tests.yml

@@ -0,0 +1,62 @@
+name: Database Tests
+concurrency:
+  group: Database-Tests-${{ github.workflow }}-${{ github.head_ref || github.event.workflow_run.head_branch || github.run_id }}
+  cancel-in-progress: true
+
+on:
+  merge_group:
+  pull_request:
+    branches:
+      - main
+      - "release/**"
+  push:
+    tags:
+      - "v*.*.*"
+
+permissions:
+  contents: read
+
+jobs:
+  database-tests:
+    runs-on:
+      - runs-on
+      - runner=2cpu-linux-arm64
+      - "run-id=${{ github.run_id }}-database-tests"
+    timeout-minutes: 45
+    steps:
+      - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2
+
+      - name: Checkout code
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6
+        with:
+          persist-credentials: false
+
+      - name: Setup Python and Install Dependencies
+        uses: ./.github/actions/setup-python-and-install-dependencies
+        with:
+          requirements: |
+            backend/requirements/default.txt
+            backend/requirements/dev.txt
+
+      - name: Generate OpenAPI schema and Python client
+        shell: bash
+        run: |
+          ods openapi all
+
+      # needed for pulling external images otherwise, we hit the "Unauthenticated users" limit
+      # https://docs.docker.com/docker-hub/usage/
+      - name: Login to Docker Hub
+        uses: docker/login-action@5e57cd118135c172c3672efd75eb46360885c0ef # ratchet:docker/login-action@v3
+        with:
+          username: ${{ secrets.DOCKER_USERNAME }}
+          password: ${{ secrets.DOCKER_TOKEN }}
+
+      - name: Start Docker containers
+        working-directory: ./deployment/docker_compose
+        run: |
+          docker compose -f docker-compose.yml -f docker-compose.dev.yml up -d \
+            relational_db
+
+      - name: Run Database Tests
+        working-directory: ./backend
+        run: pytest -m alembic tests/integration/tests/migrations/

.github/workflows/pr-external-dependency-unit-tests.yml

@@ -7,6 +7,9 @@ on:
   merge_group:
   pull_request:
     branches: [main]
+  push:
+    tags:
+      - "v*.*.*"
 
 permissions:
   contents: read
@@ -29,6 +32,9 @@ env:
   CONFLUENCE_ACCESS_TOKEN: ${{ secrets.CONFLUENCE_ACCESS_TOKEN }}
   CONFLUENCE_ACCESS_TOKEN_SCOPED: ${{ secrets.CONFLUENCE_ACCESS_TOKEN_SCOPED }}
 
+  # Jira
+  JIRA_ADMIN_API_TOKEN: ${{ secrets.JIRA_ADMIN_API_TOKEN }}
+
   # LLMs
   OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
   ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
@@ -46,7 +52,7 @@ jobs:
       test-dirs: ${{ steps.set-matrix.outputs.test-dirs }}
     steps:
       - name: Checkout code
-        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # ratchet:actions/checkout@v6
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6
         with:
           persist-credentials: false
 
@@ -74,12 +80,13 @@ jobs:
     env:
       PYTHONPATH: ./backend
       MODEL_SERVER_HOST: "disabled"
+      DISABLE_TELEMETRY: "true"
 
     steps:
       - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2
 
       - name: Checkout code
-        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # ratchet:actions/checkout@v6
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6
         with:
           persist-credentials: false
 
@@ -107,6 +114,7 @@ jobs:
         run: |
           cat <<EOF > deployment/docker_compose/.env
           CODE_INTERPRETER_BETA_ENABLED=true
+          DISABLE_TELEMETRY=true
           EOF
 
       - name: Set up Standard Dependencies
@@ -162,7 +170,7 @@ jobs:
 
       - name: Upload Docker logs
         if: failure()
-        uses: actions/upload-artifact@v4
+        uses: actions/upload-artifact@330a01c490aca151604b8cf639adc76d48f6c5d4 # ratchet:actions/upload-artifact@v5
         with:
           name: docker-logs-${{ matrix.test-dir }}
           path: docker-logs/

.github/workflows/pr-helm-chart-testing.yml

@@ -7,6 +7,9 @@ on:
   merge_group:
   pull_request:
     branches: [ main ]
+  push:
+    tags:
+      - "v*.*.*"
   workflow_dispatch:  # Allows manual triggering
 
 permissions:
@@ -21,7 +24,7 @@ jobs:
     # fetch-depth 0 is required for helm/chart-testing-action
     steps:
     - name: Checkout code
-      uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # ratchet:actions/checkout@v6
+      uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6
       with:
         fetch-depth: 0
         persist-credentials: false
@@ -80,6 +83,7 @@ jobs:
         helm repo add cloudnative-pg https://cloudnative-pg.github.io/charts
         helm repo add ot-container-kit https://ot-container-kit.github.io/helm-charts
         helm repo add minio https://charts.min.io/
+        helm repo add code-interpreter https://onyx-dot-app.github.io/code-interpreter/
         helm repo update
 
     - name: Install Redis operator

.github/workflows/pr-integration-tests.yml

@@ -9,6 +9,9 @@ on:
     branches:
       - main
       - "release/**"
+  push:
+    tags:
+      - "v*.*.*"
 
 permissions:
   contents: read
@@ -30,6 +33,11 @@ env:
   PERM_SYNC_SHAREPOINT_CERTIFICATE_PASSWORD: ${{ secrets.PERM_SYNC_SHAREPOINT_CERTIFICATE_PASSWORD }}
   PERM_SYNC_SHAREPOINT_DIRECTORY_ID: ${{ secrets.PERM_SYNC_SHAREPOINT_DIRECTORY_ID }}
   EXA_API_KEY: ${{ secrets.EXA_API_KEY }}
+  GITHUB_PERMISSION_SYNC_TEST_ACCESS_TOKEN: ${{ secrets.ONYX_GITHUB_PERMISSION_SYNC_TEST_ACCESS_TOKEN }}
+  GITHUB_PERMISSION_SYNC_TEST_ACCESS_TOKEN_CLASSIC: ${{ secrets.ONYX_GITHUB_PERMISSION_SYNC_TEST_ACCESS_TOKEN_CLASSIC }}
+  GITHUB_ADMIN_EMAIL: ${{ secrets.ONYX_GITHUB_ADMIN_EMAIL }}
+  GITHUB_TEST_USER_1_EMAIL: ${{ secrets.ONYX_GITHUB_TEST_USER_1_EMAIL }}
+  GITHUB_TEST_USER_2_EMAIL: ${{ secrets.ONYX_GITHUB_TEST_USER_2_EMAIL }}
 
 jobs:
   discover-test-dirs:
@@ -40,7 +48,7 @@ jobs:
       test-dirs: ${{ steps.set-matrix.outputs.test-dirs }}
     steps:
       - name: Checkout code
-        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # ratchet:actions/checkout@v6
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6
         with:
           persist-credentials: false
 
@@ -48,7 +56,7 @@ jobs:
         id: set-matrix
         run: |
           # Find all leaf-level directories in both test directories
-          tests_dirs=$(find backend/tests/integration/tests -mindepth 1 -maxdepth 1 -type d ! -name "__pycache__" -exec basename {} \; | sort)
+          tests_dirs=$(find backend/tests/integration/tests -mindepth 1 -maxdepth 1 -type d ! -name "__pycache__" ! -name "mcp" -exec basename {} \; | sort)
           connector_dirs=$(find backend/tests/integration/connector_job_tests -mindepth 1 -maxdepth 1 -type d ! -name "__pycache__" -exec basename {} \; | sort)
 
           # Create JSON array with directory info
@@ -64,17 +72,36 @@ jobs:
           all_dirs="[${all_dirs%,}]"
           echo "test-dirs=$all_dirs" >> $GITHUB_OUTPUT
 
-
   build-backend-image:
-    runs-on: [runs-on, runner=1cpu-linux-arm64, "run-id=${{ github.run_id }}-build-backend-image", "extras=ecr-cache"]
+    runs-on:
+      [
+        runs-on,
+        runner=1cpu-linux-arm64,
+        "run-id=${{ github.run_id }}-build-backend-image",
+        "extras=ecr-cache",
+      ]
     timeout-minutes: 45
     steps:
       - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2
       - name: Checkout code
-        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # ratchet:actions/checkout@v6
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6
         with:
           persist-credentials: false
 
+      - name: Format branch name for cache
+        id: format-branch
+        env:
+          PR_NUMBER: ${{ github.event.pull_request.number }}
+          REF_NAME: ${{ github.ref_name }}
+        run: |
+          if [ -n "${PR_NUMBER}" ]; then
+            CACHE_SUFFIX="${PR_NUMBER}"
+          else
+            # shellcheck disable=SC2001
+            CACHE_SUFFIX=$(echo "${REF_NAME}" | sed 's/[^A-Za-z0-9._-]/-/g')
+          fi
+          echo "cache-suffix=${CACHE_SUFFIX}" >> $GITHUB_OUTPUT
+
       - name: Set up Docker Buildx
         uses: docker/setup-buildx-action@e468171a9de216ec08956ac3ada2f0791b6bd435 # ratchet:docker/setup-buildx-action@v3
 
@@ -95,23 +122,46 @@ jobs:
           push: true
           tags: ${{ env.RUNS_ON_ECR_CACHE }}:integration-test-backend-test-${{ github.run_id }}
           cache-from: |
+            type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:backend-cache-${{ github.event.pull_request.head.sha || github.sha }}
+            type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:backend-cache-${{ steps.format-branch.outputs.cache-suffix }}
             type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:backend-cache
             type=registry,ref=onyxdotapp/onyx-backend:latest
           cache-to: |
+            type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:backend-cache-${{ github.event.pull_request.head.sha || github.sha }},mode=max
+            type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:backend-cache-${{ steps.format-branch.outputs.cache-suffix }},mode=max
             type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:backend-cache,mode=max
           no-cache: ${{ vars.DOCKER_NO_CACHE == 'true' }}
 
-
   build-model-server-image:
-    runs-on: [runs-on, runner=1cpu-linux-arm64, "run-id=${{ github.run_id }}-build-model-server-image", "extras=ecr-cache"]
+    runs-on:
+      [
+        runs-on,
+        runner=1cpu-linux-arm64,
+        "run-id=${{ github.run_id }}-build-model-server-image",
+        "extras=ecr-cache",
+      ]
     timeout-minutes: 45
     steps:
       - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2
       - name: Checkout code
-        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # ratchet:actions/checkout@v6
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6
         with:
           persist-credentials: false
 
+      - name: Format branch name for cache
+        id: format-branch
+        env:
+          PR_NUMBER: ${{ github.event.pull_request.number }}
+          REF_NAME: ${{ github.ref_name }}
+        run: |
+          if [ -n "${PR_NUMBER}" ]; then
+            CACHE_SUFFIX="${PR_NUMBER}"
+          else
+            # shellcheck disable=SC2001
+            CACHE_SUFFIX=$(echo "${REF_NAME}" | sed 's/[^A-Za-z0-9._-]/-/g')
+          fi
+          echo "cache-suffix=${CACHE_SUFFIX}" >> $GITHUB_OUTPUT
+
       - name: Set up Docker Buildx
         uses: docker/setup-buildx-action@e468171a9de216ec08956ac3ada2f0791b6bd435 # ratchet:docker/setup-buildx-action@v3
 
@@ -132,18 +182,28 @@ jobs:
           push: true
           tags: ${{ env.RUNS_ON_ECR_CACHE }}:integration-test-model-server-test-${{ github.run_id }}
           cache-from: |
+            type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:model-server-cache-${{ github.event.pull_request.head.sha || github.sha }}
+            type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:model-server-cache-${{ steps.format-branch.outputs.cache-suffix }}
             type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:model-server-cache
             type=registry,ref=onyxdotapp/onyx-model-server:latest
-          cache-to: type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:model-server-cache,mode=max
-
+          cache-to: |
+            type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:model-server-cache-${{ github.event.pull_request.head.sha || github.sha }},mode=max
+            type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:model-server-cache-${{ steps.format-branch.outputs.cache-suffix }},mode=max
+            type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:model-server-cache,mode=max
 
   build-integration-image:
-    runs-on: [runs-on, runner=2cpu-linux-arm64, "run-id=${{ github.run_id }}-build-integration-image", "extras=ecr-cache"]
+    runs-on:
+      [
+        runs-on,
+        runner=2cpu-linux-arm64,
+        "run-id=${{ github.run_id }}-build-integration-image",
+        "extras=ecr-cache",
+      ]
     timeout-minutes: 45
     steps:
       - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2
       - name: Checkout code
-        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # ratchet:actions/checkout@v6
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6
         with:
           persist-credentials: false
 
@@ -159,16 +219,40 @@ jobs:
           username: ${{ secrets.DOCKER_USERNAME }}
           password: ${{ secrets.DOCKER_TOKEN }}
 
+      - name: Format branch name for cache
+        id: format-branch
+        env:
+          PR_NUMBER: ${{ github.event.pull_request.number }}
+          REF_NAME: ${{ github.ref_name }}
+        run: |
+          if [ -n "${PR_NUMBER}" ]; then
+            CACHE_SUFFIX="${PR_NUMBER}"
+          else
+            # shellcheck disable=SC2001
+            CACHE_SUFFIX=$(echo "${REF_NAME}" | sed 's/[^A-Za-z0-9._-]/-/g')
+          fi
+          echo "cache-suffix=${CACHE_SUFFIX}" >> $GITHUB_OUTPUT
+
       - name: Build and push integration test image with Docker Bake
         env:
           INTEGRATION_REPOSITORY: ${{ env.RUNS_ON_ECR_CACHE }}
           TAG: integration-test-${{ github.run_id }}
+          CACHE_SUFFIX: ${{ steps.format-branch.outputs.cache-suffix }}
+          HEAD_SHA: ${{ github.event.pull_request.head.sha || github.sha }}
         run: |
-          cd backend && docker buildx bake --push \
+          docker buildx bake --push \
+            --set backend.cache-from=type=registry,ref=${RUNS_ON_ECR_CACHE}:backend-cache-${HEAD_SHA} \
+            --set backend.cache-from=type=registry,ref=${RUNS_ON_ECR_CACHE}:backend-cache-${CACHE_SUFFIX} \
             --set backend.cache-from=type=registry,ref=${RUNS_ON_ECR_CACHE}:backend-cache \
             --set backend.cache-from=type=registry,ref=onyxdotapp/onyx-backend:latest \
+            --set backend.cache-to=type=registry,ref=${RUNS_ON_ECR_CACHE}:backend-cache-${HEAD_SHA},mode=max \
+            --set backend.cache-to=type=registry,ref=${RUNS_ON_ECR_CACHE}:backend-cache-${CACHE_SUFFIX},mode=max \
             --set backend.cache-to=type=registry,ref=${RUNS_ON_ECR_CACHE}:backend-cache,mode=max \
+            --set integration.cache-from=type=registry,ref=${RUNS_ON_ECR_CACHE}:integration-cache-${HEAD_SHA} \
+            --set integration.cache-from=type=registry,ref=${RUNS_ON_ECR_CACHE}:integration-cache-${CACHE_SUFFIX} \
             --set integration.cache-from=type=registry,ref=${RUNS_ON_ECR_CACHE}:integration-cache \
+            --set integration.cache-to=type=registry,ref=${RUNS_ON_ECR_CACHE}:integration-cache-${HEAD_SHA},mode=max \
+            --set integration.cache-to=type=registry,ref=${RUNS_ON_ECR_CACHE}:integration-cache-${CACHE_SUFFIX},mode=max \
             --set integration.cache-to=type=registry,ref=${RUNS_ON_ECR_CACHE}:integration-cache,mode=max \
             integration
 
@@ -195,7 +279,7 @@ jobs:
     steps:
       - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2
       - name: Checkout code
-        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # ratchet:actions/checkout@v6
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6
         with:
           persist-credentials: false
 
@@ -210,23 +294,27 @@ jobs:
 
       # NOTE: Use pre-ping/null pool to reduce flakiness due to dropped connections
       # NOTE: don't need web server for integration tests
-      - name: Start Docker containers
+      - name: Create .env file for Docker Compose
         env:
           ECR_CACHE: ${{ env.RUNS_ON_ECR_CACHE }}
           RUN_ID: ${{ github.run_id }}
+        run: |
+          cat <<EOF > deployment/docker_compose/.env
+          ENABLE_PAID_ENTERPRISE_EDITION_FEATURES=true
+          AUTH_TYPE=basic
+          POSTGRES_POOL_PRE_PING=true
+          POSTGRES_USE_NULL_POOL=true
+          REQUIRE_EMAIL_VERIFICATION=false
+          DISABLE_TELEMETRY=true
+          ONYX_BACKEND_IMAGE=${ECR_CACHE}:integration-test-backend-test-${RUN_ID}
+          ONYX_MODEL_SERVER_IMAGE=${ECR_CACHE}:integration-test-model-server-test-${RUN_ID}
+          INTEGRATION_TESTS_MODE=true
+          CHECK_TTL_MANAGEMENT_TASK_FREQUENCY_IN_HOURS=0.001
+          EOF
+
+      - name: Start Docker containers
         run: |
           cd deployment/docker_compose
-          ENABLE_PAID_ENTERPRISE_EDITION_FEATURES=true \
-          AUTH_TYPE=basic \
-          POSTGRES_POOL_PRE_PING=true \
-          POSTGRES_USE_NULL_POOL=true \
-          REQUIRE_EMAIL_VERIFICATION=false \
-          DISABLE_TELEMETRY=true \
-          ONYX_BACKEND_IMAGE=${ECR_CACHE}:integration-test-backend-test-${RUN_ID} \
-          ONYX_MODEL_SERVER_IMAGE=${ECR_CACHE}:integration-test-model-server-test-${RUN_ID} \
-          INTEGRATION_TESTS_MODE=true \
-          CHECK_TTL_MANAGEMENT_TASK_FREQUENCY_IN_HOURS=0.001 \
-          MCP_SERVER_ENABLED=true \
           docker compose -f docker-compose.yml -f docker-compose.dev.yml up \
             relational_db \
             index \
@@ -235,7 +323,6 @@ jobs:
             api_server \
             inference_model_server \
             indexing_model_server \
-            mcp_server \
             background \
             -d
         id: start_docker
@@ -278,12 +365,6 @@ jobs:
           }
 
           wait_for_service "http://localhost:8080/health" "API server"
-          test_dir="${{ matrix.test-dir.path }}"
-          if [ "$test_dir" = "tests/mcp" ]; then
-            wait_for_service "http://localhost:8090/health" "MCP server"
-          else
-            echo "Skipping MCP server wait for non-MCP suite: $test_dir"
-          fi
           echo "Finished waiting for services."
 
       - name: Start Mock Services
@@ -313,8 +394,6 @@ jobs:
               -e VESPA_HOST=index \
               -e REDIS_HOST=cache \
               -e API_SERVER_HOST=api_server \
-              -e MCP_SERVER_HOST=mcp_server \
-              -e MCP_SERVER_PORT=8090 \
               -e OPENAI_API_KEY=${OPENAI_API_KEY} \
               -e EXA_API_KEY=${EXA_API_KEY} \
               -e SLACK_BOT_TOKEN=${SLACK_BOT_TOKEN} \
@@ -330,6 +409,11 @@ jobs:
               -e PERM_SYNC_SHAREPOINT_PRIVATE_KEY="${PERM_SYNC_SHAREPOINT_PRIVATE_KEY}" \
               -e PERM_SYNC_SHAREPOINT_CERTIFICATE_PASSWORD=${PERM_SYNC_SHAREPOINT_CERTIFICATE_PASSWORD} \
               -e PERM_SYNC_SHAREPOINT_DIRECTORY_ID=${PERM_SYNC_SHAREPOINT_DIRECTORY_ID} \
+              -e GITHUB_PERMISSION_SYNC_TEST_ACCESS_TOKEN=${GITHUB_PERMISSION_SYNC_TEST_ACCESS_TOKEN} \
+              -e GITHUB_PERMISSION_SYNC_TEST_ACCESS_TOKEN_CLASSIC=${GITHUB_PERMISSION_SYNC_TEST_ACCESS_TOKEN_CLASSIC} \
+              -e GITHUB_ADMIN_EMAIL=${GITHUB_ADMIN_EMAIL} \
+              -e GITHUB_TEST_USER_1_EMAIL=${GITHUB_TEST_USER_1_EMAIL} \
+              -e GITHUB_TEST_USER_2_EMAIL=${GITHUB_TEST_USER_2_EMAIL} \
               -e TEST_WEB_HOSTNAME=test-runner \
               -e MOCK_CONNECTOR_SERVER_HOST=mock_connector_server \
               -e MOCK_CONNECTOR_SERVER_PORT=8001 \
@@ -358,21 +442,22 @@ jobs:
           path: ${{ github.workspace }}/docker-compose.log
       # ------------------------------------------------------------
 
-
   multitenant-tests:
     needs:
+      [build-backend-image, build-model-server-image, build-integration-image]
+    runs-on:
       [
-        build-backend-image,
-        build-model-server-image,
-        build-integration-image,
+        runs-on,
+        runner=8cpu-linux-arm64,
+        "run-id=${{ github.run_id }}-multitenant-tests",
+        "extras=ecr-cache",
       ]
-    runs-on: [runs-on, runner=8cpu-linux-arm64, "run-id=${{ github.run_id }}-multitenant-tests", "extras=ecr-cache"]
     timeout-minutes: 45
 
     steps:
       - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2
       - name: Checkout code
-        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # ratchet:actions/checkout@v6
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6
         with:
           persist-credentials: false
 
@@ -396,7 +481,6 @@ jobs:
           ONYX_BACKEND_IMAGE=${ECR_CACHE}:integration-test-backend-test-${RUN_ID} \
           ONYX_MODEL_SERVER_IMAGE=${ECR_CACHE}:integration-test-model-server-test-${RUN_ID} \
           DEV_MODE=true \
-          MCP_SERVER_ENABLED=true \
           docker compose -f docker-compose.multitenant-dev.yml up \
             relational_db \
             index \
@@ -405,7 +489,6 @@ jobs:
             api_server \
             inference_model_server \
             indexing_model_server \
-            mcp_server \
             background \
             -d
         id: start_docker_multi_tenant
@@ -454,8 +537,6 @@ jobs:
             -e VESPA_HOST=index \
             -e REDIS_HOST=cache \
             -e API_SERVER_HOST=api_server \
-            -e MCP_SERVER_HOST=mcp_server \
-            -e MCP_SERVER_PORT=8090 \
             -e OPENAI_API_KEY=${OPENAI_API_KEY} \
             -e EXA_API_KEY=${EXA_API_KEY} \
             -e SLACK_BOT_TOKEN=${SLACK_BOT_TOKEN} \

.github/workflows/pr-jest-tests.yml

@@ -3,7 +3,15 @@ concurrency:
   group: Run-Jest-Tests-${{ github.workflow }}-${{ github.head_ref || github.event.workflow_run.head_branch || github.run_id }}
   cancel-in-progress: true
 
-on: push
+on:
+  merge_group:
+  pull_request:
+    branches:
+      - main
+      - "release/**"
+  push:
+    tags:
+      - "v*.*.*"
 
 permissions:
   contents: read
@@ -15,15 +23,15 @@ jobs:
     timeout-minutes: 45
     steps:
       - name: Checkout code
-        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # ratchet:actions/checkout@v6
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6
         with:
           persist-credentials: false
 
       - name: Setup node
-        uses: actions/setup-node@2028fbc5c25fe9cf00d9f06a71cc4710d4507903 # ratchet:actions/setup-node@v4
+        uses: actions/setup-node@395ad3262231945c25e8478fd5baf05154b1d79f # ratchet:actions/setup-node@v4
         with:
           node-version: 22
-          cache: 'npm'
+          cache: "npm"
           cache-dependency-path: ./web/package-lock.json
 
       - name: Install node dependencies

.github/workflows/pr-mit-integration-tests.yml

@@ -6,6 +6,9 @@ concurrency:
 on:
   merge_group:
     types: [checks_requested]
+  push:
+    tags:
+      - "v*.*.*"
 
 permissions:
   contents: read
@@ -37,7 +40,7 @@ jobs:
       test-dirs: ${{ steps.set-matrix.outputs.test-dirs }}
     steps:
       - name: Checkout code
-        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # ratchet:actions/checkout@v6
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6
         with:
           persist-credentials: false
 
@@ -45,7 +48,7 @@ jobs:
         id: set-matrix
         run: |
           # Find all leaf-level directories in both test directories
-          tests_dirs=$(find backend/tests/integration/tests -mindepth 1 -maxdepth 1 -type d ! -name "__pycache__" -exec basename {} \; | sort)
+          tests_dirs=$(find backend/tests/integration/tests -mindepth 1 -maxdepth 1 -type d ! -name "__pycache__" ! -name "mcp" -exec basename {} \; | sort)
           connector_dirs=$(find backend/tests/integration/connector_job_tests -mindepth 1 -maxdepth 1 -type d ! -name "__pycache__" -exec basename {} \; | sort)
 
           # Create JSON array with directory info
@@ -62,15 +65,35 @@ jobs:
           echo "test-dirs=$all_dirs" >> $GITHUB_OUTPUT
 
   build-backend-image:
-    runs-on: [runs-on, runner=1cpu-linux-arm64, "run-id=${{ github.run_id }}-build-backend-image", "extras=ecr-cache"]
+    runs-on:
+      [
+        runs-on,
+        runner=1cpu-linux-arm64,
+        "run-id=${{ github.run_id }}-build-backend-image",
+        "extras=ecr-cache",
+      ]
     timeout-minutes: 45
     steps:
       - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2
       - name: Checkout code
-        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # ratchet:actions/checkout@v6
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6
         with:
           persist-credentials: false
 
+      - name: Format branch name for cache
+        id: format-branch
+        env:
+          PR_NUMBER: ${{ github.event.pull_request.number }}
+          REF_NAME: ${{ github.ref_name }}
+        run: |
+          if [ -n "${PR_NUMBER}" ]; then
+            CACHE_SUFFIX="${PR_NUMBER}"
+          else
+            # shellcheck disable=SC2001
+            CACHE_SUFFIX=$(echo "${REF_NAME}" | sed 's/[^A-Za-z0-9._-]/-/g')
+          fi
+          echo "cache-suffix=${CACHE_SUFFIX}" >> $GITHUB_OUTPUT
+
       - name: Set up Docker Buildx
         uses: docker/setup-buildx-action@e468171a9de216ec08956ac3ada2f0791b6bd435 # ratchet:docker/setup-buildx-action@v3
 
@@ -91,21 +114,46 @@ jobs:
           push: true
           tags: ${{ env.RUNS_ON_ECR_CACHE }}:integration-test-backend-test-${{ github.run_id }}
           cache-from: |
+            type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:backend-cache-${{ github.event.pull_request.head.sha || github.sha }}
+            type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:backend-cache-${{ steps.format-branch.outputs.cache-suffix }}
             type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:backend-cache
             type=registry,ref=onyxdotapp/onyx-backend:latest
-          cache-to: type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:backend-cache,mode=max
+          cache-to: |
+            type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:backend-cache-${{ github.event.pull_request.head.sha || github.sha }},mode=max
+            type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:backend-cache-${{ steps.format-branch.outputs.cache-suffix }},mode=max
+            type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:backend-cache,mode=max
           no-cache: ${{ vars.DOCKER_NO_CACHE == 'true' }}
 
   build-model-server-image:
-    runs-on: [runs-on, runner=1cpu-linux-arm64, "run-id=${{ github.run_id }}-build-model-server-image", "extras=ecr-cache"]
+    runs-on:
+      [
+        runs-on,
+        runner=1cpu-linux-arm64,
+        "run-id=${{ github.run_id }}-build-model-server-image",
+        "extras=ecr-cache",
+      ]
     timeout-minutes: 45
     steps:
       - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2
       - name: Checkout code
-        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # ratchet:actions/checkout@v6
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6
         with:
           persist-credentials: false
 
+      - name: Format branch name for cache
+        id: format-branch
+        env:
+          PR_NUMBER: ${{ github.event.pull_request.number }}
+          REF_NAME: ${{ github.ref_name }}
+        run: |
+          if [ -n "${PR_NUMBER}" ]; then
+            CACHE_SUFFIX="${PR_NUMBER}"
+          else
+            # shellcheck disable=SC2001
+            CACHE_SUFFIX=$(echo "${REF_NAME}" | sed 's/[^A-Za-z0-9._-]/-/g')
+          fi
+          echo "cache-suffix=${CACHE_SUFFIX}" >> $GITHUB_OUTPUT
+
       - name: Set up Docker Buildx
         uses: docker/setup-buildx-action@e468171a9de216ec08956ac3ada2f0791b6bd435 # ratchet:docker/setup-buildx-action@v3
 
@@ -126,20 +174,45 @@ jobs:
           push: true
           tags: ${{ env.RUNS_ON_ECR_CACHE }}:integration-test-model-server-test-${{ github.run_id }}
           cache-from: |
+            type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:model-server-cache-${{ github.event.pull_request.head.sha || github.sha }}
+            type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:model-server-cache-${{ steps.format-branch.outputs.cache-suffix }}
             type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:model-server-cache
             type=registry,ref=onyxdotapp/onyx-model-server:latest
-          cache-to: type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:model-server-cache,mode=max
+          cache-to: |
+            type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:model-server-cache-${{ github.event.pull_request.head.sha || github.sha }},mode=max
+            type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:model-server-cache-${{ steps.format-branch.outputs.cache-suffix }},mode=max
+            type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:model-server-cache,mode=max
 
   build-integration-image:
-    runs-on: [runs-on, runner=2cpu-linux-arm64, "run-id=${{ github.run_id }}-build-integration-image", "extras=ecr-cache"]
+    runs-on:
+      [
+        runs-on,
+        runner=2cpu-linux-arm64,
+        "run-id=${{ github.run_id }}-build-integration-image",
+        "extras=ecr-cache",
+      ]
     timeout-minutes: 45
     steps:
       - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2
       - name: Checkout code
-        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # ratchet:actions/checkout@v6
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6
         with:
           persist-credentials: false
 
+      - name: Format branch name for cache
+        id: format-branch
+        env:
+          PR_NUMBER: ${{ github.event.pull_request.number }}
+          REF_NAME: ${{ github.ref_name }}
+        run: |
+          if [ -n "${PR_NUMBER}" ]; then
+            CACHE_SUFFIX="${PR_NUMBER}"
+          else
+            # shellcheck disable=SC2001
+            CACHE_SUFFIX=$(echo "${REF_NAME}" | sed 's/[^A-Za-z0-9._-]/-/g')
+          fi
+          echo "cache-suffix=${CACHE_SUFFIX}" >> $GITHUB_OUTPUT
+
       - name: Set up Docker Buildx
         uses: docker/setup-buildx-action@e468171a9de216ec08956ac3ada2f0791b6bd435 # ratchet:docker/setup-buildx-action@v3
 
@@ -156,12 +229,22 @@ jobs:
         env:
           INTEGRATION_REPOSITORY: ${{ env.RUNS_ON_ECR_CACHE }}
           TAG: integration-test-${{ github.run_id }}
+          CACHE_SUFFIX: ${{ steps.format-branch.outputs.cache-suffix }}
+          HEAD_SHA: ${{ github.event.pull_request.head.sha || github.sha }}
         run: |
-          cd backend && docker buildx bake --push \
+          docker buildx bake --push \
+            --set backend.cache-from=type=registry,ref=${RUNS_ON_ECR_CACHE}:backend-cache-${HEAD_SHA} \
+            --set backend.cache-from=type=registry,ref=${RUNS_ON_ECR_CACHE}:backend-cache-${CACHE_SUFFIX} \
             --set backend.cache-from=type=registry,ref=${RUNS_ON_ECR_CACHE}:backend-cache \
             --set backend.cache-from=type=registry,ref=onyxdotapp/onyx-backend:latest \
+            --set backend.cache-to=type=registry,ref=${RUNS_ON_ECR_CACHE}:backend-cache-${HEAD_SHA},mode=max \
+            --set backend.cache-to=type=registry,ref=${RUNS_ON_ECR_CACHE}:backend-cache-${CACHE_SUFFIX},mode=max \
             --set backend.cache-to=type=registry,ref=${RUNS_ON_ECR_CACHE}:backend-cache,mode=max \
+            --set integration.cache-from=type=registry,ref=${RUNS_ON_ECR_CACHE}:integration-cache-${HEAD_SHA} \
+            --set integration.cache-from=type=registry,ref=${RUNS_ON_ECR_CACHE}:integration-cache-${CACHE_SUFFIX} \
             --set integration.cache-from=type=registry,ref=${RUNS_ON_ECR_CACHE}:integration-cache \
+            --set integration.cache-to=type=registry,ref=${RUNS_ON_ECR_CACHE}:integration-cache-${HEAD_SHA},mode=max \
+            --set integration.cache-to=type=registry,ref=${RUNS_ON_ECR_CACHE}:integration-cache-${CACHE_SUFFIX},mode=max \
             --set integration.cache-to=type=registry,ref=${RUNS_ON_ECR_CACHE}:integration-cache,mode=max \
             integration
 
@@ -188,7 +271,7 @@ jobs:
     steps:
       - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2
       - name: Checkout code
-        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # ratchet:actions/checkout@v6
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6
         with:
           persist-credentials: false
 
@@ -203,21 +286,25 @@ jobs:
 
       # NOTE: Use pre-ping/null pool to reduce flakiness due to dropped connections
       # NOTE: don't need web server for integration tests
-      - name: Start Docker containers
+      - name: Create .env file for Docker Compose
         env:
           ECR_CACHE: ${{ env.RUNS_ON_ECR_CACHE }}
           RUN_ID: ${{ github.run_id }}
+        run: |
+          cat <<EOF > deployment/docker_compose/.env
+          AUTH_TYPE=basic
+          POSTGRES_POOL_PRE_PING=true
+          POSTGRES_USE_NULL_POOL=true
+          REQUIRE_EMAIL_VERIFICATION=false
+          DISABLE_TELEMETRY=true
+          ONYX_BACKEND_IMAGE=${ECR_CACHE}:integration-test-backend-test-${RUN_ID}
+          ONYX_MODEL_SERVER_IMAGE=${ECR_CACHE}:integration-test-model-server-test-${RUN_ID}
+          INTEGRATION_TESTS_MODE=true
+          EOF
+
+      - name: Start Docker containers
         run: |
           cd deployment/docker_compose
-          AUTH_TYPE=basic \
-          POSTGRES_POOL_PRE_PING=true \
-          POSTGRES_USE_NULL_POOL=true \
-          REQUIRE_EMAIL_VERIFICATION=false \
-          DISABLE_TELEMETRY=true \
-          ONYX_BACKEND_IMAGE=${ECR_CACHE}:integration-test-backend-test-${RUN_ID} \
-          ONYX_MODEL_SERVER_IMAGE=${ECR_CACHE}:integration-test-model-server-test-${RUN_ID} \
-          INTEGRATION_TESTS_MODE=true \
-          MCP_SERVER_ENABLED=true \
           docker compose -f docker-compose.yml -f docker-compose.dev.yml up \
             relational_db \
             index \
@@ -226,7 +313,6 @@ jobs:
             api_server \
             inference_model_server \
             indexing_model_server \
-            mcp_server \
             background \
             -d
         id: start_docker
@@ -269,12 +355,6 @@ jobs:
           }
 
           wait_for_service "http://localhost:8080/health" "API server"
-          test_dir="${{ matrix.test-dir.path }}"
-          if [ "$test_dir" = "tests/mcp" ]; then
-            wait_for_service "http://localhost:8090/health" "MCP server"
-          else
-            echo "Skipping MCP server wait for non-MCP suite: $test_dir"
-          fi
           echo "Finished waiting for services."
 
       - name: Start Mock Services
@@ -305,8 +385,6 @@ jobs:
               -e VESPA_HOST=index \
               -e REDIS_HOST=cache \
               -e API_SERVER_HOST=api_server \
-              -e MCP_SERVER_HOST=mcp_server \
-              -e MCP_SERVER_PORT=8090 \
               -e OPENAI_API_KEY=${OPENAI_API_KEY} \
               -e EXA_API_KEY=${EXA_API_KEY} \
               -e SLACK_BOT_TOKEN=${SLACK_BOT_TOKEN} \
@@ -350,7 +428,6 @@ jobs:
           path: ${{ github.workspace }}/docker-compose.log
       # ------------------------------------------------------------
 
-
   required:
     # NOTE: Github-hosted runners have about 20s faster queue times and are preferred here.
     runs-on: ubuntu-slim

.github/workflows/pr-playwright-tests.yml

@@ -3,7 +3,15 @@ concurrency:
   group: Run-Playwright-Tests-${{ github.workflow }}-${{ github.head_ref || github.event.workflow_run.head_branch || github.run_id }}
   cancel-in-progress: true
 
-on: push
+on:
+  merge_group:
+  pull_request:
+    branches:
+      - main
+      - "release/**"
+  push:
+    tags:
+      - "v*.*.*"
 
 permissions:
   contents: read
@@ -46,16 +54,36 @@ env:
 
 jobs:
   build-web-image:
-    runs-on: [runs-on, runner=4cpu-linux-arm64, "run-id=${{ github.run_id }}-build-web-image", "extras=ecr-cache"]
+    runs-on:
+      [
+        runs-on,
+        runner=4cpu-linux-arm64,
+        "run-id=${{ github.run_id }}-build-web-image",
+        "extras=ecr-cache",
+      ]
     timeout-minutes: 45
     steps:
       - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2
 
       - name: Checkout code
-        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # ratchet:actions/checkout@v6
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6
         with:
           persist-credentials: false
 
+      - name: Format branch name for cache
+        id: format-branch
+        env:
+          PR_NUMBER: ${{ github.event.pull_request.number }}
+          REF_NAME: ${{ github.ref_name }}
+        run: |
+          if [ -n "${PR_NUMBER}" ]; then
+            CACHE_SUFFIX="${PR_NUMBER}"
+          else
+            # shellcheck disable=SC2001
+            CACHE_SUFFIX=$(echo "${REF_NAME}" | sed 's/[^A-Za-z0-9._-]/-/g')
+          fi
+          echo "cache-suffix=${CACHE_SUFFIX}" >> $GITHUB_OUTPUT
+
       - name: Set up Docker Buildx
         uses: docker/setup-buildx-action@e468171a9de216ec08956ac3ada2f0791b6bd435 # ratchet:docker/setup-buildx-action@v3
 
@@ -76,22 +104,47 @@ jobs:
           tags: ${{ env.RUNS_ON_ECR_CACHE }}:playwright-test-web-${{ github.run_id }}
           push: true
           cache-from: |
+            type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:web-cache-${{ github.event.pull_request.head.sha || github.sha }}
+            type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:web-cache-${{ steps.format-branch.outputs.cache-suffix }}
             type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:web-cache
             type=registry,ref=onyxdotapp/onyx-web-server:latest
-          cache-to: type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:web-cache,mode=max
+          cache-to: |
+            type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:web-cache-${{ github.event.pull_request.head.sha || github.sha }},mode=max
+            type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:web-cache-${{ steps.format-branch.outputs.cache-suffix }},mode=max
+            type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:web-cache,mode=max
           no-cache: ${{ vars.DOCKER_NO_CACHE == 'true' }}
 
   build-backend-image:
-    runs-on: [runs-on, runner=1cpu-linux-arm64, "run-id=${{ github.run_id }}-build-backend-image", "extras=ecr-cache"]
+    runs-on:
+      [
+        runs-on,
+        runner=1cpu-linux-arm64,
+        "run-id=${{ github.run_id }}-build-backend-image",
+        "extras=ecr-cache",
+      ]
     timeout-minutes: 45
     steps:
       - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2
 
       - name: Checkout code
-        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # ratchet:actions/checkout@v6
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6
         with:
           persist-credentials: false
 
+      - name: Format branch name for cache
+        id: format-branch
+        env:
+          PR_NUMBER: ${{ github.event.pull_request.number }}
+          REF_NAME: ${{ github.ref_name }}
+        run: |
+          if [ -n "${PR_NUMBER}" ]; then
+            CACHE_SUFFIX="${PR_NUMBER}"
+          else
+            # shellcheck disable=SC2001
+            CACHE_SUFFIX=$(echo "${REF_NAME}" | sed 's/[^A-Za-z0-9._-]/-/g')
+          fi
+          echo "cache-suffix=${CACHE_SUFFIX}" >> $GITHUB_OUTPUT
+
       - name: Set up Docker Buildx
         uses: docker/setup-buildx-action@e468171a9de216ec08956ac3ada2f0791b6bd435 # ratchet:docker/setup-buildx-action@v3
 
@@ -112,23 +165,47 @@ jobs:
           tags: ${{ env.RUNS_ON_ECR_CACHE }}:playwright-test-backend-${{ github.run_id }}
           push: true
           cache-from: |
+            type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:backend-cache-${{ github.event.pull_request.head.sha || github.sha }}
+            type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:backend-cache-${{ steps.format-branch.outputs.cache-suffix }}
             type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:backend-cache
             type=registry,ref=onyxdotapp/onyx-backend:latest
           cache-to: |
+            type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:backend-cache-${{ github.event.pull_request.head.sha || github.sha }},mode=max
+            type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:backend-cache-${{ steps.format-branch.outputs.cache-suffix }},mode=max
             type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:backend-cache,mode=max
           no-cache: ${{ vars.DOCKER_NO_CACHE == 'true' }}
 
   build-model-server-image:
-    runs-on: [runs-on, runner=1cpu-linux-arm64, "run-id=${{ github.run_id }}-build-model-server-image", "extras=ecr-cache"]
+    runs-on:
+      [
+        runs-on,
+        runner=1cpu-linux-arm64,
+        "run-id=${{ github.run_id }}-build-model-server-image",
+        "extras=ecr-cache",
+      ]
     timeout-minutes: 45
     steps:
       - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2
 
       - name: Checkout code
-        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # ratchet:actions/checkout@v6
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6
         with:
           persist-credentials: false
 
+      - name: Format branch name for cache
+        id: format-branch
+        env:
+          PR_NUMBER: ${{ github.event.pull_request.number }}
+          REF_NAME: ${{ github.ref_name }}
+        run: |
+          if [ -n "${PR_NUMBER}" ]; then
+            CACHE_SUFFIX="${PR_NUMBER}"
+          else
+            # shellcheck disable=SC2001
+            CACHE_SUFFIX=$(echo "${REF_NAME}" | sed 's/[^A-Za-z0-9._-]/-/g')
+          fi
+          echo "cache-suffix=${CACHE_SUFFIX}" >> $GITHUB_OUTPUT
+
       - name: Set up Docker Buildx
         uses: docker/setup-buildx-action@e468171a9de216ec08956ac3ada2f0791b6bd435 # ratchet:docker/setup-buildx-action@v3
 
@@ -149,9 +226,14 @@ jobs:
           tags: ${{ env.RUNS_ON_ECR_CACHE }}:playwright-test-model-server-${{ github.run_id }}
           push: true
           cache-from: |
+            type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:model-server-cache-${{ github.event.pull_request.head.sha || github.sha }}
+            type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:model-server-cache-${{ steps.format-branch.outputs.cache-suffix }}
             type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:model-server-cache
             type=registry,ref=onyxdotapp/onyx-model-server:latest
-          cache-to: type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:model-server-cache,mode=max
+          cache-to: |
+            type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:model-server-cache-${{ github.event.pull_request.head.sha || github.sha }},mode=max
+            type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:model-server-cache-${{ steps.format-branch.outputs.cache-suffix }},mode=max
+            type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:model-server-cache,mode=max
           no-cache: ${{ vars.DOCKER_NO_CACHE == 'true' }}
 
   playwright-tests:
@@ -172,16 +254,15 @@ jobs:
       - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2
 
       - name: Checkout code
-        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # ratchet:actions/checkout@v6
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6
         with:
-          fetch-depth: 0
           persist-credentials: false
 
       - name: Setup node
-        uses: actions/setup-node@2028fbc5c25fe9cf00d9f06a71cc4710d4507903 # ratchet:actions/setup-node@v4
+        uses: actions/setup-node@395ad3262231945c25e8478fd5baf05154b1d79f # ratchet:actions/setup-node@v4
         with:
           node-version: 22
-          cache: 'npm'
+          cache: "npm"
           cache-dependency-path: ./web/package-lock.json
 
       - name: Install node dependencies
@@ -390,7 +471,6 @@ jobs:
         if: ${{ contains(needs.*.result, 'failure') || contains(needs.*.result, 'cancelled') || contains(needs.*.result, 'skipped') }}
         run: exit 1
 
-
 # NOTE: Chromatic UI diff testing is currently disabled.
 # We are using Playwright for local and CI testing without visual regression checks.
 # Chromatic may be reintroduced in the future for UI diff testing if needed.
@@ -408,12 +488,12 @@ jobs:
 #     ]
 #   steps:
 #     - name: Checkout code
-#       uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # ratchet:actions/checkout@v6
+#       uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6
 #       with:
 #         fetch-depth: 0
 
 #     - name: Setup node
-#       uses: actions/setup-node@2028fbc5c25fe9cf00d9f06a71cc4710d4507903 # ratchet:actions/setup-node@v4
+#       uses: actions/setup-node@395ad3262231945c25e8478fd5baf05154b1d79f # ratchet:actions/setup-node@v4
 #       with:
 #         node-version: 22
 

.github/workflows/pr-python-checks.yml

@@ -9,29 +9,14 @@ on:
     branches:
       - main
       - 'release/**'
+  push:
+    tags:
+      - "v*.*.*"
 
 permissions:
   contents: read
 
 jobs:
-  validate-requirements:
-    runs-on: ubuntu-slim
-    timeout-minutes: 45
-    steps:
-      - name: Checkout code
-        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # ratchet:actions/checkout@v6
-        with:
-          persist-credentials: false
-
-      - name: Setup uv
-        uses: astral-sh/setup-uv@caf0cab7a618c569241d31dcd442f54681755d39 # ratchet:astral-sh/setup-uv@v3
-        # TODO: Enable caching once there is a uv.lock file checked in.
-        # with:
-        #   enable-cache: true
-
-      - name: Validate requirements lock files
-        run: ./backend/scripts/compile_requirements.py --check
-
   mypy-check:
     # See https://runs-on.com/runners/linux/
     # Note: Mypy seems quite optimized for x64 compared to arm64.
@@ -42,7 +27,7 @@ jobs:
     steps:
       - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2
       - name: Checkout code
-        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # ratchet:actions/checkout@v6
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6
         with:
           persist-credentials: false
 
@@ -55,35 +40,10 @@ jobs:
             backend/requirements/model_server.txt
             backend/requirements/ee.txt
 
-      - name: Generate OpenAPI schema
-        shell: bash
-        working-directory: backend
-        env:
-          PYTHONPATH: "."
-        run: |
-          python scripts/onyx_openapi_schema.py --filename generated/openapi.json
-
-      # needed for pulling openapitools/openapi-generator-cli
-      # otherwise, we hit the "Unauthenticated users" limit
-      # https://docs.docker.com/docker-hub/usage/
-      - name: Login to Docker Hub
-        uses: docker/login-action@5e57cd118135c172c3672efd75eb46360885c0ef # ratchet:docker/login-action@v3
-        with:
-          username: ${{ secrets.DOCKER_USERNAME }}
-          password: ${{ secrets.DOCKER_TOKEN }}
-
-      - name: Generate OpenAPI Python client
+      - name: Generate OpenAPI schema and Python client
         shell: bash
         run: |
-          docker run --rm \
-            -v "${{ github.workspace }}/backend/generated:/local" \
-            openapitools/openapi-generator-cli generate \
-            -i /local/openapi.json \
-            -g python \
-            -o /local/onyx_openapi_client \
-            --package-name onyx_openapi_client \
-            --skip-validate-spec \
-            --openapi-normalizer "SIMPLIFY_ONEOF_ANYOF=true,SET_OAS3_NULLABLE=true"
+          ods openapi all
 
       - name: Cache mypy cache
         if: ${{ vars.DISABLE_MYPY_CACHE != 'true' }}
@@ -101,11 +61,8 @@ jobs:
           TERM: xterm-256color
         run: mypy .
 
-      - name: Check import order with reorder-python-imports
-        working-directory: ./backend
-        run: |
-          find ./onyx -name "*.py" | xargs reorder-python-imports --py311-plus
-
-      - name: Check code formatting with Black
-        working-directory: ./backend
-        run: black --check .
+      - name: Run MyPy (tools/)
+        env:
+          MYPY_FORCE_COLOR: 1
+          TERM: xterm-256color
+        run: mypy tools/

.github/workflows/pr-python-connector-tests.yml

@@ -7,6 +7,9 @@ on:
   merge_group:
   pull_request:
     branches: [main]
+  push:
+    tags:
+      - "v*.*.*"
   schedule:
     # This cron expression runs the job daily at 16:00 UTC (9am PT)
     - cron: "0 16 * * *"
@@ -130,12 +133,13 @@ jobs:
 
     env:
       PYTHONPATH: ./backend
+      DISABLE_TELEMETRY: "true"
 
     steps:
       - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2
 
       - name: Checkout code
-        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # ratchet:actions/checkout@v6
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6
         with:
           persist-credentials: false
 
@@ -157,16 +161,20 @@ jobs:
             hubspot:
               - 'backend/onyx/connectors/hubspot/**'
               - 'backend/tests/daily/connectors/hubspot/**'
+              - 'uv.lock'
             salesforce:
               - 'backend/onyx/connectors/salesforce/**'
               - 'backend/tests/daily/connectors/salesforce/**'
+              - 'uv.lock'
             github:
               - 'backend/onyx/connectors/github/**'
               - 'backend/tests/daily/connectors/github/**'
+              - 'uv.lock'
             file_processing:
               - 'backend/onyx/file_processing/**'
+              - 'uv.lock'
 
-      - name: Run Tests (excluding HubSpot, Salesforce, and GitHub)
+      - name: Run Tests (excluding HubSpot, Salesforce, GitHub, and Coda)
         shell: script -q -e -c "bash --noprofile --norc -eo pipefail {0}"
         run: |
           py.test \
@@ -179,7 +187,8 @@ jobs:
             backend/tests/daily/connectors \
             --ignore backend/tests/daily/connectors/hubspot \
             --ignore backend/tests/daily/connectors/salesforce \
-            --ignore backend/tests/daily/connectors/github
+            --ignore backend/tests/daily/connectors/github \
+            --ignore backend/tests/daily/connectors/coda
 
       - name: Run HubSpot Connector Tests
         if: ${{ github.event_name == 'schedule' || steps.changes.outputs.hubspot == 'true' || steps.changes.outputs.file_processing == 'true' }}

.github/workflows/pr-python-model-tests.yml

@@ -39,7 +39,7 @@ jobs:
 
     steps:
       - name: Checkout code
-        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # ratchet:actions/checkout@v6
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6
         with:
           persist-credentials: false
 
@@ -61,7 +61,7 @@ jobs:
           docker tag onyxdotapp/onyx-model-server:latest onyxdotapp/onyx-model-server:test
 
       - name: Set up Python
-        uses: actions/setup-python@e797f83bcb11b83ae66e0230d6156d7c80228e7c # ratchet:actions/setup-python@v6
+        uses: actions/setup-python@83679a892e2d95755f2dac6acb0bfd1e9ac5d548 # ratchet:actions/setup-python@v6
         with:
           python-version: "3.11"
           cache: "pip"

.github/workflows/pr-python-tests.yml

@@ -9,6 +9,9 @@ on:
     branches:
       - main
       - 'release/**'
+  push:
+    tags:
+      - "v*.*.*"
 
 permissions:
   contents: read
@@ -23,15 +26,13 @@ jobs:
     env:
       PYTHONPATH: ./backend
       REDIS_CLOUD_PYTEST_PASSWORD: ${{ secrets.REDIS_CLOUD_PYTEST_PASSWORD }}
-      SF_USERNAME: ${{ secrets.SF_USERNAME }}
-      SF_PASSWORD: ${{ secrets.SF_PASSWORD }}
-      SF_SECURITY_TOKEN: ${{ secrets.SF_SECURITY_TOKEN }}
+      DISABLE_TELEMETRY: "true"
 
     steps:
     - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2
 
     - name: Checkout code
-      uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # ratchet:actions/checkout@v6
+      uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6
       with:
         persist-credentials: false
 

.github/workflows/pr-quality-checks.yml

@@ -6,32 +6,42 @@ concurrency:
 on:
   merge_group:
   pull_request: null
+  push:
+    branches:
+      - main
+    tags:
+      - "v*.*.*"
 
 permissions:
   contents: read
 
 jobs:
   quality-checks:
-    # See https://runs-on.com/runners/linux/
-    runs-on: [runs-on, runner=1cpu-linux-arm64, "run-id=${{ github.run_id }}-quality-checks"]
+    runs-on: ubuntu-latest
     timeout-minutes: 45
     steps:
-      - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2
-      - uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # ratchet:actions/checkout@v6
+      - uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6
         with:
           fetch-depth: 0
           persist-credentials: false
-      - uses: actions/setup-python@e797f83bcb11b83ae66e0230d6156d7c80228e7c # ratchet:actions/setup-python@v6
+      - uses: actions/setup-python@83679a892e2d95755f2dac6acb0bfd1e9ac5d548 # ratchet:actions/setup-python@v6
         with:
           python-version: "3.11"
       - name: Setup Terraform
         uses: hashicorp/setup-terraform@b9cd54a3c349d3f38e8881555d616ced269862dd # ratchet:hashicorp/setup-terraform@v3
-      - uses: pre-commit/action@2c7b3805fd2a0fd8c1884dcaebf91fc102a13ecd # ratchet:pre-commit/action@v3.0.1
-        env:
-          # uv-run is mypy's id and mypy is covered by the Python Checks which caches dependencies better.
-          SKIP: uv-run
+      - name: Setup node
+        uses: actions/setup-node@395ad3262231945c25e8478fd5baf05154b1d79f # ratchet:actions/setup-node@v6
+        with: # zizmor: ignore[cache-poisoning]
+          node-version: 22
+          cache: "npm"
+          cache-dependency-path: ./web/package-lock.json
+      - name: Install node dependencies
+        working-directory: ./web
+        run: npm ci
+      - uses: j178/prek-action@91fd7d7cf70ae1dee9f4f44e7dfa5d1073fe6623 # ratchet:j178/prek-action@v1
         with:
-          extra_args: ${{ github.event_name == 'pull_request' && format('--from-ref {0} --to-ref {1}', github.event.pull_request.base.sha, github.event.pull_request.head.sha) || '' }}
+          prek-version: '0.2.21'
+          extra-args: ${{ github.event_name == 'pull_request' && format('--from-ref {0} --to-ref {1}', github.event.pull_request.base.sha, github.event.pull_request.head.sha) || github.event_name == 'merge_group' && format('--from-ref {0} --to-ref {1}', github.event.merge_group.base_sha, github.event.merge_group.head_sha) || github.ref_name == 'main' && '--all-files' || '' }}
       - name: Check Actions
         uses: giner/check-actions@28d366c7cbbe235f9624a88aa31a628167eee28c # ratchet:giner/check-actions@v1.0.1
         with:

.github/workflows/release-devtools.yml

@@ -0,0 +1,41 @@
+name: Release Devtools
+
+on:
+  push:
+    tags:
+      - "ods/v*.*.*"
+
+jobs:
+  pypi:
+    runs-on: ubuntu-latest
+    environment:
+      name: release-devtools
+    permissions:
+      id-token: write
+    timeout-minutes: 10
+    strategy:
+      matrix:
+        os-arch:
+          - { goos: "linux", goarch: "amd64" }
+          - { goos: "linux", goarch: "arm64" }
+          - { goos: "windows", goarch: "amd64" }
+          - { goos: "windows", goarch: "arm64" }
+          - { goos: "darwin", goarch: "amd64" }
+          - { goos: "darwin", goarch: "arm64" }
+          - { goos: "", goarch: "" }
+    steps:
+      - uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6
+        with:
+          persist-credentials: false
+          fetch-depth: 0
+      - uses: astral-sh/setup-uv@ed21f2f24f8dd64503750218de024bcf64c7250a # ratchet:astral-sh/setup-uv@v7
+        with:
+          enable-cache: false
+          version: "0.9.9"
+      - run: |
+          GOOS="${{ matrix.os-arch.goos }}" \
+          GOARCH="${{ matrix.os-arch.goarch }}" \
+          uv build --wheel
+        working-directory: tools/ods
+      - run: uv publish
+        working-directory: tools/ods

.github/workflows/sync_foss.yml

@@ -14,7 +14,7 @@ jobs:
       contents: read
     steps:
       - name: Checkout main Onyx repo
-        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # ratchet:actions/checkout@v6
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6
         with:
           fetch-depth: 0
           persist-credentials: false

.github/workflows/tag-nightly.yml

@@ -18,7 +18,7 @@ jobs:
       # see https://github.com/orgs/community/discussions/27028#discussioncomment-3254367 for the workaround we
       # implement here which needs an actual user's deploy key
       - name: Checkout code
-        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # ratchet:actions/checkout@v6
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6
         with:
           ssh-key: "${{ secrets.DEPLOY_KEY }}"
           persist-credentials: true

.github/workflows/zizmor.yml

@@ -17,15 +17,18 @@ jobs:
       security-events: write # needed for SARIF uploads
     steps:
       - name: Checkout repository
-        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # ratchet:actions/checkout@v6.0.0
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # ratchet:actions/checkout@v6.0.1
         with:
           persist-credentials: false
 
       - name: Install the latest version of uv
-        uses: astral-sh/setup-uv@5a7eac68fb9809dea845d802897dc5c723910fa3 # ratchet:astral-sh/setup-uv@v7.1.3
+        uses: astral-sh/setup-uv@ed21f2f24f8dd64503750218de024bcf64c7250a # ratchet:astral-sh/setup-uv@v7
+        with:
+          enable-cache: false
+          version: "0.9.9"
 
       - name: Run zizmor
-        run: uvx zizmor==1.16.3 --format=sarif . > results.sarif
+        run: uv run --no-sync --with zizmor zizmor --format=sarif . > results.sarif
         env:
           GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
 

.gitignore

@@ -49,5 +49,10 @@ CLAUDE.md
 # Local .terraform.lock.hcl file
 .terraform.lock.hcl
 
+node_modules
+
 # MCP configs
 .playwright-mcp
+
+# plans
+plans/

.pre-commit-config.yaml

@@ -5,73 +5,138 @@ default_install_hook_types:
   - post-rewrite
 repos:
   - repo: https://github.com/astral-sh/uv-pre-commit
-    # This revision is from https://github.com/astral-sh/uv-pre-commit/pull/53
+    # From: https://github.com/astral-sh/uv-pre-commit/pull/53/commits/d30b4298e4fb63ce8609e29acdbcf4c9018a483c
     rev: d30b4298e4fb63ce8609e29acdbcf4c9018a483c
     hooks:
       - id: uv-sync
+        args: ["--active", "--locked", "--all-extras"]
+      - id: uv-lock
+        files: ^pyproject\.toml$
+      - id: uv-export
+        name: uv-export default.txt
+        args:
+          [
+            "--no-emit-project",
+            "--no-default-groups",
+            "--no-hashes",
+            "--extra",
+            "backend",
+            "-o",
+            "backend/requirements/default.txt",
+          ]
+        files: ^(pyproject\.toml|uv\.lock|backend/requirements/.*\.txt)$
+      - id: uv-export
+        name: uv-export dev.txt
+        args:
+          [
+            "--no-emit-project",
+            "--no-default-groups",
+            "--no-hashes",
+            "--extra",
+            "dev",
+            "-o",
+            "backend/requirements/dev.txt",
+          ]
+        files: ^(pyproject\.toml|uv\.lock|backend/requirements/.*\.txt)$
+      - id: uv-export
+        name: uv-export ee.txt
+        args:
+          [
+            "--no-emit-project",
+            "--no-default-groups",
+            "--no-hashes",
+            "--extra",
+            "ee",
+            "-o",
+            "backend/requirements/ee.txt",
+          ]
+        files: ^(pyproject\.toml|uv\.lock|backend/requirements/.*\.txt)$
+      - id: uv-export
+        name: uv-export model_server.txt
+        args:
+          [
+            "--no-emit-project",
+            "--no-default-groups",
+            "--no-hashes",
+            "--extra",
+            "model_server",
+            "-o",
+            "backend/requirements/model_server.txt",
+          ]
+        files: ^(pyproject\.toml|uv\.lock|backend/requirements/.*\.txt)$
       - id: uv-run
-        name: mypy
-        args: ["mypy"]
-        pass_filenames: true
-        files: ^backend/.*\.py$
-
-  - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.6.0
-    hooks:
-      - id: check-yaml
-        files: ^.github/
+        name: Check lazy imports
+        args: ["--active", "--with=onyx-devtools", "ods", "check-lazy-imports"]
+        files: ^backend/(?!\.venv/).*\.py$
+      # NOTE: This takes ~6s on a single, large module which is prohibitively slow.
+      # - id: uv-run
+      #   name: mypy
+      #   args: ["--all-extras", "mypy"]
+      #   pass_filenames: true
+      #   files: ^backend/.*\.py$
 
   - repo: https://github.com/rhysd/actionlint
-    rev: v1.7.8
+    rev: a443f344ff32813837fa49f7aa6cbc478d770e62 # frozen: v1.7.9
     hooks:
       - id: actionlint
 
   - repo: https://github.com/psf/black
-    rev: 25.1.0
+    rev: 8a737e727ac5ab2f1d4cf5876720ed276dc8dc4b # frozen: 25.1.0
     hooks:
-    - id: black
-      language_version: python3.11
+      - id: black
+        language_version: python3.11
 
   # this is a fork which keeps compatibility with black
   - repo: https://github.com/wimglenn/reorder-python-imports-black
-    rev: v3.14.0
+    rev: f55cd27f90f0cf0ee775002c2383ce1c7820013d # frozen: v3.14.0
     hooks:
-    - id: reorder-python-imports
-      args: ['--py311-plus', '--application-directories=backend/']
-      # need to ignore alembic files, since reorder-python-imports gets confused
-      # and thinks that alembic is a local package since there is a folder
-      # in the backend directory called `alembic`
-      exclude: ^backend/alembic/
+      - id: reorder-python-imports
+        args: ["--py311-plus", "--application-directories=backend/"]
+        # need to ignore alembic files, since reorder-python-imports gets confused
+        # and thinks that alembic is a local package since there is a folder
+        # in the backend directory called `alembic`
+        exclude: ^backend/alembic/
 
   # These settings will remove unused imports with side effects
   # Note: The repo currently does not and should not have imports with side effects
   - repo: https://github.com/PyCQA/autoflake
-    rev: v2.3.1
+    rev: 0544741e2b4a22b472d9d93e37d4ea9153820bb1 # frozen: v2.3.1
     hooks:
       - id: autoflake
-        args: [ '--remove-all-unused-imports', '--remove-unused-variables', '--in-place' , '--recursive']
+        args:
+          [
+            "--remove-all-unused-imports",
+            "--remove-unused-variables",
+            "--in-place",
+            "--recursive",
+          ]
+
+  - repo: https://github.com/golangci/golangci-lint
+    rev: 9f61b0f53f80672872fced07b6874397c3ed197b # frozen: v2.7.2
+    hooks:
+      - id: golangci-lint
+        entry: bash -c "find tools/ -name go.mod -print0 | xargs -0 -I{} bash -c 'cd \"$(dirname {})\" && golangci-lint run ./...'"
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.11.4
+    rev: 971923581912ef60a6b70dbf0c3e9a39563c9d47 # frozen: v0.11.4
     hooks:
       - id: ruff
 
   - repo: https://github.com/pre-commit/mirrors-prettier
-    rev: v3.1.0
+    rev: ffb6a759a979008c0e6dff86e39f4745a2d9eac4 # frozen: v3.1.0
     hooks:
-    - id: prettier
-      types_or: [html, css, javascript, ts, tsx]
-      language_version: system
+      - id: prettier
+        types_or: [html, css, javascript, ts, tsx]
+        language_version: system
 
   - repo: https://github.com/sirwart/ripsecrets
-    rev: v0.1.11
+    rev: 7d94620933e79b8acaa0cd9e60e9864b07673d86 # frozen: v0.1.11
     hooks:
       - id: ripsecrets
         args:
-        - --additional-pattern
-        - ^sk-[A-Za-z0-9_\-]{20,}$
-
+          - --additional-pattern
+          - ^sk-[A-Za-z0-9_\-]{20,}$
 
   - repo: local
     hooks:
@@ -82,8 +147,13 @@ repos:
         pass_filenames: false
         files: \.tf$
 
-      - id: check-lazy-imports
-        name: Check lazy imports
-        entry: python3 backend/scripts/check_lazy_imports.py
+      # Uses tsgo (TypeScript's native Go compiler) for ~10x faster type checking.
+      # This is a preview package - if it breaks:
+      #   1. Try updating: cd web && npm update @typescript/native-preview
+      #   2. Or fallback to tsc: replace 'tsgo' with 'tsc' below
+      - id: typescript-check
+        name: TypeScript type check
+        entry: bash -c 'cd web && npx tsgo --noEmit --project tsconfig.types.json'
         language: system
-        files: ^backend/(?!\.venv/).*\.py$
+        pass_filenames: false
+        files: ^web/.*\.(ts|tsx)$

.vscode/env_template.txt

@@ -5,11 +5,8 @@
 # For local dev, often user Authentication is not needed
 AUTH_TYPE=disabled
 
-# Skip warm up for dev
-SKIP_WARM_UP=True
-
 # Always keep these on for Dev
-# Logs all model prompts to stdout
+# Logs model prompts, reasoning, and answer to stdout
 LOG_ONYX_MODEL_INTERACTIONS=True
 # More verbose logging
 LOG_LEVEL=debug
@@ -37,31 +34,16 @@ OPENAI_API_KEY=<REPLACE THIS>
 GEN_AI_MODEL_VERSION=gpt-4o
 FAST_GEN_AI_MODEL_VERSION=gpt-4o
 
-# For Onyx Slack Bot, overrides the UI values so no need to set this up via UI every time
-# Only needed if using OnyxBot
-#ONYX_BOT_SLACK_APP_TOKEN=<REPLACE THIS>
-#ONYX_BOT_SLACK_BOT_TOKEN=<REPLACE THIS>
-
 
 # Python stuff
 PYTHONPATH=../backend
 PYTHONUNBUFFERED=1
 
 
-# Internet Search
-EXA_API_KEY=<REPLACE THIS>
-
-
 # Enable the full set of Danswer Enterprise Edition features
 # NOTE: DO NOT ENABLE THIS UNLESS YOU HAVE A PAID ENTERPRISE LICENSE (or if you are using this for local testing/development)
 ENABLE_PAID_ENTERPRISE_EDITION_FEATURES=False
 
-# Agent Search configs  # TODO: Remove give proper namings
-AGENT_RETRIEVAL_STATS=False   # Note: This setting will incur substantial re-ranking effort
-AGENT_RERANKING_STATS=True
-AGENT_MAX_QUERY_RETRIEVAL_RESULTS=20
-AGENT_RERANKING_MAX_QUERY_RETRIEVAL_RESULTS=20
-
 # S3 File Store Configuration (MinIO for local development)
 S3_ENDPOINT_URL=http://localhost:9004
 S3_FILE_STORE_BUCKET_NAME=onyx-file-store-bucket

.vscode/launch.template.jsonc

@@ -133,8 +133,6 @@
       },
       "consoleTitle": "API Server Console"
     },
-    // For the listener to access the Slack API,
-    // ONYX_BOT_SLACK_APP_TOKEN & ONYX_BOT_SLACK_BOT_TOKEN need to be set in .env file located in the root of the project
     {
       "name": "Slack Bot",
       "consoleName": "Slack Bot",
@@ -510,7 +508,6 @@
       ],
       "cwd": "${workspaceFolder}",
       "console": "integratedTerminal",
-      "stopOnEntry": true,
       "presentation": {
         "group": "3"
       }
@@ -556,10 +553,10 @@
       "name": "Install Python Requirements",
       "type": "node",
       "request": "launch",
-      "runtimeExecutable": "bash",
+      "runtimeExecutable": "uv",
       "runtimeArgs": [
-        "-c",
-        "pip install -r backend/requirements/default.txt && pip install -r backend/requirements/dev.txt && pip install -r backend/requirements/ee.txt && pip install -r backend/requirements/model_server.txt"
+        "sync",
+        "--all-extras"
       ],
       "cwd": "${workspaceFolder}",
       "console": "integratedTerminal",
@@ -572,14 +569,14 @@
       "name": "Onyx OpenAPI Schema Generator",
       "type": "debugpy",
       "request": "launch",
-      "program": "scripts/onyx_openapi_schema.py",
-      "cwd": "${workspaceFolder}/backend",
+      "program": "backend/scripts/onyx_openapi_schema.py",
+      "cwd": "${workspaceFolder}",
       "envFile": "${workspaceFolder}/.env",
       "env": {
         "PYTHONUNBUFFERED": "1",
-        "PYTHONPATH": "."
+        "PYTHONPATH": "backend"
       },
-      "args": ["--filename", "generated/openapi.json"]
+      "args": ["--filename", "backend/generated/openapi.json", "--generate-python-client"]
     },
     {
       // script to debug multi tenant db issues

AGENTS.md.template

@@ -1,10 +1,10 @@
 # AGENTS.md
 
-This file provides guidance to Codex when working with code in this repository.
+This file provides guidance to AI agents when working with code in this repository.
 
 ## KEY NOTES
 
-- If you run into any missing python dependency errors, try running your command with `source backend/.venv/bin/activate` \
+- If you run into any missing python dependency errors, try running your command with `source .venv/bin/activate` \
 to assume the python venv.
 - To make tests work, check the `.env` file at the root of the project to find an OpenAI key.
 - If using `playwright` to explore the frontend, you can usually log in with username `a@test.com` and password
@@ -181,6 +181,286 @@ web/
 └── src/lib/                     # Utilities & business logic
 ```
 
+## Frontend Standards
+
+### 1. Import Standards
+
+**Always use absolute imports with the `@` prefix.**
+
+**Reason:** Moving files around becomes easier since you don't also have to update those import statements. This makes modifications to the codebase much nicer.
+
+```typescript
+// ✅ Good
+import { Button } from "@/components/ui/button";
+import { useAuth } from "@/hooks/useAuth";
+import { Text } from "@/refresh-components/texts/Text";
+
+// ❌ Bad
+import { Button } from "../../../components/ui/button";
+import { useAuth } from "./hooks/useAuth";
+```
+
+### 2. React Component Functions
+
+**Prefer regular functions over arrow functions for React components.**
+
+**Reason:** Functions just become easier to read.
+
+```typescript
+// ✅ Good
+function UserProfile({ userId }: UserProfileProps) {
+  return <div>User Profile</div>
+}
+
+// ❌ Bad
+const UserProfile = ({ userId }: UserProfileProps) => {
+  return <div>User Profile</div>
+}
+```
+
+### 3. Props Interface Extraction
+
+**Extract prop types into their own interface definitions.**
+
+**Reason:** Functions just become easier to read.
+
+```typescript
+// ✅ Good
+interface UserCardProps {
+  user: User
+  showActions?: boolean
+  onEdit?: (userId: string) => void
+}
+
+function UserCard({ user, showActions = false, onEdit }: UserCardProps) {
+  return <div>User Card</div>
+}
+
+// ❌ Bad
+function UserCard({
+  user,
+  showActions = false,
+  onEdit
+}: {
+  user: User
+  showActions?: boolean
+  onEdit?: (userId: string) => void
+}) {
+  return <div>User Card</div>
+}
+```
+
+### 4. Spacing Guidelines
+
+**Prefer padding over margins for spacing.**
+
+**Reason:** We want to consolidate usage to paddings instead of margins.
+
+```typescript
+// ✅ Good
+<div className="p-4 space-y-2">
+  <div className="p-2">Content</div>
+</div>
+
+// ❌ Bad
+<div className="m-4 space-y-2">
+  <div className="m-2">Content</div>
+</div>
+```
+
+### 5. Tailwind Dark Mode
+
+**Strictly forbid using the `dark:` modifier in Tailwind classes, except for logo icon handling.**
+
+**Reason:** The `colors.css` file already, VERY CAREFULLY, defines what the exact opposite colour of each light-mode colour is. Overriding this behaviour is VERY bad and will lead to horrible UI breakages.
+
+**Exception:** The `createLogoIcon` helper in `web/src/components/icons/icons.tsx` uses `dark:` modifiers (`dark:invert`, `dark:hidden`, `dark:block`) to handle third-party logo icons that cannot automatically adapt through `colors.css`. This is the ONLY acceptable use of dark mode modifiers.
+
+```typescript
+// ✅ Good - Standard components use `web/tailwind-themes/tailwind.config.js` / `web/src/app/css/colors.css`
+<div className="bg-background-neutral-03 text-text-02">
+  Content
+</div>
+
+// ✅ Good - Logo icons with dark mode handling via createLogoIcon
+export const GithubIcon = createLogoIcon(githubLightIcon, {
+  monochromatic: true,  // Will apply dark:invert internally
+});
+
+export const GitbookIcon = createLogoIcon(gitbookLightIcon, {
+  darkSrc: gitbookDarkIcon,  // Will use dark:hidden/dark:block internally
+});
+
+// ❌ Bad - Manual dark mode overrides
+<div className="bg-white dark:bg-black text-black dark:text-white">
+  Content
+</div>
+```
+
+### 6. Class Name Utilities
+
+**Use the `cn` utility instead of raw string formatting for classNames.**
+
+**Reason:** `cn`s are easier to read. They also allow for more complex types (i.e., string-arrays) to get formatted properly (it flattens each element in that string array down). As a result, it can allow things such as conditionals (i.e., `myCondition && "some-tailwind-class"`, which evaluates to `false` when `myCondition` is `false`) to get filtered out.
+
+```typescript
+import { cn } from '@/lib/utils'
+
+// ✅ Good
+<div className={cn(
+  'base-class',
+  isActive && 'active-class',
+  className
+)}>
+  Content
+</div>
+
+// ❌ Bad
+<div className={`base-class ${isActive ? 'active-class' : ''} ${className}`}>
+  Content
+</div>
+```
+
+### 7. Custom Hooks Organization
+
+**Follow a "hook-per-file" layout. Each hook should live in its own file within `web/src/hooks`.**
+
+**Reason:** This is just a layout preference. Keeps code clean.
+
+```typescript
+// web/src/hooks/useUserData.ts
+export function useUserData(userId: string) {
+  // hook implementation
+}
+
+// web/src/hooks/useLocalStorage.ts
+export function useLocalStorage<T>(key: string, initialValue: T) {
+  // hook implementation
+}
+```
+
+### 8. Icon Usage
+
+**ONLY use icons from the `web/src/icons` directory. Do NOT use icons from `react-icons`, `lucide`, or other external libraries.**
+
+**Reason:** We have a very carefully curated selection of icons that match our Onyx guidelines. We do NOT want to muddy those up with different aesthetic stylings.
+
+```typescript
+// ✅ Good
+import SvgX from "@/icons/x";
+import SvgMoreHorizontal from "@/icons/more-horizontal";
+
+// ❌ Bad
+import { User } from "lucide-react";
+import { FiSearch } from "react-icons/fi";
+```
+
+**Missing Icons**: If an icon is needed but doesn't exist in the `web/src/icons` directory, import it from Figma using the Figma MCP tool and add it to the icons directory.
+If you need help with this step, reach out to `raunak@onyx.app`.
+
+### 9. Text Rendering
+
+**Prefer using the `refresh-components/texts/Text` component for all text rendering. Avoid "naked" text nodes.**
+
+**Reason:** The `Text` component is fully compliant with the stylings provided in Figma. It provides easy utilities to specify the text-colour and font-size in the form of flags. Super duper easy.
+
+```typescript
+// ✅ Good
+import { Text } from '@/refresh-components/texts/Text'
+
+function UserCard({ name }: { name: string }) {
+  return (
+    <Text
+      {/* The `text03` flag makes the text it renders to be coloured the 3rd-scale grey */}
+      text03
+      {/* The `mainAction` flag makes the text it renders to be "main-action" font + line-height + weightage, as described in the Figma */}
+      mainAction
+    >
+      {name}
+    </Text>
+  )
+}
+
+// ❌ Bad
+function UserCard({ name }: { name: string }) {
+  return (
+    <div>
+      <h2>{name}</h2>
+      <p>User details</p>
+    </div>
+  )
+}
+```
+
+### 10. Component Usage
+
+**Heavily avoid raw HTML input components. Always use components from the `web/src/refresh-components` or `web/lib/opal/src` directory.**
+
+**Reason:** We've put in a lot of effort to unify the components that are rendered in the Onyx app. Using raw components breaks the entire UI of the application, and leaves it in a muddier state than before.
+
+```typescript
+// ✅ Good
+import Button from '@/refresh-components/buttons/Button'
+import InputTypeIn from '@/refresh-components/inputs/InputTypeIn'
+import SvgPlusCircle from '@/icons/plus-circle'
+
+function ContactForm() {
+  return (
+    <form>
+      <InputTypeIn placeholder="Search..." />
+      <Button type="submit" leftIcon={SvgPlusCircle}>Submit</Button>
+    </form>
+  )
+}
+
+// ❌ Bad
+function ContactForm() {
+  return (
+    <form>
+      <input placeholder="Name" />
+      <textarea placeholder="Message" />
+      <button type="submit">Submit</button>
+    </form>
+  )
+}
+```
+
+### 11. Colors
+
+**Always use custom overrides for colors and borders rather than built in Tailwind CSS colors. These overrides live in `web/tailwind-themes/tailwind.config.js`.**
+
+**Reason:** Our custom color system uses CSS variables that automatically handle dark mode and maintain design consistency across the app. Standard Tailwind colors bypass this system.
+
+**Available color categories:**
+- **Text:** `text-01` through `text-05`, `text-inverted-XX`
+- **Backgrounds:** `background-neutral-XX`, `background-tint-XX` (and inverted variants)
+- **Borders:** `border-01` through `border-05`, `border-inverted-XX`
+- **Actions:** `action-link-XX`, `action-danger-XX`
+- **Status:** `status-info-XX`, `status-success-XX`, `status-warning-XX`, `status-error-XX`
+- **Theme:** `theme-primary-XX`, `theme-red-XX`, `theme-blue-XX`, etc.
+
+```typescript
+// ✅ Good - Use custom Onyx color classes
+<div className="bg-background-neutral-01 border border-border-02" />
+<div className="bg-background-tint-02 border border-border-01" />
+<div className="bg-status-success-01" />
+<div className="bg-action-link-01" />
+<div className="bg-theme-primary-05" />
+
+// ❌ Bad - Do NOT use standard Tailwind colors
+<div className="bg-gray-100 border border-gray-300 text-gray-600" />
+<div className="bg-white border border-slate-200" />
+<div className="bg-green-100 text-green-700" />
+<div className="bg-blue-100 text-blue-600" />
+<div className="bg-indigo-500" />
+```
+
+### 12. Data Fetching
+
+**Prefer using `useSWR` for data fetching. Data should generally be fetched on the client side. Components that need data should display a loader / placeholder while waiting for that data. Prefer loading data within the component that needs it rather than at the top level and passing it down.**
+
+**Reason:** Client side fetching allows us to load the skeleton of the page without waiting for data to load, leading to a snappier UX. Loading data where needed reduces dependencies between a component and its parent component(s).
+
 ## Database & Migrations
 
 ### Running Migrations
@@ -295,14 +575,6 @@ will be tailing their logs to this file.
 - Token management and rate limiting
 - Custom prompts and agent actions
 
-## UI/UX Patterns
-
-- Tailwind CSS with design system in `web/src/components/ui/`
-- Radix UI and Headless UI for accessible components
-- SWR for data fetching and caching
-- Form validation with react-hook-form
-- Error handling with popup notifications
-
 ## Creating a Plan
 When creating a plan in the `plans` directory, make sure to include at least these elements:
 

CLAUDE.md.template

@@ -4,7 +4,7 @@ This file provides guidance to Claude Code (claude.ai/code) when working with co
 
 ## KEY NOTES
 
-- If you run into any missing python dependency errors, try running your command with `source backend/.venv/bin/activate` \
+- If you run into any missing python dependency errors, try running your command with `source .venv/bin/activate` \
 to assume the python venv.
 - To make tests work, check the `.env` file at the root of the project to find an OpenAI key.
 - If using `playwright` to explore the frontend, you can usually log in with username `a@test.com` and password
@@ -184,6 +184,286 @@ web/
 └── src/lib/                     # Utilities & business logic
 ```
 
+## Frontend Standards
+
+### 1. Import Standards
+
+**Always use absolute imports with the `@` prefix.**
+
+**Reason:** Moving files around becomes easier since you don't also have to update those import statements. This makes modifications to the codebase much nicer.
+
+```typescript
+// ✅ Good
+import { Button } from "@/components/ui/button";
+import { useAuth } from "@/hooks/useAuth";
+import { Text } from "@/refresh-components/texts/Text";
+
+// ❌ Bad
+import { Button } from "../../../components/ui/button";
+import { useAuth } from "./hooks/useAuth";
+```
+
+### 2. React Component Functions
+
+**Prefer regular functions over arrow functions for React components.**
+
+**Reason:** Functions just become easier to read.
+
+```typescript
+// ✅ Good
+function UserProfile({ userId }: UserProfileProps) {
+  return <div>User Profile</div>
+}
+
+// ❌ Bad
+const UserProfile = ({ userId }: UserProfileProps) => {
+  return <div>User Profile</div>
+}
+```
+
+### 3. Props Interface Extraction
+
+**Extract prop types into their own interface definitions.**
+
+**Reason:** Functions just become easier to read.
+
+```typescript
+// ✅ Good
+interface UserCardProps {
+  user: User
+  showActions?: boolean
+  onEdit?: (userId: string) => void
+}
+
+function UserCard({ user, showActions = false, onEdit }: UserCardProps) {
+  return <div>User Card</div>
+}
+
+// ❌ Bad
+function UserCard({
+  user,
+  showActions = false,
+  onEdit
+}: {
+  user: User
+  showActions?: boolean
+  onEdit?: (userId: string) => void
+}) {
+  return <div>User Card</div>
+}
+```
+
+### 4. Spacing Guidelines
+
+**Prefer padding over margins for spacing.**
+
+**Reason:** We want to consolidate usage to paddings instead of margins.
+
+```typescript
+// ✅ Good
+<div className="p-4 space-y-2">
+  <div className="p-2">Content</div>
+</div>
+
+// ❌ Bad
+<div className="m-4 space-y-2">
+  <div className="m-2">Content</div>
+</div>
+```
+
+### 5. Tailwind Dark Mode
+
+**Strictly forbid using the `dark:` modifier in Tailwind classes, except for logo icon handling.**
+
+**Reason:** The `colors.css` file already, VERY CAREFULLY, defines what the exact opposite colour of each light-mode colour is. Overriding this behaviour is VERY bad and will lead to horrible UI breakages.
+
+**Exception:** The `createLogoIcon` helper in `web/src/components/icons/icons.tsx` uses `dark:` modifiers (`dark:invert`, `dark:hidden`, `dark:block`) to handle third-party logo icons that cannot automatically adapt through `colors.css`. This is the ONLY acceptable use of dark mode modifiers.
+
+```typescript
+// ✅ Good - Standard components use `tailwind-themes/tailwind.config.js` / `src/app/css/colors.css`
+<div className="bg-background-neutral-03 text-text-02">
+  Content
+</div>
+
+// ✅ Good - Logo icons with dark mode handling via createLogoIcon
+export const GithubIcon = createLogoIcon(githubLightIcon, {
+  monochromatic: true,  // Will apply dark:invert internally
+});
+
+export const GitbookIcon = createLogoIcon(gitbookLightIcon, {
+  darkSrc: gitbookDarkIcon,  // Will use dark:hidden/dark:block internally
+});
+
+// ❌ Bad - Manual dark mode overrides
+<div className="bg-white dark:bg-black text-black dark:text-white">
+  Content
+</div>
+```
+
+### 6. Class Name Utilities
+
+**Use the `cn` utility instead of raw string formatting for classNames.**
+
+**Reason:** `cn`s are easier to read. They also allow for more complex types (i.e., string-arrays) to get formatted properly (it flattens each element in that string array down). As a result, it can allow things such as conditionals (i.e., `myCondition && "some-tailwind-class"`, which evaluates to `false` when `myCondition` is `false`) to get filtered out.
+
+```typescript
+import { cn } from '@/lib/utils'
+
+// ✅ Good
+<div className={cn(
+  'base-class',
+  isActive && 'active-class',
+  className
+)}>
+  Content
+</div>
+
+// ❌ Bad
+<div className={`base-class ${isActive ? 'active-class' : ''} ${className}`}>
+  Content
+</div>
+```
+
+### 7. Custom Hooks Organization
+
+**Follow a "hook-per-file" layout. Each hook should live in its own file within `web/src/hooks`.**
+
+**Reason:** This is just a layout preference. Keeps code clean.
+
+```typescript
+// web/src/hooks/useUserData.ts
+export function useUserData(userId: string) {
+  // hook implementation
+}
+
+// web/src/hooks/useLocalStorage.ts
+export function useLocalStorage<T>(key: string, initialValue: T) {
+  // hook implementation
+}
+```
+
+### 8. Icon Usage
+
+**ONLY use icons from the `web/src/icons` directory. Do NOT use icons from `react-icons`, `lucide`, or other external libraries.**
+
+**Reason:** We have a very carefully curated selection of icons that match our Onyx guidelines. We do NOT want to muddy those up with different aesthetic stylings.
+
+```typescript
+// ✅ Good
+import SvgX from "@/icons/x";
+import SvgMoreHorizontal from "@/icons/more-horizontal";
+
+// ❌ Bad
+import { User } from "lucide-react";
+import { FiSearch } from "react-icons/fi";
+```
+
+**Missing Icons**: If an icon is needed but doesn't exist in the `web/src/icons` directory, import it from Figma using the Figma MCP tool and add it to the icons directory.
+If you need help with this step, reach out to `raunak@onyx.app`.
+
+### 9. Text Rendering
+
+**Prefer using the `refresh-components/texts/Text` component for all text rendering. Avoid "naked" text nodes.**
+
+**Reason:** The `Text` component is fully compliant with the stylings provided in Figma. It provides easy utilities to specify the text-colour and font-size in the form of flags. Super duper easy.
+
+```typescript
+// ✅ Good
+import { Text } from '@/refresh-components/texts/Text'
+
+function UserCard({ name }: { name: string }) {
+  return (
+    <Text
+      {/* The `text03` flag makes the text it renders to be coloured the 3rd-scale grey */}
+      text03
+      {/* The `mainAction` flag makes the text it renders to be "main-action" font + line-height + weightage, as described in the Figma */}
+      mainAction
+    >
+      {name}
+    </Text>
+  )
+}
+
+// ❌ Bad
+function UserCard({ name }: { name: string }) {
+  return (
+    <div>
+      <h2>{name}</h2>
+      <p>User details</p>
+    </div>
+  )
+}
+```
+
+### 10. Component Usage
+
+**Heavily avoid raw HTML input components. Always use components from the `web/src/refresh-components` or `web/lib/opal/src` directory.**
+
+**Reason:** We've put in a lot of effort to unify the components that are rendered in the Onyx app. Using raw components breaks the entire UI of the application, and leaves it in a muddier state than before.
+
+```typescript
+// ✅ Good
+import Button from '@/refresh-components/buttons/Button'
+import InputTypeIn from '@/refresh-components/inputs/InputTypeIn'
+import SvgPlusCircle from '@/icons/plus-circle'
+
+function ContactForm() {
+  return (
+    <form>
+      <InputTypeIn placeholder="Search..." />
+      <Button type="submit" leftIcon={SvgPlusCircle}>Submit</Button>
+    </form>
+  )
+}
+
+// ❌ Bad
+function ContactForm() {
+  return (
+    <form>
+      <input placeholder="Name" />
+      <textarea placeholder="Message" />
+      <button type="submit">Submit</button>
+    </form>
+  )
+}
+```
+
+### 11. Colors
+
+**Always use custom overrides for colors and borders rather than built in Tailwind CSS colors. These overrides live in `web/tailwind-themes/tailwind.config.js`.**
+
+**Reason:** Our custom color system uses CSS variables that automatically handle dark mode and maintain design consistency across the app. Standard Tailwind colors bypass this system.
+
+**Available color categories:**
+- **Text:** `text-01` through `text-05`, `text-inverted-XX`
+- **Backgrounds:** `background-neutral-XX`, `background-tint-XX` (and inverted variants)
+- **Borders:** `border-01` through `border-05`, `border-inverted-XX`
+- **Actions:** `action-link-XX`, `action-danger-XX`
+- **Status:** `status-info-XX`, `status-success-XX`, `status-warning-XX`, `status-error-XX`
+- **Theme:** `theme-primary-XX`, `theme-red-XX`, `theme-blue-XX`, etc.
+
+```typescript
+// ✅ Good - Use custom Onyx color classes
+<div className="bg-background-neutral-01 border border-border-02" />
+<div className="bg-background-tint-02 border border-border-01" />
+<div className="bg-status-success-01" />
+<div className="bg-action-link-01" />
+<div className="bg-theme-primary-05" />
+
+// ❌ Bad - Do NOT use standard Tailwind colors
+<div className="bg-gray-100 border border-gray-300 text-gray-600" />
+<div className="bg-white border border-slate-200" />
+<div className="bg-green-100 text-green-700" />
+<div className="bg-blue-100 text-blue-600" />
+<div className="bg-indigo-500" />
+```
+
+### 12. Data Fetching
+
+**Prefer using `useSWR` for data fetching. Data should generally be fetched on the client side. Components that need data should display a loader / placeholder while waiting for that data. Prefer loading data within the component that needs it rather than at the top level and passing it down.**
+
+**Reason:** Client side fetching allows us to load the skeleton of the page without waiting for data to load, leading to a snappier UX. Loading data where needed reduces dependencies between a component and its parent component(s).
+
 ## Database & Migrations
 
 ### Running Migrations
@@ -300,14 +580,6 @@ will be tailing their logs to this file.
 - Token management and rate limiting
 - Custom prompts and agent actions
 
-## UI/UX Patterns
-
-- Tailwind CSS with design system in `web/src/components/ui/`
-- Radix UI and Headless UI for accessible components
-- SWR for data fetching and caching
-- Form validation with react-hook-form
-- Error handling with popup notifications
-
 ## Creating a Plan
 When creating a plan in the `plans` directory, make sure to include at least these elements:
 

CONTRIBUTING.md

@@ -71,12 +71,12 @@ If using a higher version, sometimes some libraries will not be available (i.e.
 
 #### Backend: Python requirements
 
-Currently, we use pip and recommend creating a virtual environment.
+Currently, we use [uv](https://docs.astral.sh/uv/) and recommend creating a [virtual environment](https://docs.astral.sh/uv/pip/environments/#using-a-virtual-environment).
 
 For convenience here's a command for it:
 
 ```bash
-python -m venv .venv
+uv venv .venv --python 3.11
 source .venv/bin/activate
 ```
 
@@ -95,33 +95,15 @@ If using PowerShell, the command slightly differs:
 Install the required python dependencies:
 
 ```bash
-pip install -r backend/requirements/combined.txt
+uv sync --all-extras
 ```
 
-or
+Install Playwright for Python (headless browser required by the Web Connector):
 
 ```bash
-pip install -r backend/requirements/default.txt
-pip install -r backend/requirements/dev.txt
-pip install -r backend/requirements/ee.txt
-pip install -r backend/requirements/model_server.txt
+uv run playwright install
 ```
 
-Fix vscode/cursor auto-imports:
-```bash
-pip install -e .
-```
-
-Install Playwright for Python (headless browser required by the Web Connector)
-
-In the activated Python virtualenv, install Playwright for Python by running:
-
-```bash
-playwright install
-```
-
-You may have to deactivate and reactivate your virtualenv for `playwright` to appear on your path.
-
 #### Frontend: Node dependencies
 
 Onyx uses Node v22.20.0. We highly recommend you use [Node Version Manager (nvm)](https://github.com/nvm-sh/nvm)
@@ -130,7 +112,7 @@ to manage your Node installations. Once installed, you can run
 ```bash
 nvm install 22 && nvm use 22
 node -v # verify your active version
-``` 
+```
 
 Navigate to `onyx/web` and run:
 
@@ -144,21 +126,15 @@ npm i
 
 For the backend, you'll need to setup pre-commit hooks (black / reorder-python-imports).
 
-With the virtual environment active, install the pre-commit library with:
+Then run:
 
 ```bash
-pip install pre-commit
-```
-
-Then, from the `onyx/backend` directory, run:
-
-```bash
-pre-commit install
+uv run pre-commit install
 ```
 
 Additionally, we use `mypy` for static type checking.
 Onyx is fully type-annotated, and we want to keep it that way!
-To run the mypy checks manually, run `python -m mypy .` from the `onyx/backend` directory.
+To run the mypy checks manually, run `uv run mypy .` from the `onyx/backend` directory.
 
 ### Web
 
@@ -185,7 +161,7 @@ You will need Docker installed to run these containers.
 First navigate to `onyx/deployment/docker_compose`, then start up Postgres/Vespa/Redis/MinIO with:
 
 ```bash
-docker compose up -d index relational_db cache minio
+docker compose -f docker-compose.yml -f docker-compose.dev.yml up -d index relational_db cache minio
 ```
 
 (index refers to Vespa, relational_db refers to Postgres, and cache refers to Redis)

backend/.dockerignore

@@ -15,3 +15,4 @@ build/
 dist/
 .coverage
 htmlcov/
+model_server/legacy/

backend/Dockerfile.model_server

@@ -13,23 +13,10 @@ RUN uv pip install --system --no-cache-dir --upgrade \
         -r /tmp/requirements.txt && \
     rm -rf ~/.cache/uv /tmp/*.txt
 
-# Stage for downloading tokenizers
-FROM base AS tokenizers
-RUN python -c "from transformers import AutoTokenizer; \
-AutoTokenizer.from_pretrained('distilbert-base-uncased'); \
-AutoTokenizer.from_pretrained('mixedbread-ai/mxbai-rerank-xsmall-v1');"
-
-# Stage for downloading Onyx models
-FROM base AS onyx-models
-RUN python -c "from huggingface_hub import snapshot_download; \
-snapshot_download(repo_id='onyx-dot-app/hybrid-intent-token-classifier'); \
-snapshot_download(repo_id='onyx-dot-app/information-content-model');"
-
-# Stage for downloading embedding and reranking models
+# Stage for downloading embedding models
 FROM base AS embedding-models
 RUN python -c "from huggingface_hub import snapshot_download; \
-snapshot_download('nomic-ai/nomic-embed-text-v1'); \
-snapshot_download('mixedbread-ai/mxbai-rerank-xsmall-v1');"
+snapshot_download('nomic-ai/nomic-embed-text-v1');"
 
 # Initialize SentenceTransformer to cache the custom architecture
 RUN python -c "from sentence_transformers import SentenceTransformer; \
@@ -54,8 +41,6 @@ RUN groupadd -g 1001 onyx && \
 # In case the user has volumes mounted to /app/.cache/huggingface that they've downloaded while
 # running Onyx, move the current contents of the cache folder to a temporary location to ensure
 # it's preserved in order to combine with the user's cache contents
-COPY --chown=onyx:onyx --from=tokenizers /app/.cache/huggingface /app/.cache/temp_huggingface
-COPY --chown=onyx:onyx --from=onyx-models /app/.cache/huggingface /app/.cache/temp_huggingface
 COPY --chown=onyx:onyx --from=embedding-models /app/.cache/huggingface /app/.cache/temp_huggingface
 
 WORKDIR /app

backend/alembic/README.md

@@ -7,8 +7,12 @@ Onyx migrations use a generic single-database configuration with an async dbapi.
 
 ## To generate new migrations:
 
-run from onyx/backend:
-`alembic revision --autogenerate -m <DESCRIPTION_OF_MIGRATION>`
+From onyx/backend, run:
+`alembic revision -m <DESCRIPTION_OF_MIGRATION>`
+
+Note: you cannot use the `--autogenerate` flag as the automatic schema parsing does not work.
+
+Manually populate the upgrade and downgrade in your new migration.
 
 More info can be found here: https://alembic.sqlalchemy.org/en/latest/autogenerate.html
 

backend/alembic/env.py

@@ -39,7 +39,9 @@ config = context.config
 if config.config_file_name is not None and config.attributes.get(
     "configure_logger", True
 ):
-    fileConfig(config.config_file_name)
+    # disable_existing_loggers=False prevents breaking pytest's caplog fixture
+    # See: https://pytest-alembic.readthedocs.io/en/latest/setup.html#caplog-issues
+    fileConfig(config.config_file_name, disable_existing_loggers=False)
 
 target_metadata = [Base.metadata, ResultModelBase.metadata]
 
@@ -460,8 +462,49 @@ def run_migrations_offline() -> None:
 
 
 def run_migrations_online() -> None:
-    logger.info("run_migrations_online starting.")
-    asyncio.run(run_async_migrations())
+    """Run migrations in 'online' mode.
+
+    Supports pytest-alembic by checking for a pre-configured connection
+    in context.config.attributes["connection"]. If present, uses that
+    connection/engine directly instead of creating a new async engine.
+    """
+    # Check if pytest-alembic is providing a connection/engine
+    connectable = context.config.attributes.get("connection", None)
+
+    if connectable is not None:
+        # pytest-alembic is providing an engine - use it directly
+        logger.info("run_migrations_online starting (pytest-alembic mode).")
+
+        # For pytest-alembic, we use the default schema (public)
+        schema_name = context.config.attributes.get(
+            "schema_name", POSTGRES_DEFAULT_SCHEMA
+        )
+
+        # pytest-alembic passes an Engine, we need to get a connection from it
+        with connectable.connect() as connection:
+            # Set search path for the schema
+            connection.execute(text(f'SET search_path TO "{schema_name}"'))
+
+            context.configure(
+                connection=connection,
+                target_metadata=target_metadata,  # type: ignore
+                include_object=include_object,
+                version_table_schema=schema_name,
+                include_schemas=True,
+                compare_type=True,
+                compare_server_default=True,
+                script_location=config.get_main_option("script_location"),
+            )
+
+            with context.begin_transaction():
+                context.run_migrations()
+
+            # Commit the transaction to ensure changes are visible to next migration
+            connection.commit()
+    else:
+        # Normal operation - use async migrations
+        logger.info("run_migrations_online starting.")
+        asyncio.run(run_async_migrations())
 
 
 if context.is_offline_mode():

backend/alembic/versions/18b5b2524446_add_is_clarification_to_chat_message.py

@@ -0,0 +1,29 @@
+"""add is_clarification to chat_message
+
+Revision ID: 18b5b2524446
+Revises: 87c52ec39f84
+Create Date: 2025-01-16
+
+"""
+
+from alembic import op
+import sqlalchemy as sa
+
+# revision identifiers, used by Alembic.
+revision = "18b5b2524446"
+down_revision = "87c52ec39f84"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    op.add_column(
+        "chat_message",
+        sa.Column(
+            "is_clarification", sa.Boolean(), nullable=False, server_default="false"
+        ),
+    )
+
+
+def downgrade() -> None:
+    op.drop_column("chat_message", "is_clarification")

backend/alembic/versions/23957775e5f5_remove_feedback_foreignkey_constraint.py

@@ -12,8 +12,8 @@ import sqlalchemy as sa
 # revision identifiers, used by Alembic.
 revision = "23957775e5f5"
 down_revision = "bc9771dccadf"
-branch_labels = None  # type: ignore
-depends_on = None  # type: ignore
+branch_labels = None
+depends_on = None
 
 
 def upgrade() -> None:

backend/alembic/versions/2a391f840e85_add_last_refreshed_at_mcp_server.py

@@ -0,0 +1,27 @@
+"""add last refreshed at mcp server
+
+Revision ID: 2a391f840e85
+Revises: 4cebcbc9b2ae
+Create Date: 2025-12-06 15:19:59.766066
+
+"""
+
+from alembic import op
+import sqlalchemy as sa
+
+# revision identifiers, used by Alembi.
+revision = "2a391f840e85"
+down_revision = "4cebcbc9b2ae"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    op.add_column(
+        "mcp_server",
+        sa.Column("last_refreshed_at", sa.DateTime(timezone=True), nullable=True),
+    )
+
+
+def downgrade() -> None:
+    op.drop_column("mcp_server", "last_refreshed_at")

backend/alembic/versions/4cebcbc9b2ae_add_tab_index_to_tool_call.py

@@ -0,0 +1,27 @@
+"""add tab_index to tool_call
+
+Revision ID: 4cebcbc9b2ae
+Revises: a1b2c3d4e5f6
+Create Date: 2025-12-16
+
+"""
+
+from alembic import op
+import sqlalchemy as sa
+
+# revision identifiers, used by Alembic.
+revision = "4cebcbc9b2ae"
+down_revision = "a1b2c3d4e5f6"
+branch_labels: None = None
+depends_on: None = None
+
+
+def upgrade() -> None:
+    op.add_column(
+        "tool_call",
+        sa.Column("tab_index", sa.Integer(), nullable=False, server_default="0"),
+    )
+
+
+def downgrade() -> None:
+    op.drop_column("tool_call", "tab_index")

backend/alembic/versions/4ea2c93919c1_add_type_to_credentials.py

@@ -62,6 +62,11 @@ def upgrade() -> None:
     )
     """
     )
+
+    # Drop the temporary table to avoid conflicts if migration runs again
+    # (e.g., during upgrade -> downgrade -> upgrade cycles in tests)
+    op.execute("DROP TABLE IF EXISTS temp_connector_credential")
+
     # If no exception was raised, alter the column
     op.alter_column("credential", "source", nullable=True)  # TODO modify
     # # ### end Alembic commands ###

backend/alembic/versions/4f8a2b3c1d9e_add_open_url_tool.py

@@ -0,0 +1,104 @@
+"""add_open_url_tool
+
+Revision ID: 4f8a2b3c1d9e
+Revises: a852cbe15577
+Create Date: 2025-11-24 12:00:00.000000
+
+"""
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = "4f8a2b3c1d9e"
+down_revision = "a852cbe15577"
+branch_labels = None
+depends_on = None
+
+
+OPEN_URL_TOOL = {
+    "name": "OpenURLTool",
+    "display_name": "Open URL",
+    "description": (
+        "The Open URL Action allows the agent to fetch and read contents of web pages."
+    ),
+    "in_code_tool_id": "OpenURLTool",
+    "enabled": True,
+}
+
+
+def upgrade() -> None:
+    conn = op.get_bind()
+
+    # Check if tool already exists
+    existing = conn.execute(
+        sa.text("SELECT id FROM tool WHERE in_code_tool_id = :in_code_tool_id"),
+        {"in_code_tool_id": OPEN_URL_TOOL["in_code_tool_id"]},
+    ).fetchone()
+
+    if existing:
+        tool_id = existing[0]
+        # Update existing tool
+        conn.execute(
+            sa.text(
+                """
+                UPDATE tool
+                SET name = :name,
+                    display_name = :display_name,
+                    description = :description
+                WHERE in_code_tool_id = :in_code_tool_id
+                """
+            ),
+            OPEN_URL_TOOL,
+        )
+    else:
+        # Insert new tool
+        conn.execute(
+            sa.text(
+                """
+                INSERT INTO tool (name, display_name, description, in_code_tool_id, enabled)
+                VALUES (:name, :display_name, :description, :in_code_tool_id, :enabled)
+                """
+            ),
+            OPEN_URL_TOOL,
+        )
+        # Get the newly inserted tool's id
+        result = conn.execute(
+            sa.text("SELECT id FROM tool WHERE in_code_tool_id = :in_code_tool_id"),
+            {"in_code_tool_id": OPEN_URL_TOOL["in_code_tool_id"]},
+        ).fetchone()
+        tool_id = result[0]  # type: ignore
+
+    # Associate the tool with all existing personas
+    # Get all persona IDs
+    persona_ids = conn.execute(sa.text("SELECT id FROM persona")).fetchall()
+
+    for (persona_id,) in persona_ids:
+        # Check if association already exists
+        exists = conn.execute(
+            sa.text(
+                """
+                SELECT 1 FROM persona__tool
+                WHERE persona_id = :persona_id AND tool_id = :tool_id
+                """
+            ),
+            {"persona_id": persona_id, "tool_id": tool_id},
+        ).fetchone()
+
+        if not exists:
+            conn.execute(
+                sa.text(
+                    """
+                    INSERT INTO persona__tool (persona_id, tool_id)
+                    VALUES (:persona_id, :tool_id)
+                    """
+                ),
+                {"persona_id": persona_id, "tool_id": tool_id},
+            )
+
+
+def downgrade() -> None:
+    # We don't remove the tool on downgrade since it's fine to have it around.
+    # If we upgrade again, it will be a no-op.
+    pass

backend/alembic/versions/5e6f7a8b9c0d_update_default_persona_prompt.py

@@ -0,0 +1,55 @@
+"""update_default_persona_prompt
+
+Revision ID: 5e6f7a8b9c0d
+Revises: 4f8a2b3c1d9e
+Create Date: 2025-11-30 12:00:00.000000
+
+"""
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = "5e6f7a8b9c0d"
+down_revision = "4f8a2b3c1d9e"
+branch_labels = None
+depends_on = None
+
+
+DEFAULT_PERSONA_ID = 0
+
+# ruff: noqa: E501, W605 start
+DEFAULT_SYSTEM_PROMPT = """
+You are a highly capable, thoughtful, and precise assistant. Your goal is to deeply understand the user's intent, ask clarifying questions when needed, think step-by-step through complex problems, provide clear and accurate answers, and proactively anticipate helpful follow-up information. Always prioritize being truthful, nuanced, insightful, and efficient.
+
+The current date is [[CURRENT_DATETIME]].{citation_reminder_or_empty}
+
+# Response Style
+You use different text styles, bolding, emojis (sparingly), block quotes, and other formatting to make your responses more readable and engaging.
+You use proper Markdown and LaTeX to format your responses for math, scientific, and chemical formulas, symbols, etc.: '$$\\n[expression]\\n$$' for standalone cases and '\\( [expression] \\)' when inline.
+For code you prefer to use Markdown and specify the language.
+You can use horizontal rules (---) to separate sections of your responses.
+You can use Markdown tables to format your responses for data, lists, and other structured information.
+""".lstrip()
+# ruff: noqa: E501, W605 end
+
+
+def upgrade() -> None:
+    conn = op.get_bind()
+    conn.execute(
+        sa.text(
+            """
+            UPDATE persona
+            SET system_prompt = :system_prompt
+            WHERE id = :persona_id
+            """
+        ),
+        {"system_prompt": DEFAULT_SYSTEM_PROMPT, "persona_id": DEFAULT_PERSONA_ID},
+    )
+
+
+def downgrade() -> None:
+    # We don't revert the system prompt on downgrade since we don't know
+    # what the previous value was. The new prompt is a reasonable default.
+    pass

backend/alembic/versions/7bd55f264e1b_add_display_name_to_model_configuration.py

@@ -0,0 +1,27 @@
+"""Add display_name to model_configuration
+
+Revision ID: 7bd55f264e1b
+Revises: e8f0d2a38171
+Create Date: 2025-12-04
+
+"""
+
+from alembic import op
+import sqlalchemy as sa
+
+# revision identifiers, used by Alembic.
+revision = "7bd55f264e1b"
+down_revision = "e8f0d2a38171"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    op.add_column(
+        "model_configuration",
+        sa.Column("display_name", sa.String(), nullable=True),
+    )
+
+
+def downgrade() -> None:
+    op.drop_column("model_configuration", "display_name")

backend/alembic/versions/7ed603b64d5a_add_mcp_server_and_connection_config_.py

@@ -42,13 +42,13 @@ def upgrade() -> None:
         sa.Column(
             "created_at",
             sa.DateTime(timezone=True),
-            server_default=sa.text("now()"),  # type: ignore
+            server_default=sa.text("now()"),
             nullable=False,
         ),
         sa.Column(
             "updated_at",
             sa.DateTime(timezone=True),
-            server_default=sa.text("now()"),  # type: ignore
+            server_default=sa.text("now()"),
             nullable=False,
         ),
     )
@@ -63,13 +63,13 @@ def upgrade() -> None:
         sa.Column(
             "created_at",
             sa.DateTime(timezone=True),
-            server_default=sa.text("now()"),  # type: ignore
+            server_default=sa.text("now()"),
             nullable=False,
         ),
         sa.Column(
             "updated_at",
             sa.DateTime(timezone=True),
-            server_default=sa.text("now()"),  # type: ignore
+            server_default=sa.text("now()"),
             nullable=False,
         ),
         sa.ForeignKeyConstraint(

backend/alembic/versions/87c52ec39f84_update_default_system_prompt.py

@@ -0,0 +1,55 @@
+"""update_default_system_prompt
+
+Revision ID: 87c52ec39f84
+Revises: 7bd55f264e1b
+Create Date: 2025-12-05 15:54:06.002452
+
+"""
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = "87c52ec39f84"
+down_revision = "7bd55f264e1b"
+branch_labels = None
+depends_on = None
+
+
+DEFAULT_PERSONA_ID = 0
+
+# ruff: noqa: E501, W605 start
+DEFAULT_SYSTEM_PROMPT = """
+You are a highly capable, thoughtful, and precise assistant. Your goal is to deeply understand the user's intent, ask clarifying questions when needed, think step-by-step through complex problems, provide clear and accurate answers, and proactively anticipate helpful follow-up information. Always prioritize being truthful, nuanced, insightful, and efficient.
+
+The current date is [[CURRENT_DATETIME]].[[CITATION_GUIDANCE]]
+
+# Response Style
+You use different text styles, bolding, emojis (sparingly), block quotes, and other formatting to make your responses more readable and engaging.
+You use proper Markdown and LaTeX to format your responses for math, scientific, and chemical formulas, symbols, etc.: '$$\\n[expression]\\n$$' for standalone cases and '\\( [expression] \\)' when inline.
+For code you prefer to use Markdown and specify the language.
+You can use horizontal rules (---) to separate sections of your responses.
+You can use Markdown tables to format your responses for data, lists, and other structured information.
+""".lstrip()
+# ruff: noqa: E501, W605 end
+
+
+def upgrade() -> None:
+    conn = op.get_bind()
+    conn.execute(
+        sa.text(
+            """
+            UPDATE persona
+            SET system_prompt = :system_prompt
+            WHERE id = :persona_id
+            """
+        ),
+        {"system_prompt": DEFAULT_SYSTEM_PROMPT, "persona_id": DEFAULT_PERSONA_ID},
+    )
+
+
+def downgrade() -> None:
+    # We don't revert the system prompt on downgrade since we don't know
+    # what the previous value was. The new prompt is a reasonable default.
+    pass

backend/alembic/versions/a01bf2971c5d_update_default_tool_descriptions.py

@@ -0,0 +1,62 @@
+"""update_default_tool_descriptions
+
+Revision ID: a01bf2971c5d
+Revises: 87c52ec39f84
+Create Date: 2025-12-16 15:21:25.656375
+
+"""
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = "a01bf2971c5d"
+down_revision = "18b5b2524446"
+branch_labels = None
+depends_on = None
+
+# new tool descriptions (12/2025)
+TOOL_DESCRIPTIONS = {
+    "SearchTool": "The Search Action allows the agent to search through connected knowledge to help build an answer.",
+    "ImageGenerationTool": (
+        "The Image Generation Action allows the agent to use DALL-E 3 or GPT-IMAGE-1 to generate images. "
+        "The action will be used when the user asks the agent to generate an image."
+    ),
+    "WebSearchTool": (
+        "The Web Search Action allows the agent "
+        "to perform internet searches for up-to-date information."
+    ),
+    "KnowledgeGraphTool": (
+        "The Knowledge Graph Search Action allows the agent to search the "
+        "Knowledge Graph for information. This tool can (for now) only be active in the KG Beta Agent, "
+        "and it requires the Knowledge Graph to be enabled."
+    ),
+    "OktaProfileTool": (
+        "The Okta Profile Action allows the agent to fetch the current user's information from Okta. "
+        "This may include the user's name, email, phone number, address, and other details such as their "
+        "manager and direct reports."
+    ),
+}
+
+
+def upgrade() -> None:
+    conn = op.get_bind()
+    conn.execute(sa.text("BEGIN"))
+
+    try:
+        for tool_id, description in TOOL_DESCRIPTIONS.items():
+            conn.execute(
+                sa.text(
+                    "UPDATE tool SET description = :description WHERE in_code_tool_id = :tool_id"
+                ),
+                {"description": description, "tool_id": tool_id},
+            )
+        conn.execute(sa.text("COMMIT"))
+    except Exception as e:
+        conn.execute(sa.text("ROLLBACK"))
+        raise e
+
+
+def downgrade() -> None:
+    pass

backend/alembic/versions/a1b2c3d4e5f6_add_license_table.py

@@ -0,0 +1,49 @@
+"""add license table
+
+Revision ID: a1b2c3d4e5f6
+Revises: a01bf2971c5d
+Create Date: 2025-12-04 10:00:00.000000
+
+"""
+
+from alembic import op
+import sqlalchemy as sa
+
+# revision identifiers, used by Alembic.
+revision = "a1b2c3d4e5f6"
+down_revision = "a01bf2971c5d"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    op.create_table(
+        "license",
+        sa.Column("id", sa.Integer(), primary_key=True),
+        sa.Column("license_data", sa.Text(), nullable=False),
+        sa.Column(
+            "created_at",
+            sa.DateTime(timezone=True),
+            server_default=sa.func.now(),
+            nullable=False,
+        ),
+        sa.Column(
+            "updated_at",
+            sa.DateTime(timezone=True),
+            server_default=sa.func.now(),
+            nullable=False,
+        ),
+    )
+
+    # Singleton pattern - only ever one row in this table
+    op.create_index(
+        "idx_license_singleton",
+        "license",
+        [sa.text("(true)")],
+        unique=True,
+    )
+
+
+def downgrade() -> None:
+    op.drop_index("idx_license_singleton", table_name="license")
+    op.drop_table("license")

backend/alembic/versions/a2b3c4d5e6f7_remove_fast_default_model_name.py

@@ -0,0 +1,27 @@
+"""Remove fast_default_model_name from llm_provider
+
+Revision ID: a2b3c4d5e6f7
+Revises: 2a391f840e85
+Create Date: 2024-12-17
+
+"""
+
+from alembic import op
+import sqlalchemy as sa
+
+# revision identifiers, used by Alembic.
+revision = "a2b3c4d5e6f7"
+down_revision = "2a391f840e85"
+branch_labels: None = None
+depends_on: None = None
+
+
+def upgrade() -> None:
+    op.drop_column("llm_provider", "fast_default_model_name")
+
+
+def downgrade() -> None:
+    op.add_column(
+        "llm_provider",
+        sa.Column("fast_default_model_name", sa.String(), nullable=True),
+    )

backend/alembic/versions/a852cbe15577_new_chat_history.py

@@ -0,0 +1,425 @@
+"""New Chat History
+
+Revision ID: a852cbe15577
+Revises: 6436661d5b65
+Create Date: 2025-11-08 15:16:37.781308
+
+"""
+
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = "a852cbe15577"
+down_revision = "6436661d5b65"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # 1. Drop old research/agent tables (CASCADE handles dependencies)
+    op.execute("DROP TABLE IF EXISTS research_agent_iteration_sub_step CASCADE")
+    op.execute("DROP TABLE IF EXISTS research_agent_iteration CASCADE")
+    op.execute("DROP TABLE IF EXISTS agent__sub_query__search_doc CASCADE")
+    op.execute("DROP TABLE IF EXISTS agent__sub_query CASCADE")
+    op.execute("DROP TABLE IF EXISTS agent__sub_question CASCADE")
+
+    # 2. ChatMessage table changes
+    # Rename columns and add FKs
+    op.alter_column(
+        "chat_message", "parent_message", new_column_name="parent_message_id"
+    )
+    op.create_foreign_key(
+        "fk_chat_message_parent_message_id",
+        "chat_message",
+        "chat_message",
+        ["parent_message_id"],
+        ["id"],
+    )
+    op.alter_column(
+        "chat_message",
+        "latest_child_message",
+        new_column_name="latest_child_message_id",
+    )
+    op.create_foreign_key(
+        "fk_chat_message_latest_child_message_id",
+        "chat_message",
+        "chat_message",
+        ["latest_child_message_id"],
+        ["id"],
+    )
+
+    # Add new column
+    op.add_column(
+        "chat_message", sa.Column("reasoning_tokens", sa.Text(), nullable=True)
+    )
+
+    # Drop old columns
+    op.drop_column("chat_message", "rephrased_query")
+    op.drop_column("chat_message", "alternate_assistant_id")
+    op.drop_column("chat_message", "overridden_model")
+    op.drop_column("chat_message", "is_agentic")
+    op.drop_column("chat_message", "refined_answer_improvement")
+    op.drop_column("chat_message", "research_type")
+    op.drop_column("chat_message", "research_plan")
+    op.drop_column("chat_message", "research_answer_purpose")
+
+    # 3. ToolCall table changes
+    # Drop the unique constraint first
+    op.drop_constraint("uq_tool_call_message_id", "tool_call", type_="unique")
+
+    # Delete orphaned tool_call rows (those without valid chat_message)
+    op.execute(
+        "DELETE FROM tool_call WHERE message_id NOT IN (SELECT id FROM chat_message)"
+    )
+
+    # Add chat_session_id as nullable first, populate, then make NOT NULL
+    op.add_column(
+        "tool_call",
+        sa.Column("chat_session_id", postgresql.UUID(as_uuid=True), nullable=True),
+    )
+
+    # Populate chat_session_id from the related chat_message
+    op.execute(
+        """
+        UPDATE tool_call
+        SET chat_session_id = chat_message.chat_session_id
+        FROM chat_message
+        WHERE tool_call.message_id = chat_message.id
+    """
+    )
+
+    # Now make it NOT NULL and add FK
+    op.alter_column("tool_call", "chat_session_id", nullable=False)
+    op.create_foreign_key(
+        "fk_tool_call_chat_session_id",
+        "tool_call",
+        "chat_session",
+        ["chat_session_id"],
+        ["id"],
+        ondelete="CASCADE",
+    )
+
+    # Rename message_id and make nullable, recreate FK with CASCADE
+    op.drop_constraint("tool_call_message_id_fkey", "tool_call", type_="foreignkey")
+    op.alter_column(
+        "tool_call",
+        "message_id",
+        new_column_name="parent_chat_message_id",
+        nullable=True,
+    )
+    op.create_foreign_key(
+        "fk_tool_call_parent_chat_message_id",
+        "tool_call",
+        "chat_message",
+        ["parent_chat_message_id"],
+        ["id"],
+        ondelete="CASCADE",
+    )
+
+    # Add parent_tool_call_id with FK
+    op.add_column(
+        "tool_call", sa.Column("parent_tool_call_id", sa.Integer(), nullable=True)
+    )
+    op.create_foreign_key(
+        "fk_tool_call_parent_tool_call_id",
+        "tool_call",
+        "tool_call",
+        ["parent_tool_call_id"],
+        ["id"],
+        ondelete="CASCADE",
+    )
+
+    # Add other new columns
+    op.add_column(
+        "tool_call",
+        sa.Column("turn_number", sa.Integer(), nullable=False, server_default="0"),
+    )
+    op.add_column(
+        "tool_call",
+        sa.Column("tool_call_id", sa.String(), nullable=False, server_default=""),
+    )
+    op.add_column("tool_call", sa.Column("reasoning_tokens", sa.Text(), nullable=True))
+    op.add_column(
+        "tool_call",
+        sa.Column("tool_call_tokens", sa.Integer(), nullable=False, server_default="0"),
+    )
+    op.add_column(
+        "tool_call",
+        sa.Column("generated_images", postgresql.JSONB(), nullable=True),
+    )
+
+    # Rename columns
+    op.alter_column(
+        "tool_call", "tool_arguments", new_column_name="tool_call_arguments"
+    )
+    op.alter_column("tool_call", "tool_result", new_column_name="tool_call_response")
+
+    # Change tool_call_response type from JSONB to Text
+    op.execute(
+        """
+        ALTER TABLE tool_call
+        ALTER COLUMN tool_call_response TYPE TEXT
+        USING tool_call_response::text
+    """
+    )
+
+    # Drop old columns
+    op.drop_column("tool_call", "tool_name")
+
+    # 4. Create new association table
+    op.create_table(
+        "tool_call__search_doc",
+        sa.Column("tool_call_id", sa.Integer(), nullable=False),
+        sa.Column("search_doc_id", sa.Integer(), nullable=False),
+        sa.ForeignKeyConstraint(["tool_call_id"], ["tool_call.id"], ondelete="CASCADE"),
+        sa.ForeignKeyConstraint(
+            ["search_doc_id"], ["search_doc.id"], ondelete="CASCADE"
+        ),
+        sa.PrimaryKeyConstraint("tool_call_id", "search_doc_id"),
+    )
+
+    # 5. Persona table change
+    op.add_column(
+        "persona",
+        sa.Column(
+            "replace_base_system_prompt",
+            sa.Boolean(),
+            nullable=False,
+            server_default="false",
+        ),
+    )
+
+
+def downgrade() -> None:
+    # Reverse persona changes
+    op.drop_column("persona", "replace_base_system_prompt")
+
+    # Drop new association table
+    op.drop_table("tool_call__search_doc")
+
+    # Reverse ToolCall changes
+    op.add_column(
+        "tool_call",
+        sa.Column("tool_name", sa.String(), nullable=False, server_default=""),
+    )
+
+    # Change tool_call_response back to JSONB
+    op.execute(
+        """
+        ALTER TABLE tool_call
+        ALTER COLUMN tool_call_response TYPE JSONB
+        USING tool_call_response::jsonb
+    """
+    )
+
+    op.alter_column("tool_call", "tool_call_response", new_column_name="tool_result")
+    op.alter_column(
+        "tool_call", "tool_call_arguments", new_column_name="tool_arguments"
+    )
+
+    op.drop_column("tool_call", "generated_images")
+    op.drop_column("tool_call", "tool_call_tokens")
+    op.drop_column("tool_call", "reasoning_tokens")
+    op.drop_column("tool_call", "tool_call_id")
+    op.drop_column("tool_call", "turn_number")
+
+    op.drop_constraint(
+        "fk_tool_call_parent_tool_call_id", "tool_call", type_="foreignkey"
+    )
+    op.drop_column("tool_call", "parent_tool_call_id")
+
+    op.drop_constraint(
+        "fk_tool_call_parent_chat_message_id", "tool_call", type_="foreignkey"
+    )
+    op.alter_column(
+        "tool_call",
+        "parent_chat_message_id",
+        new_column_name="message_id",
+        nullable=False,
+    )
+    op.create_foreign_key(
+        "tool_call_message_id_fkey",
+        "tool_call",
+        "chat_message",
+        ["message_id"],
+        ["id"],
+    )
+
+    op.drop_constraint("fk_tool_call_chat_session_id", "tool_call", type_="foreignkey")
+    op.drop_column("tool_call", "chat_session_id")
+
+    op.create_unique_constraint("uq_tool_call_message_id", "tool_call", ["message_id"])
+
+    # Reverse ChatMessage changes
+    # Note: research_answer_purpose and research_type were originally String columns,
+    # not Enum types (see migrations 5ae8240accb3 and f8a9b2c3d4e5)
+    op.add_column(
+        "chat_message",
+        sa.Column("research_answer_purpose", sa.String(), nullable=True),
+    )
+    op.add_column(
+        "chat_message", sa.Column("research_plan", postgresql.JSONB(), nullable=True)
+    )
+    op.add_column(
+        "chat_message",
+        sa.Column("research_type", sa.String(), nullable=True),
+    )
+    op.add_column(
+        "chat_message",
+        sa.Column("refined_answer_improvement", sa.Boolean(), nullable=True),
+    )
+    op.add_column(
+        "chat_message",
+        sa.Column("is_agentic", sa.Boolean(), nullable=False, server_default="false"),
+    )
+    op.add_column(
+        "chat_message", sa.Column("overridden_model", sa.String(), nullable=True)
+    )
+    op.add_column(
+        "chat_message", sa.Column("alternate_assistant_id", sa.Integer(), nullable=True)
+    )
+    # Recreate the FK constraint that was implicitly dropped when the column was dropped
+    op.create_foreign_key(
+        "fk_chat_message_persona",
+        "chat_message",
+        "persona",
+        ["alternate_assistant_id"],
+        ["id"],
+    )
+    op.add_column(
+        "chat_message", sa.Column("rephrased_query", sa.Text(), nullable=True)
+    )
+
+    op.drop_column("chat_message", "reasoning_tokens")
+
+    op.drop_constraint(
+        "fk_chat_message_latest_child_message_id", "chat_message", type_="foreignkey"
+    )
+    op.alter_column(
+        "chat_message",
+        "latest_child_message_id",
+        new_column_name="latest_child_message",
+    )
+
+    op.drop_constraint(
+        "fk_chat_message_parent_message_id", "chat_message", type_="foreignkey"
+    )
+    op.alter_column(
+        "chat_message", "parent_message_id", new_column_name="parent_message"
+    )
+
+    # Recreate agent sub question and sub query tables
+    op.create_table(
+        "agent__sub_question",
+        sa.Column("id", sa.Integer(), primary_key=True),
+        sa.Column("primary_question_id", sa.Integer(), nullable=False),
+        sa.Column("chat_session_id", postgresql.UUID(as_uuid=True), nullable=False),
+        sa.Column("sub_question", sa.Text(), nullable=False),
+        sa.Column("level", sa.Integer(), nullable=False),
+        sa.Column("level_question_num", sa.Integer(), nullable=False),
+        sa.Column(
+            "time_created",
+            sa.DateTime(timezone=True),
+            server_default=sa.text("now()"),
+            nullable=False,
+        ),
+        sa.Column("sub_answer", sa.Text(), nullable=False),
+        sa.Column("sub_question_doc_results", postgresql.JSONB(), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["primary_question_id"], ["chat_message.id"], ondelete="CASCADE"
+        ),
+        sa.ForeignKeyConstraint(["chat_session_id"], ["chat_session.id"]),
+        sa.PrimaryKeyConstraint("id"),
+    )
+
+    op.create_table(
+        "agent__sub_query",
+        sa.Column("id", sa.Integer(), primary_key=True),
+        sa.Column("parent_question_id", sa.Integer(), nullable=False),
+        sa.Column("chat_session_id", postgresql.UUID(as_uuid=True), nullable=False),
+        sa.Column("sub_query", sa.Text(), nullable=False),
+        sa.Column(
+            "time_created",
+            sa.DateTime(timezone=True),
+            server_default=sa.text("now()"),
+            nullable=False,
+        ),
+        sa.ForeignKeyConstraint(
+            ["parent_question_id"], ["agent__sub_question.id"], ondelete="CASCADE"
+        ),
+        sa.ForeignKeyConstraint(["chat_session_id"], ["chat_session.id"]),
+        sa.PrimaryKeyConstraint("id"),
+    )
+
+    op.create_table(
+        "agent__sub_query__search_doc",
+        sa.Column("sub_query_id", sa.Integer(), nullable=False),
+        sa.Column("search_doc_id", sa.Integer(), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["sub_query_id"], ["agent__sub_query.id"], ondelete="CASCADE"
+        ),
+        sa.ForeignKeyConstraint(["search_doc_id"], ["search_doc.id"]),
+        sa.PrimaryKeyConstraint("sub_query_id", "search_doc_id"),
+    )
+
+    # Recreate research agent tables
+    op.create_table(
+        "research_agent_iteration",
+        sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
+        sa.Column("primary_question_id", sa.Integer(), nullable=False),
+        sa.Column("iteration_nr", sa.Integer(), nullable=False),
+        sa.Column(
+            "created_at",
+            sa.DateTime(timezone=True),
+            server_default=sa.text("now()"),
+            nullable=False,
+        ),
+        sa.Column("purpose", sa.String(), nullable=True),
+        sa.Column("reasoning", sa.String(), nullable=True),
+        sa.ForeignKeyConstraint(
+            ["primary_question_id"], ["chat_message.id"], ondelete="CASCADE"
+        ),
+        sa.PrimaryKeyConstraint("id"),
+        sa.UniqueConstraint(
+            "primary_question_id",
+            "iteration_nr",
+            name="_research_agent_iteration_unique_constraint",
+        ),
+    )
+
+    op.create_table(
+        "research_agent_iteration_sub_step",
+        sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
+        sa.Column("primary_question_id", sa.Integer(), nullable=False),
+        sa.Column("iteration_nr", sa.Integer(), nullable=False),
+        sa.Column("iteration_sub_step_nr", sa.Integer(), nullable=False),
+        sa.Column(
+            "created_at",
+            sa.DateTime(timezone=True),
+            server_default=sa.text("now()"),
+            nullable=False,
+        ),
+        sa.Column("sub_step_instructions", sa.String(), nullable=True),
+        sa.Column("sub_step_tool_id", sa.Integer(), nullable=True),
+        sa.Column("reasoning", sa.String(), nullable=True),
+        sa.Column("sub_answer", sa.String(), nullable=True),
+        sa.Column("cited_doc_results", postgresql.JSONB(), nullable=False),
+        sa.Column("claims", postgresql.JSONB(), nullable=True),
+        sa.Column("is_web_fetch", sa.Boolean(), nullable=True),
+        sa.Column("queries", postgresql.JSONB(), nullable=True),
+        sa.Column("generated_images", postgresql.JSONB(), nullable=True),
+        sa.Column("additional_data", postgresql.JSONB(), nullable=True),
+        sa.Column("file_ids", postgresql.JSONB(), nullable=True),
+        sa.ForeignKeyConstraint(
+            ["primary_question_id", "iteration_nr"],
+            [
+                "research_agent_iteration.primary_question_id",
+                "research_agent_iteration.iteration_nr",
+            ],
+            ondelete="CASCADE",
+        ),
+        sa.ForeignKeyConstraint(["sub_step_tool_id"], ["tool.id"], ondelete="SET NULL"),
+        sa.PrimaryKeyConstraint("id"),
+    )

backend/alembic/versions/b8c9d0e1f2a3_drop_milestone_table.py

@@ -0,0 +1,46 @@
+"""Drop milestone table
+
+Revision ID: b8c9d0e1f2a3
+Revises: a2b3c4d5e6f7
+Create Date: 2025-12-18
+
+"""
+
+from alembic import op
+import sqlalchemy as sa
+import fastapi_users_db_sqlalchemy
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = "b8c9d0e1f2a3"
+down_revision = "a2b3c4d5e6f7"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    op.drop_table("milestone")
+
+
+def downgrade() -> None:
+    op.create_table(
+        "milestone",
+        sa.Column("id", sa.UUID(), nullable=False),
+        sa.Column("tenant_id", sa.String(), nullable=True),
+        sa.Column(
+            "user_id",
+            fastapi_users_db_sqlalchemy.generics.GUID(),
+            nullable=True,
+        ),
+        sa.Column("event_type", sa.String(), nullable=False),
+        sa.Column(
+            "time_created",
+            sa.DateTime(timezone=True),
+            server_default=sa.text("now()"),
+            nullable=False,
+        ),
+        sa.Column("event_tracker", postgresql.JSONB(), nullable=True),
+        sa.ForeignKeyConstraint(["user_id"], ["user.id"], ondelete="CASCADE"),
+        sa.PrimaryKeyConstraint("id"),
+        sa.UniqueConstraint("event_type", name="uq_milestone_event_type"),
+    )

backend/alembic/versions/c1d2e3f4a5b6_add_deep_research_tool.py

@@ -0,0 +1,52 @@
+"""add_deep_research_tool
+
+Revision ID: c1d2e3f4a5b6
+Revises: b8c9d0e1f2a3
+Create Date: 2025-12-18 16:00:00.000000
+
+"""
+
+from alembic import op
+from onyx.deep_research.dr_mock_tools import RESEARCH_AGENT_DB_NAME
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = "c1d2e3f4a5b6"
+down_revision = "b8c9d0e1f2a3"
+branch_labels = None
+depends_on = None
+
+
+DEEP_RESEARCH_TOOL = {
+    "name": RESEARCH_AGENT_DB_NAME,
+    "display_name": "Research Agent",
+    "description": "The Research Agent is a sub-agent that conducts research on a specific topic.",
+    "in_code_tool_id": "ResearchAgent",
+}
+
+
+def upgrade() -> None:
+    conn = op.get_bind()
+    conn.execute(
+        sa.text(
+            """
+            INSERT INTO tool (name, display_name, description, in_code_tool_id, enabled)
+            VALUES (:name, :display_name, :description, :in_code_tool_id, false)
+            """
+        ),
+        DEEP_RESEARCH_TOOL,
+    )
+
+
+def downgrade() -> None:
+    conn = op.get_bind()
+    conn.execute(
+        sa.text(
+            """
+            DELETE FROM tool
+            WHERE in_code_tool_id = :in_code_tool_id
+            """
+        ),
+        {"in_code_tool_id": DEEP_RESEARCH_TOOL["in_code_tool_id"]},
+    )

backend/alembic/versions/c9e2cd766c29_add_s3_file_store_table.py

@@ -257,8 +257,8 @@ def _migrate_files_to_external_storage() -> None:
             print(f"File {file_id} not found in PostgreSQL storage.")
             continue
 
-        lobj_id = cast(int, file_record.lobj_oid)  # type: ignore
-        file_metadata = cast(Any, file_record.file_metadata)  # type: ignore
+        lobj_id = cast(int, file_record.lobj_oid)
+        file_metadata = cast(Any, file_record.file_metadata)
 
         # Read file content from PostgreSQL
         try:
@@ -280,7 +280,7 @@ def _migrate_files_to_external_storage() -> None:
             else:
                 # Convert other types to dict if possible, otherwise None
                 try:
-                    file_metadata = dict(file_record.file_metadata)  # type: ignore
+                    file_metadata = dict(file_record.file_metadata)
                 except (TypeError, ValueError):
                     file_metadata = None
 

backend/alembic/versions/e209dc5a8156_added_prune_frequency.py

@@ -11,8 +11,8 @@ import sqlalchemy as sa
 
 revision = "e209dc5a8156"
 down_revision = "48d14957fe80"
-branch_labels = None  # type: ignore
-depends_on = None  # type: ignore
+branch_labels = None
+depends_on = None
 
 
 def upgrade() -> None:

backend/alembic/versions/e8f0d2a38171_add_status_to_mcp_server_and_make_auth_.py

@@ -0,0 +1,115 @@
+"""add status to mcp server and make auth fields nullable
+
+Revision ID: e8f0d2a38171
+Revises: ed9e44312505
+Create Date: 2025-11-28 11:15:37.667340
+
+"""
+
+from alembic import op
+import sqlalchemy as sa
+from onyx.db.enums import (
+    MCPTransport,
+    MCPAuthenticationType,
+    MCPAuthenticationPerformer,
+    MCPServerStatus,
+)
+
+# revision identifiers, used by Alembic.
+revision = "e8f0d2a38171"
+down_revision = "ed9e44312505"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # Make auth fields nullable
+    op.alter_column(
+        "mcp_server",
+        "transport",
+        existing_type=sa.Enum(MCPTransport, name="mcp_transport", native_enum=False),
+        nullable=True,
+    )
+
+    op.alter_column(
+        "mcp_server",
+        "auth_type",
+        existing_type=sa.Enum(
+            MCPAuthenticationType, name="mcp_authentication_type", native_enum=False
+        ),
+        nullable=True,
+    )
+
+    op.alter_column(
+        "mcp_server",
+        "auth_performer",
+        existing_type=sa.Enum(
+            MCPAuthenticationPerformer,
+            name="mcp_authentication_performer",
+            native_enum=False,
+        ),
+        nullable=True,
+    )
+
+    # Add status column with default
+    op.add_column(
+        "mcp_server",
+        sa.Column(
+            "status",
+            sa.Enum(MCPServerStatus, name="mcp_server_status", native_enum=False),
+            nullable=False,
+            server_default="CREATED",
+        ),
+    )
+
+    # For existing records, mark status as CONNECTED
+    bind = op.get_bind()
+    bind.execute(
+        sa.text(
+            """
+        UPDATE mcp_server
+        SET status = 'CONNECTED'
+        WHERE status != 'CONNECTED'
+        and admin_connection_config_id IS NOT NULL
+        """
+        )
+    )
+
+
+def downgrade() -> None:
+    # Remove status column
+    op.drop_column("mcp_server", "status")
+
+    # Make auth fields non-nullable (set defaults first)
+    op.execute(
+        "UPDATE mcp_server SET transport = 'STREAMABLE_HTTP' WHERE transport IS NULL"
+    )
+    op.execute("UPDATE mcp_server SET auth_type = 'NONE' WHERE auth_type IS NULL")
+    op.execute(
+        "UPDATE mcp_server SET auth_performer = 'ADMIN' WHERE auth_performer IS NULL"
+    )
+
+    op.alter_column(
+        "mcp_server",
+        "transport",
+        existing_type=sa.Enum(MCPTransport, name="mcp_transport", native_enum=False),
+        nullable=False,
+    )
+    op.alter_column(
+        "mcp_server",
+        "auth_type",
+        existing_type=sa.Enum(
+            MCPAuthenticationType, name="mcp_authentication_type", native_enum=False
+        ),
+        nullable=False,
+    )
+    op.alter_column(
+        "mcp_server",
+        "auth_performer",
+        existing_type=sa.Enum(
+            MCPAuthenticationPerformer,
+            name="mcp_authentication_performer",
+            native_enum=False,
+        ),
+        nullable=False,
+    )

backend/alembic/versions/ed9e44312505_add_icon_name_field.py

@@ -0,0 +1,34 @@
+"""Add icon_name field
+
+Revision ID: ed9e44312505
+Revises: 5e6f7a8b9c0d
+Create Date: 2025-12-03 16:35:07.828393
+
+"""
+
+from alembic import op
+import sqlalchemy as sa
+
+# revision identifiers, used by Alembic.
+revision = "ed9e44312505"
+down_revision = "5e6f7a8b9c0d"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # Add icon_name column
+    op.add_column("persona", sa.Column("icon_name", sa.String(), nullable=True))
+
+    # Remove old icon columns
+    op.drop_column("persona", "icon_shape")
+    op.drop_column("persona", "icon_color")
+
+
+def downgrade() -> None:
+    # Re-add old icon columns
+    op.add_column("persona", sa.Column("icon_color", sa.String(), nullable=True))
+    op.add_column("persona", sa.Column("icon_shape", sa.Integer(), nullable=True))
+
+    # Remove icon_name column
+    op.drop_column("persona", "icon_name")

backend/alembic_tenants/env.py

@@ -20,7 +20,9 @@ config = context.config
 if config.config_file_name is not None and config.attributes.get(
     "configure_logger", True
 ):
-    fileConfig(config.config_file_name)
+    # disable_existing_loggers=False prevents breaking pytest's caplog fixture
+    # See: https://pytest-alembic.readthedocs.io/en/latest/setup.html#caplog-issues
+    fileConfig(config.config_file_name, disable_existing_loggers=False)
 
 # add your model's MetaData object here
 # for 'autogenerate' support
@@ -82,9 +84,9 @@ def run_migrations_offline() -> None:
 def do_run_migrations(connection: Connection) -> None:
     context.configure(
         connection=connection,
-        target_metadata=target_metadata,  # type: ignore
+        target_metadata=target_metadata,  # type: ignore[arg-type]
         include_object=include_object,
-    )  # type: ignore
+    )
 
     with context.begin_transaction():
         context.run_migrations()
@@ -108,9 +110,24 @@ async def run_async_migrations() -> None:
 
 
 def run_migrations_online() -> None:
-    """Run migrations in 'online' mode."""
+    """Run migrations in 'online' mode.
 
-    asyncio.run(run_async_migrations())
+    Supports pytest-alembic by checking for a pre-configured connection
+    in context.config.attributes["connection"]. If present, uses that
+    connection/engine directly instead of creating a new async engine.
+    """
+    # Check if pytest-alembic is providing a connection/engine
+    connectable = context.config.attributes.get("connection", None)
+
+    if connectable is not None:
+        # pytest-alembic is providing an engine - use it directly
+        with connectable.connect() as connection:
+            do_run_migrations(connection)
+            # Commit to ensure changes are visible to next migration
+            connection.commit()
+    else:
+        # Normal operation - use async migrations
+        asyncio.run(run_async_migrations())
 
 
 if context.is_offline_mode():

backend/ee/onyx/configs/app_configs.py

@@ -41,6 +41,10 @@ CONFLUENCE_ANONYMOUS_ACCESS_IS_PUBLIC = (
 JIRA_PERMISSION_DOC_SYNC_FREQUENCY = int(
     os.environ.get("JIRA_PERMISSION_DOC_SYNC_FREQUENCY") or 30 * 60
 )
+# In seconds, default is 30 minutes
+JIRA_PERMISSION_GROUP_SYNC_FREQUENCY = int(
+    os.environ.get("JIRA_PERMISSION_GROUP_SYNC_FREQUENCY") or 30 * 60
+)
 
 
 #####

backend/ee/onyx/db/analytics.py

@@ -199,10 +199,7 @@ def fetch_persona_message_analytics(
             ChatMessage.chat_session_id == ChatSession.id,
         )
         .where(
-            or_(
-                ChatMessage.alternate_assistant_id == persona_id,
-                ChatSession.persona_id == persona_id,
-            ),
+            ChatSession.persona_id == persona_id,
             ChatMessage.time_sent >= start,
             ChatMessage.time_sent <= end,
             ChatMessage.message_type == MessageType.ASSISTANT,
@@ -231,10 +228,7 @@ def fetch_persona_unique_users(
             ChatMessage.chat_session_id == ChatSession.id,
         )
         .where(
-            or_(
-                ChatMessage.alternate_assistant_id == persona_id,
-                ChatSession.persona_id == persona_id,
-            ),
+            ChatSession.persona_id == persona_id,
             ChatMessage.time_sent >= start,
             ChatMessage.time_sent <= end,
             ChatMessage.message_type == MessageType.ASSISTANT,
@@ -265,10 +259,7 @@ def fetch_assistant_message_analytics(
             ChatMessage.chat_session_id == ChatSession.id,
         )
         .where(
-            or_(
-                ChatMessage.alternate_assistant_id == assistant_id,
-                ChatSession.persona_id == assistant_id,
-            ),
+            ChatSession.persona_id == assistant_id,
             ChatMessage.time_sent >= start,
             ChatMessage.time_sent <= end,
             ChatMessage.message_type == MessageType.ASSISTANT,
@@ -299,10 +290,7 @@ def fetch_assistant_unique_users(
             ChatMessage.chat_session_id == ChatSession.id,
         )
         .where(
-            or_(
-                ChatMessage.alternate_assistant_id == assistant_id,
-                ChatSession.persona_id == assistant_id,
-            ),
+            ChatSession.persona_id == assistant_id,
             ChatMessage.time_sent >= start,
             ChatMessage.time_sent <= end,
             ChatMessage.message_type == MessageType.ASSISTANT,
@@ -332,10 +320,7 @@ def fetch_assistant_unique_users_total(
             ChatMessage.chat_session_id == ChatSession.id,
         )
         .where(
-            or_(
-                ChatMessage.alternate_assistant_id == assistant_id,
-                ChatSession.persona_id == assistant_id,
-            ),
+            ChatSession.persona_id == assistant_id,
             ChatMessage.time_sent >= start,
             ChatMessage.time_sent <= end,
             ChatMessage.message_type == MessageType.ASSISTANT,

backend/ee/onyx/db/document_set.py

@@ -118,6 +118,6 @@ def fetch_document_sets(
             .all()
         )
 
-        document_set_with_cc_pairs.append((document_set, cc_pairs))  # type: ignore
+        document_set_with_cc_pairs.append((document_set, cc_pairs))
 
     return document_set_with_cc_pairs

backend/ee/onyx/db/license.py

@@ -0,0 +1,278 @@
+"""Database and cache operations for the license table."""
+
+from datetime import datetime
+
+from sqlalchemy import func
+from sqlalchemy import select
+from sqlalchemy.orm import Session
+
+from ee.onyx.server.license.models import LicenseMetadata
+from ee.onyx.server.license.models import LicensePayload
+from ee.onyx.server.license.models import LicenseSource
+from onyx.db.models import License
+from onyx.db.models import User
+from onyx.redis.redis_pool import get_redis_client
+from onyx.redis.redis_pool import get_redis_replica_client
+from onyx.utils.logger import setup_logger
+from shared_configs.configs import MULTI_TENANT
+from shared_configs.contextvars import get_current_tenant_id
+
+logger = setup_logger()
+
+LICENSE_METADATA_KEY = "license:metadata"
+LICENSE_CACHE_TTL_SECONDS = 86400  # 24 hours
+
+
+# -----------------------------------------------------------------------------
+# Database CRUD Operations
+# -----------------------------------------------------------------------------
+
+
+def get_license(db_session: Session) -> License | None:
+    """
+    Get the current license (singleton pattern - only one row).
+
+    Args:
+        db_session: Database session
+
+    Returns:
+        License object if exists, None otherwise
+    """
+    return db_session.execute(select(License)).scalars().first()
+
+
+def upsert_license(db_session: Session, license_data: str) -> License:
+    """
+    Insert or update the license (singleton pattern).
+
+    Args:
+        db_session: Database session
+        license_data: Base64-encoded signed license blob
+
+    Returns:
+        The created or updated License object
+    """
+    existing = get_license(db_session)
+
+    if existing:
+        existing.license_data = license_data
+        db_session.commit()
+        db_session.refresh(existing)
+        logger.info("License updated")
+        return existing
+
+    new_license = License(license_data=license_data)
+    db_session.add(new_license)
+    db_session.commit()
+    db_session.refresh(new_license)
+    logger.info("License created")
+    return new_license
+
+
+def delete_license(db_session: Session) -> bool:
+    """
+    Delete the current license.
+
+    Args:
+        db_session: Database session
+
+    Returns:
+        True if deleted, False if no license existed
+    """
+    existing = get_license(db_session)
+    if existing:
+        db_session.delete(existing)
+        db_session.commit()
+        logger.info("License deleted")
+        return True
+    return False
+
+
+# -----------------------------------------------------------------------------
+# Seat Counting
+# -----------------------------------------------------------------------------
+
+
+def get_used_seats(tenant_id: str | None = None) -> int:
+    """
+    Get current seat usage.
+
+    For multi-tenant: counts users in UserTenantMapping for this tenant.
+    For self-hosted: counts all active users (includes both Onyx UI users
+    and Slack users who have been converted to Onyx users).
+    """
+    if MULTI_TENANT:
+        from ee.onyx.server.tenants.user_mapping import get_tenant_count
+
+        return get_tenant_count(tenant_id or get_current_tenant_id())
+    else:
+        # Self-hosted: count all active users (Onyx + converted Slack users)
+        from onyx.db.engine.sql_engine import get_session_with_current_tenant
+
+        with get_session_with_current_tenant() as db_session:
+            result = db_session.execute(
+                select(func.count()).select_from(User).where(User.is_active)  # type: ignore
+            )
+            return result.scalar() or 0
+
+
+# -----------------------------------------------------------------------------
+# Redis Cache Operations
+# -----------------------------------------------------------------------------
+
+
+def get_cached_license_metadata(tenant_id: str | None = None) -> LicenseMetadata | None:
+    """
+    Get license metadata from Redis cache.
+
+    Args:
+        tenant_id: Tenant ID (for multi-tenant deployments)
+
+    Returns:
+        LicenseMetadata if cached, None otherwise
+    """
+    tenant = tenant_id or get_current_tenant_id()
+    redis_client = get_redis_replica_client(tenant_id=tenant)
+
+    cached = redis_client.get(LICENSE_METADATA_KEY)
+    if cached:
+        try:
+            cached_str: str
+            if isinstance(cached, bytes):
+                cached_str = cached.decode("utf-8")
+            else:
+                cached_str = str(cached)
+            return LicenseMetadata.model_validate_json(cached_str)
+        except Exception as e:
+            logger.warning(f"Failed to parse cached license metadata: {e}")
+            return None
+    return None
+
+
+def invalidate_license_cache(tenant_id: str | None = None) -> None:
+    """
+    Invalidate the license metadata cache (not the license itself).
+
+    This deletes the cached LicenseMetadata from Redis. The actual license
+    in the database is not affected. Redis delete is idempotent - if the
+    key doesn't exist, this is a no-op.
+
+    Args:
+        tenant_id: Tenant ID (for multi-tenant deployments)
+    """
+    tenant = tenant_id or get_current_tenant_id()
+    redis_client = get_redis_client(tenant_id=tenant)
+    redis_client.delete(LICENSE_METADATA_KEY)
+    logger.info("License cache invalidated")
+
+
+def update_license_cache(
+    payload: LicensePayload,
+    source: LicenseSource | None = None,
+    grace_period_end: datetime | None = None,
+    tenant_id: str | None = None,
+) -> LicenseMetadata:
+    """
+    Update the Redis cache with license metadata.
+
+    We cache all license statuses (ACTIVE, GRACE_PERIOD, GATED_ACCESS) because:
+    1. Frontend needs status to show appropriate UI/banners
+    2. Caching avoids repeated DB + crypto verification on every request
+    3. Status enforcement happens at the feature level, not here
+
+    Args:
+        payload: Verified license payload
+        source: How the license was obtained
+        grace_period_end: Optional grace period end time
+        tenant_id: Tenant ID (for multi-tenant deployments)
+
+    Returns:
+        The cached LicenseMetadata
+    """
+    from ee.onyx.utils.license import get_license_status
+
+    tenant = tenant_id or get_current_tenant_id()
+    redis_client = get_redis_client(tenant_id=tenant)
+
+    used_seats = get_used_seats(tenant)
+    status = get_license_status(payload, grace_period_end)
+
+    metadata = LicenseMetadata(
+        tenant_id=payload.tenant_id,
+        organization_name=payload.organization_name,
+        seats=payload.seats,
+        used_seats=used_seats,
+        plan_type=payload.plan_type,
+        issued_at=payload.issued_at,
+        expires_at=payload.expires_at,
+        grace_period_end=grace_period_end,
+        status=status,
+        source=source,
+        stripe_subscription_id=payload.stripe_subscription_id,
+    )
+
+    redis_client.setex(
+        LICENSE_METADATA_KEY,
+        LICENSE_CACHE_TTL_SECONDS,
+        metadata.model_dump_json(),
+    )
+
+    logger.info(f"License cache updated: {metadata.seats} seats, status={status.value}")
+    return metadata
+
+
+def refresh_license_cache(
+    db_session: Session,
+    tenant_id: str | None = None,
+) -> LicenseMetadata | None:
+    """
+    Refresh the license cache from the database.
+
+    Args:
+        db_session: Database session
+        tenant_id: Tenant ID (for multi-tenant deployments)
+
+    Returns:
+        LicenseMetadata if license exists, None otherwise
+    """
+    from ee.onyx.utils.license import verify_license_signature
+
+    license_record = get_license(db_session)
+    if not license_record:
+        invalidate_license_cache(tenant_id)
+        return None
+
+    try:
+        payload = verify_license_signature(license_record.license_data)
+        return update_license_cache(
+            payload,
+            source=LicenseSource.AUTO_FETCH,
+            tenant_id=tenant_id,
+        )
+    except ValueError as e:
+        logger.error(f"Failed to verify license during cache refresh: {e}")
+        invalidate_license_cache(tenant_id)
+        return None
+
+
+def get_license_metadata(
+    db_session: Session,
+    tenant_id: str | None = None,
+) -> LicenseMetadata | None:
+    """
+    Get license metadata, using cache if available.
+
+    Args:
+        db_session: Database session
+        tenant_id: Tenant ID (for multi-tenant deployments)
+
+    Returns:
+        LicenseMetadata if license exists, None otherwise
+    """
+    # Try cache first
+    cached = get_cached_license_metadata(tenant_id)
+    if cached:
+        return cached
+
+    # Refresh from database
+    return refresh_license_cache(db_session, tenant_id)

backend/ee/onyx/db/usage_export.py

@@ -55,18 +55,7 @@ def get_empty_chat_messages_entries__paginated(
 
             # Get assistant name (from session persona, or alternate if specified)
             assistant_name = None
-            if message.alternate_assistant_id:
-                # If there's an alternate assistant, we need to fetch it
-                from onyx.db.models import Persona
-
-                alternate_persona = (
-                    db_session.query(Persona)
-                    .filter(Persona.id == message.alternate_assistant_id)
-                    .first()
-                )
-                if alternate_persona:
-                    assistant_name = alternate_persona.name
-            elif chat_session.persona:
+            if chat_session.persona:
                 assistant_name = chat_session.persona.name
 
             message_skeletons.append(

backend/ee/onyx/db/user_group.py

@@ -8,6 +8,7 @@ from sqlalchemy import func
 from sqlalchemy import Select
 from sqlalchemy import select
 from sqlalchemy import update
+from sqlalchemy.dialects.postgresql import insert
 from sqlalchemy.orm import Session
 
 from ee.onyx.server.user_group.models import SetCuratorRequest
@@ -362,14 +363,29 @@ def _check_user_group_is_modifiable(user_group: UserGroup) -> None:
 
 def _add_user__user_group_relationships__no_commit(
     db_session: Session, user_group_id: int, user_ids: list[UUID]
-) -> list[User__UserGroup]:
-    """NOTE: does not commit the transaction."""
-    relationships = [
-        User__UserGroup(user_id=user_id, user_group_id=user_group_id)
-        for user_id in user_ids
-    ]
-    db_session.add_all(relationships)
-    return relationships
+) -> None:
+    """NOTE: does not commit the transaction.
+
+    This function is idempotent - it will skip users who are already in the group
+    to avoid duplicate key violations during concurrent operations or re-syncs.
+    Uses ON CONFLICT DO NOTHING to keep inserts atomic under concurrency.
+    """
+    if not user_ids:
+        return
+
+    insert_stmt = (
+        insert(User__UserGroup)
+        .values(
+            [
+                {"user_id": user_id, "user_group_id": user_group_id}
+                for user_id in user_ids
+            ]
+        )
+        .on_conflict_do_nothing(
+            index_elements=[User__UserGroup.user_group_id, User__UserGroup.user_id]
+        )
+    )
+    db_session.execute(insert_stmt)
 
 
 def _add_user_group__cc_pair_relationships__no_commit(

backend/ee/onyx/external_permissions/confluence/group_sync.py

@@ -3,12 +3,15 @@ from collections.abc import Generator
 from ee.onyx.db.external_perm import ExternalUserGroup
 from ee.onyx.external_permissions.confluence.constants import ALL_CONF_EMAILS_GROUP_NAME
 from onyx.background.error_logging import emit_background_error
+from onyx.configs.app_configs import CONFLUENCE_USE_ONYX_USERS_FOR_GROUP_SYNC
 from onyx.connectors.confluence.onyx_confluence import (
     get_user_email_from_username__server,
 )
 from onyx.connectors.confluence.onyx_confluence import OnyxConfluence
 from onyx.connectors.credentials_provider import OnyxDBCredentialsProvider
+from onyx.db.engine.sql_engine import get_session_with_current_tenant
 from onyx.db.models import ConnectorCredentialPair
+from onyx.db.users import get_all_users
 from onyx.utils.logger import setup_logger
 
 logger = setup_logger()
@@ -19,7 +22,7 @@ def _build_group_member_email_map(
 ) -> dict[str, set[str]]:
     group_member_emails: dict[str, set[str]] = {}
     for user in confluence_client.paginated_cql_user_retrieval():
-        logger.debug(f"Processing groups for user: {user}")
+        logger.info(f"Processing groups for user: {user}")
 
         email = user.email
         if not email:
@@ -31,6 +34,8 @@ def _build_group_member_email_map(
                     confluence_client=confluence_client,
                     user_name=user_name,
                 )
+            else:
+                logger.error(f"user result missing username field: {user}")
 
         if not email:
             # If we still don't have an email, skip this user
@@ -64,6 +69,92 @@ def _build_group_member_email_map(
     return group_member_emails
 
 
+def _build_group_member_email_map_from_onyx_users(
+    confluence_client: OnyxConfluence,
+) -> dict[str, set[str]]:
+    """Hacky, but it's the only way to do this as long as the
+    Confluence APIs are broken.
+
+    This is fixed in Confluence Data Center 10.1.0, so first choice
+    is to tell users to upgrade to 10.1.0.
+    https://jira.atlassian.com/browse/CONFSERVER-95999
+    """
+    with get_session_with_current_tenant() as db_session:
+        # don't include external since they are handled by the "through confluence"
+        # user fetching mechanism
+        user_emails = [
+            user.email for user in get_all_users(db_session, include_external=False)
+        ]
+
+    def _infer_username_from_email(email: str) -> str:
+        return email.split("@")[0]
+
+    group_member_emails: dict[str, set[str]] = {}
+    for email in user_emails:
+        logger.info(f"Processing groups for user with email: {email}")
+        try:
+            user_name = _infer_username_from_email(email)
+            response = confluence_client.get_user_details_by_username(user_name)
+            user_key = response.get("userKey")
+            if not user_key:
+                logger.error(f"User key not found for user with email {email}")
+                continue
+
+            all_users_groups: set[str] = set()
+            for group in confluence_client.paginated_groups_by_user_retrieval(user_key):
+                # group name uniqueness is enforced by Confluence, so we can use it as a group ID
+                group_id = group["name"]
+                group_member_emails.setdefault(group_id, set()).add(email)
+                all_users_groups.add(group_id)
+
+            if not all_users_groups:
+                msg = f"No groups found for user with email: {email}"
+                logger.error(msg)
+            else:
+                logger.info(
+                    f"Found groups {all_users_groups} for user with email {email}"
+                )
+        except Exception:
+            logger.exception(f"Error getting user details for user with email {email}")
+
+    return group_member_emails
+
+
+def _build_final_group_to_member_email_map(
+    confluence_client: OnyxConfluence,
+    cc_pair_id: int,
+    # if set, will infer confluence usernames from onyx users in addition to using the
+    # confluence users API. This is a hacky workaround for the fact that the Confluence
+    # users API is broken before Confluence Data Center 10.1.0.
+    use_onyx_users: bool = CONFLUENCE_USE_ONYX_USERS_FOR_GROUP_SYNC,
+) -> dict[str, set[str]]:
+    group_to_member_email_map = _build_group_member_email_map(
+        confluence_client=confluence_client,
+        cc_pair_id=cc_pair_id,
+    )
+    group_to_member_email_map_from_onyx_users = (
+        (
+            _build_group_member_email_map_from_onyx_users(
+                confluence_client=confluence_client,
+            )
+        )
+        if use_onyx_users
+        else {}
+    )
+
+    all_group_ids = set(group_to_member_email_map.keys()) | set(
+        group_to_member_email_map_from_onyx_users.keys()
+    )
+    final_group_to_member_email_map = {}
+    for group_id in all_group_ids:
+        group_member_emails = group_to_member_email_map.get(
+            group_id, set()
+        ) | group_to_member_email_map_from_onyx_users.get(group_id, set())
+        final_group_to_member_email_map[group_id] = group_member_emails
+
+    return final_group_to_member_email_map
+
+
 def confluence_group_sync(
     tenant_id: str,
     cc_pair: ConnectorCredentialPair,
@@ -87,13 +178,12 @@ def confluence_group_sync(
     confluence_client._probe_connection(**probe_kwargs)
     confluence_client._initialize_connection(**final_kwargs)
 
-    group_member_email_map = _build_group_member_email_map(
-        confluence_client=confluence_client,
-        cc_pair_id=cc_pair.id,
+    group_to_member_email_map = _build_final_group_to_member_email_map(
+        confluence_client, cc_pair.id
     )
 
     all_found_emails = set()
-    for group_id, group_member_emails in group_member_email_map.items():
+    for group_id, group_member_emails in group_to_member_email_map.items():
         yield (
             ExternalUserGroup(
                 id=group_id,

backend/ee/onyx/external_permissions/jira/group_sync.py

@@ -0,0 +1,136 @@
+from collections.abc import Generator
+
+from jira import JIRA
+
+from ee.onyx.db.external_perm import ExternalUserGroup
+from onyx.connectors.jira.utils import build_jira_client
+from onyx.db.models import ConnectorCredentialPair
+from onyx.utils.logger import setup_logger
+
+logger = setup_logger()
+
+
+def _get_jira_group_members_email(
+    jira_client: JIRA,
+    group_name: str,
+) -> list[str]:
+    """Get all member emails for a Jira group.
+
+    Filters out app accounts (bots, integrations) and only returns real user emails.
+    """
+    emails: list[str] = []
+
+    try:
+        # group_members returns an OrderedDict of account_id -> member_info
+        members = jira_client.group_members(group=group_name)
+
+        if not members:
+            logger.warning(f"No members found for group {group_name}")
+            return emails
+
+        for account_id, member_info in members.items():
+            # member_info is a dict with keys like 'fullname', 'email', 'active'
+            email = member_info.get("email")
+
+            # Skip "hidden" emails - these are typically app accounts
+            if email and email != "hidden":
+                emails.append(email)
+            else:
+                # For cloud, we might need to fetch user details separately
+                try:
+                    user = jira_client.user(id=account_id)
+
+                    # Skip app accounts (bots, integrations, etc.)
+                    if hasattr(user, "accountType") and user.accountType == "app":
+                        logger.info(
+                            f"Skipping app account {account_id} for group {group_name}"
+                        )
+                        continue
+
+                    if hasattr(user, "emailAddress") and user.emailAddress:
+                        emails.append(user.emailAddress)
+                    else:
+                        logger.warning(f"User {account_id} has no email address")
+                except Exception as e:
+                    logger.warning(
+                        f"Could not fetch email for user {account_id} in group {group_name}: {e}"
+                    )
+
+    except Exception as e:
+        logger.error(f"Error fetching members for group {group_name}: {e}")
+
+    return emails
+
+
+def _build_group_member_email_map(
+    jira_client: JIRA,
+) -> dict[str, set[str]]:
+    """Build a map of group names to member emails."""
+    group_member_emails: dict[str, set[str]] = {}
+
+    try:
+        # Get all groups from Jira - returns a list of group name strings
+        group_names = jira_client.groups()
+
+        if not group_names:
+            logger.warning("No groups found in Jira")
+            return group_member_emails
+
+        logger.info(f"Found {len(group_names)} groups in Jira")
+
+        for group_name in group_names:
+            if not group_name:
+                continue
+
+            member_emails = _get_jira_group_members_email(
+                jira_client=jira_client,
+                group_name=group_name,
+            )
+
+            if member_emails:
+                group_member_emails[group_name] = set(member_emails)
+                logger.debug(
+                    f"Found {len(member_emails)} members for group {group_name}"
+                )
+            else:
+                logger.debug(f"No members found for group {group_name}")
+
+    except Exception as e:
+        logger.error(f"Error building group member email map: {e}")
+
+    return group_member_emails
+
+
+def jira_group_sync(
+    tenant_id: str,
+    cc_pair: ConnectorCredentialPair,
+) -> Generator[ExternalUserGroup, None, None]:
+    """
+    Sync Jira groups and their members.
+
+    This function fetches all groups from Jira and yields ExternalUserGroup
+    objects containing the group ID and member emails.
+    """
+    jira_base_url = cc_pair.connector.connector_specific_config.get("jira_base_url", "")
+    scoped_token = cc_pair.connector.connector_specific_config.get(
+        "scoped_token", False
+    )
+
+    if not jira_base_url:
+        raise ValueError("No jira_base_url found in connector config")
+
+    jira_client = build_jira_client(
+        credentials=cc_pair.credential.credential_json,
+        jira_base=jira_base_url,
+        scoped_token=scoped_token,
+    )
+
+    group_member_email_map = _build_group_member_email_map(jira_client=jira_client)
+    if not group_member_email_map:
+        raise ValueError(f"No groups with members found for cc_pair_id={cc_pair.id}")
+
+    for group_id, group_member_emails in group_member_email_map.items():
+        yield ExternalUserGroup(
+            id=group_id,
+            user_emails=list(group_member_emails),
+        )

backend/ee/onyx/external_permissions/jira/page_access.py

@@ -16,6 +16,10 @@ HolderMap = dict[str, list[Holder]]
 logger = setup_logger()
 
 
+def _get_role_id(holder: Holder) -> str | None:
+    return holder.get("value") or holder.get("parameter")
+
+
 def _build_holder_map(permissions: list[dict]) -> dict[str, list[Holder]]:
     """
     A "Holder" in JIRA is a person / entity who "holds" the corresponding permission.
@@ -110,80 +114,137 @@ def _get_user_emails(user_holders: list[Holder]) -> list[str]:
     return emails
 
 
-def _get_user_emails_from_project_roles(
+def _get_user_emails_and_groups_from_project_roles(
     jira_client: JIRA,
     jira_project: str,
     project_role_holders: list[Holder],
-) -> list[str]:
-    # NOTE (@raunakab) a `parallel_yield` may be helpful here...?
+) -> tuple[list[str], list[str]]:
+    """
+    Get user emails and group names from project roles.
+    Returns a tuple of (emails, group_names).
+    """
+    # Get role IDs - Cloud uses "value", Data Center uses "parameter"
+    role_ids = []
+    for holder in project_role_holders:
+        role_id = _get_role_id(holder)
+        if role_id:
+            role_ids.append(role_id)
+        else:
+            logger.warning(f"No value or parameter in projectRole holder: {holder}")
+
     roles = [
-        jira_client.project_role(project=jira_project, id=project_role_holder["value"])
-        for project_role_holder in project_role_holders
-        if "value" in project_role_holder
+        jira_client.project_role(project=jira_project, id=role_id)
+        for role_id in role_ids
     ]
 
     emails = []
+    groups = []
 
     for role in roles:
         if not hasattr(role, "actors"):
+            logger.warning(f"Project role {role} has no actors attribute")
             continue
 
         for actor in role.actors:
-            if not hasattr(actor, "actorUser") or not hasattr(
-                actor.actorUser, "accountId"
-            ):
+            # Handle group actors
+            if hasattr(actor, "actorGroup"):
+                group_name = getattr(actor.actorGroup, "name", None) or getattr(
+                    actor.actorGroup, "displayName", None
+                )
+                if group_name:
+                    groups.append(group_name)
                 continue
 
-            user = jira_client.user(id=actor.actorUser.accountId)
-            if not hasattr(user, "accountType") or user.accountType != "atlassian":
+            # Handle user actors
+            if hasattr(actor, "actorUser"):
+                account_id = getattr(actor.actorUser, "accountId", None)
+                if not account_id:
+                    logger.error(f"No accountId in actorUser: {actor.actorUser}")
+                    continue
+
+                user = jira_client.user(id=account_id)
+                if not hasattr(user, "accountType") or user.accountType != "atlassian":
+                    logger.info(
+                        f"Skipping user {account_id} because it is not an atlassian user"
+                    )
+                    continue
+
+                if not hasattr(user, "emailAddress"):
+                    msg = f"User's email address was not able to be retrieved;  {actor.actorUser.accountId=}"
+                    if hasattr(user, "displayName"):
+                        msg += f" {actor.displayName=}"
+                    logger.warning(msg)
+                    continue
+
+                emails.append(user.emailAddress)
                 continue
 
-            if not hasattr(user, "emailAddress"):
-                msg = f"User's email address was not able to be retrieved;  {actor.actorUser.accountId=}"
-                if hasattr(user, "displayName"):
-                    msg += f" {actor.displayName=}"
-                logger.warn(msg)
-                continue
+            logger.debug(f"Skipping actor type: {actor}")
 
-            emails.append(user.emailAddress)
-
-    return emails
+    return emails, groups
 
 
 def _build_external_access_from_holder_map(
     jira_client: JIRA, jira_project: str, holder_map: HolderMap
 ) -> ExternalAccess:
     """
-    # Note:
-        If the `holder_map` contains an instance of "anyone", then this is a public JIRA project.
-        Otherwise, we fetch the "projectRole"s (i.e., the user-groups in JIRA speak), and the user emails.
-    """
+    Build ExternalAccess from the holder map.
 
+    Holder types handled:
+        - "anyone": Public project, anyone can access
+        - "applicationRole": All users with a Jira license can access (treated as public)
+        - "user": Specific users with access
+        - "projectRole": Project roles containing users and/or groups
+        - "group": Groups directly assigned in the permission scheme
+    """
+    # Public access - anyone can view
     if "anyone" in holder_map:
         return ExternalAccess(
             external_user_emails=set(), external_user_group_ids=set(), is_public=True
         )
 
+    # applicationRole means all users with a Jira license can access - treat as public
+    if "applicationRole" in holder_map:
+        return ExternalAccess(
+            external_user_emails=set(), external_user_group_ids=set(), is_public=True
+        )
+
+    # Get emails from explicit user holders
     user_emails = (
         _get_user_emails(user_holders=holder_map["user"])
         if "user" in holder_map
         else []
     )
-    project_role_user_emails = (
-        _get_user_emails_from_project_roles(
-            jira_client=jira_client,
-            jira_project=jira_project,
-            project_role_holders=holder_map["projectRole"],
+
+    # Get emails and groups from project roles
+    project_role_user_emails: list[str] = []
+    project_role_groups: list[str] = []
+    if "projectRole" in holder_map:
+        project_role_user_emails, project_role_groups = (
+            _get_user_emails_and_groups_from_project_roles(
+                jira_client=jira_client,
+                jira_project=jira_project,
+                project_role_holders=holder_map["projectRole"],
+            )
         )
-        if "projectRole" in holder_map
-        else []
-    )
+
+    # Get groups directly assigned in permission scheme (common in Data Center)
+    # Format: {'type': 'group', 'parameter': 'group-name', 'expand': 'group'}
+    direct_groups: list[str] = []
+    if "group" in holder_map:
+        for group_holder in holder_map["group"]:
+            group_name = _get_role_id(group_holder)
+            if group_name:
+                direct_groups.append(group_name)
+            else:
+                logger.error(f"No parameter/value in group holder: {group_holder}")
 
     external_user_emails = set(user_emails + project_role_user_emails)
+    external_user_group_ids = set(project_role_groups + direct_groups)
 
     return ExternalAccess(
         external_user_emails=external_user_emails,
-        external_user_group_ids=set(),
+        external_user_group_ids=external_user_group_ids,
         is_public=False,
     )
 
@@ -197,9 +258,11 @@ def get_project_permissions(
     )
 
     if not hasattr(project_permissions, "permissions"):
+        logger.error(f"Project {jira_project} has no permissions attribute")
         return None
 
     if not isinstance(project_permissions.permissions, list):
+        logger.error(f"Project {jira_project} permissions is not a list")
         return None
 
     holder_map = _build_holder_map(permissions=project_permissions.permissions)

backend/ee/onyx/external_permissions/sharepoint/permission_utils.py

@@ -15,6 +15,7 @@ from ee.onyx.db.external_perm import ExternalUserGroup
 from onyx.access.models import ExternalAccess
 from onyx.access.utils import build_ext_group_name_for_onyx
 from onyx.configs.constants import DocumentSource
+from onyx.connectors.sharepoint.connector import SHARED_DOCUMENTS_MAP_REVERSE
 from onyx.connectors.sharepoint.connector import sleep_and_retry
 from onyx.utils.logger import setup_logger
 
@@ -511,8 +512,8 @@ def get_external_access_from_sharepoint(
                 f"Failed to get SharePoint list item ID for item {drive_item.id}"
             )
 
-        if drive_name == "Shared Documents":
-            drive_name = "Documents"
+        if drive_name in SHARED_DOCUMENTS_MAP_REVERSE:
+            drive_name = SHARED_DOCUMENTS_MAP_REVERSE[drive_name]
 
         item = client_context.web.lists.get_by_title(drive_name).items.get_by_id(
             item_id

backend/ee/onyx/external_permissions/sync_params.py

@@ -11,6 +11,7 @@ from ee.onyx.configs.app_configs import GITHUB_PERMISSION_DOC_SYNC_FREQUENCY
 from ee.onyx.configs.app_configs import GITHUB_PERMISSION_GROUP_SYNC_FREQUENCY
 from ee.onyx.configs.app_configs import GOOGLE_DRIVE_PERMISSION_GROUP_SYNC_FREQUENCY
 from ee.onyx.configs.app_configs import JIRA_PERMISSION_DOC_SYNC_FREQUENCY
+from ee.onyx.configs.app_configs import JIRA_PERMISSION_GROUP_SYNC_FREQUENCY
 from ee.onyx.configs.app_configs import SHAREPOINT_PERMISSION_DOC_SYNC_FREQUENCY
 from ee.onyx.configs.app_configs import SHAREPOINT_PERMISSION_GROUP_SYNC_FREQUENCY
 from ee.onyx.configs.app_configs import SLACK_PERMISSION_DOC_SYNC_FREQUENCY
@@ -23,6 +24,7 @@ from ee.onyx.external_permissions.gmail.doc_sync import gmail_doc_sync
 from ee.onyx.external_permissions.google_drive.doc_sync import gdrive_doc_sync
 from ee.onyx.external_permissions.google_drive.group_sync import gdrive_group_sync
 from ee.onyx.external_permissions.jira.doc_sync import jira_doc_sync
+from ee.onyx.external_permissions.jira.group_sync import jira_group_sync
 from ee.onyx.external_permissions.perm_sync_types import CensoringFuncType
 from ee.onyx.external_permissions.perm_sync_types import DocSyncFuncType
 from ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsFunction
@@ -110,6 +112,11 @@ _SOURCE_TO_SYNC_CONFIG: dict[DocumentSource, SyncConfig] = {
             doc_sync_func=jira_doc_sync,
             initial_index_should_sync=True,
         ),
+        group_sync_config=GroupSyncConfig(
+            group_sync_frequency=JIRA_PERMISSION_GROUP_SYNC_FREQUENCY,
+            group_sync_func=jira_group_sync,
+            group_sync_is_cc_pair_agnostic=True,
+        ),
     ),
     # Groups are not needed for Slack.
     # All channel access is done at the individual user level.

backend/ee/onyx/main.py

@@ -14,6 +14,7 @@ from ee.onyx.server.enterprise_settings.api import (
     basic_router as enterprise_settings_router,
 )
 from ee.onyx.server.evals.api import router as evals_router
+from ee.onyx.server.license.api import router as license_router
 from ee.onyx.server.manage.standard_answer import router as standard_answer_router
 from ee.onyx.server.middleware.tenant_tracking import (
     add_api_server_tenant_id_middleware,
@@ -139,6 +140,8 @@ def get_application() -> FastAPI:
     )
     include_router_with_global_prefix_prepended(application, enterprise_settings_router)
     include_router_with_global_prefix_prepended(application, usage_export_router)
+    # License management
+    include_router_with_global_prefix_prepended(application, license_router)
 
     if MULTI_TENANT:
         # Tenant management

backend/ee/onyx/server/license/api.py

@@ -0,0 +1,246 @@
+"""License API endpoints."""
+
+import requests
+from fastapi import APIRouter
+from fastapi import Depends
+from fastapi import File
+from fastapi import HTTPException
+from fastapi import UploadFile
+from sqlalchemy.orm import Session
+
+from ee.onyx.auth.users import current_admin_user
+from ee.onyx.db.license import delete_license as db_delete_license
+from ee.onyx.db.license import get_license_metadata
+from ee.onyx.db.license import invalidate_license_cache
+from ee.onyx.db.license import refresh_license_cache
+from ee.onyx.db.license import update_license_cache
+from ee.onyx.db.license import upsert_license
+from ee.onyx.server.license.models import LicenseResponse
+from ee.onyx.server.license.models import LicenseSource
+from ee.onyx.server.license.models import LicenseStatusResponse
+from ee.onyx.server.license.models import LicenseUploadResponse
+from ee.onyx.server.license.models import SeatUsageResponse
+from ee.onyx.server.tenants.access import generate_data_plane_token
+from ee.onyx.utils.license import verify_license_signature
+from onyx.auth.users import User
+from onyx.configs.app_configs import CONTROL_PLANE_API_BASE_URL
+from onyx.db.engine.sql_engine import get_session
+from onyx.utils.logger import setup_logger
+from shared_configs.contextvars import get_current_tenant_id
+
+logger = setup_logger()
+
+router = APIRouter(prefix="/license")
+
+
+@router.get("")
+async def get_license_status(
+    _: User = Depends(current_admin_user),
+    db_session: Session = Depends(get_session),
+) -> LicenseStatusResponse:
+    """Get current license status and seat usage."""
+    metadata = get_license_metadata(db_session)
+
+    if not metadata:
+        return LicenseStatusResponse(has_license=False)
+
+    return LicenseStatusResponse(
+        has_license=True,
+        seats=metadata.seats,
+        used_seats=metadata.used_seats,
+        plan_type=metadata.plan_type,
+        issued_at=metadata.issued_at,
+        expires_at=metadata.expires_at,
+        grace_period_end=metadata.grace_period_end,
+        status=metadata.status,
+        source=metadata.source,
+    )
+
+
+@router.get("/seats")
+async def get_seat_usage(
+    _: User = Depends(current_admin_user),
+    db_session: Session = Depends(get_session),
+) -> SeatUsageResponse:
+    """Get detailed seat usage information."""
+    metadata = get_license_metadata(db_session)
+
+    if not metadata:
+        return SeatUsageResponse(
+            total_seats=0,
+            used_seats=0,
+            available_seats=0,
+        )
+
+    return SeatUsageResponse(
+        total_seats=metadata.seats,
+        used_seats=metadata.used_seats,
+        available_seats=max(0, metadata.seats - metadata.used_seats),
+    )
+
+
+@router.post("/fetch")
+async def fetch_license(
+    _: User = Depends(current_admin_user),
+    db_session: Session = Depends(get_session),
+) -> LicenseResponse:
+    """
+    Fetch license from control plane.
+    Used after Stripe checkout completion to retrieve the new license.
+    """
+    tenant_id = get_current_tenant_id()
+
+    try:
+        token = generate_data_plane_token()
+    except ValueError as e:
+        logger.error(f"Failed to generate data plane token: {e}")
+        raise HTTPException(
+            status_code=500, detail="Authentication configuration error"
+        )
+
+    try:
+        headers = {
+            "Authorization": f"Bearer {token}",
+            "Content-Type": "application/json",
+        }
+        url = f"{CONTROL_PLANE_API_BASE_URL}/license/{tenant_id}"
+        response = requests.get(url, headers=headers, timeout=10)
+        response.raise_for_status()
+
+        data = response.json()
+        if not isinstance(data, dict) or "license" not in data:
+            raise HTTPException(
+                status_code=502, detail="Invalid response from control plane"
+            )
+
+        license_data = data["license"]
+        if not license_data:
+            raise HTTPException(status_code=404, detail="No license found")
+
+        # Verify signature before persisting
+        payload = verify_license_signature(license_data)
+
+        # Verify the fetched license is for this tenant
+        if payload.tenant_id != tenant_id:
+            logger.error(
+                f"License tenant mismatch: expected {tenant_id}, got {payload.tenant_id}"
+            )
+            raise HTTPException(
+                status_code=400,
+                detail="License tenant ID mismatch - control plane returned wrong license",
+            )
+
+        # Persist to DB and update cache atomically
+        upsert_license(db_session, license_data)
+        try:
+            update_license_cache(payload, source=LicenseSource.AUTO_FETCH)
+        except Exception as cache_error:
+            # Log but don't fail - DB is source of truth, cache will refresh on next read
+            logger.warning(f"Failed to update license cache: {cache_error}")
+
+        return LicenseResponse(success=True, license=payload)
+
+    except requests.HTTPError as e:
+        status_code = e.response.status_code if e.response is not None else 502
+        logger.error(f"Control plane returned error: {status_code}")
+        raise HTTPException(
+            status_code=status_code,
+            detail="Failed to fetch license from control plane",
+        )
+    except ValueError as e:
+        logger.error(f"License verification failed: {type(e).__name__}")
+        raise HTTPException(status_code=400, detail=str(e))
+    except requests.RequestException:
+        logger.exception("Failed to fetch license from control plane")
+        raise HTTPException(
+            status_code=502, detail="Failed to connect to control plane"
+        )
+
+
+@router.post("/upload")
+async def upload_license(
+    license_file: UploadFile = File(...),
+    _: User = Depends(current_admin_user),
+    db_session: Session = Depends(get_session),
+) -> LicenseUploadResponse:
+    """
+    Upload a license file manually.
+    Used for air-gapped deployments where control plane is not accessible.
+    """
+    try:
+        content = await license_file.read()
+        license_data = content.decode("utf-8").strip()
+    except UnicodeDecodeError:
+        raise HTTPException(status_code=400, detail="Invalid license file format")
+
+    try:
+        payload = verify_license_signature(license_data)
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail=str(e))
+
+    tenant_id = get_current_tenant_id()
+    if payload.tenant_id != tenant_id:
+        raise HTTPException(
+            status_code=400,
+            detail=f"License tenant ID mismatch. Expected {tenant_id}, got {payload.tenant_id}",
+        )
+
+    # Persist to DB and update cache
+    upsert_license(db_session, license_data)
+    try:
+        update_license_cache(payload, source=LicenseSource.MANUAL_UPLOAD)
+    except Exception as cache_error:
+        # Log but don't fail - DB is source of truth, cache will refresh on next read
+        logger.warning(f"Failed to update license cache: {cache_error}")
+
+    return LicenseUploadResponse(
+        success=True,
+        message=f"License uploaded successfully. {payload.seats} seats, expires {payload.expires_at.date()}",
+    )
+
+
+@router.post("/refresh")
+async def refresh_license_cache_endpoint(
+    _: User = Depends(current_admin_user),
+    db_session: Session = Depends(get_session),
+) -> LicenseStatusResponse:
+    """
+    Force refresh the license cache from the database.
+    Useful after manual database changes or to verify license validity.
+    """
+    metadata = refresh_license_cache(db_session)
+
+    if not metadata:
+        return LicenseStatusResponse(has_license=False)
+
+    return LicenseStatusResponse(
+        has_license=True,
+        seats=metadata.seats,
+        used_seats=metadata.used_seats,
+        plan_type=metadata.plan_type,
+        issued_at=metadata.issued_at,
+        expires_at=metadata.expires_at,
+        grace_period_end=metadata.grace_period_end,
+        status=metadata.status,
+        source=metadata.source,
+    )
+
+
+@router.delete("")
+async def delete_license(
+    _: User = Depends(current_admin_user),
+    db_session: Session = Depends(get_session),
+) -> dict[str, bool]:
+    """
+    Delete the current license.
+    Admin only - removes license and invalidates cache.
+    """
+    # Invalidate cache first - if DB delete fails, stale cache is worse than no cache
+    try:
+        invalidate_license_cache()
+    except Exception as cache_error:
+        logger.warning(f"Failed to invalidate license cache: {cache_error}")
+
+    deleted = db_delete_license(db_session)
+
+    return {"deleted": deleted}

backend/ee/onyx/server/license/models.py

@@ -0,0 +1,92 @@
+from datetime import datetime
+from enum import Enum
+
+from pydantic import BaseModel
+
+from onyx.server.settings.models import ApplicationStatus
+
+
+class PlanType(str, Enum):
+    MONTHLY = "monthly"
+    ANNUAL = "annual"
+
+
+class LicenseSource(str, Enum):
+    AUTO_FETCH = "auto_fetch"
+    MANUAL_UPLOAD = "manual_upload"
+
+
+class LicensePayload(BaseModel):
+    """The payload portion of a signed license."""
+
+    version: str
+    tenant_id: str
+    organization_name: str | None = None
+    issued_at: datetime
+    expires_at: datetime
+    seats: int
+    plan_type: PlanType
+    billing_cycle: str | None = None
+    grace_period_days: int = 30
+    stripe_subscription_id: str | None = None
+    stripe_customer_id: str | None = None
+
+
+class LicenseData(BaseModel):
+    """Full signed license structure."""
+
+    payload: LicensePayload
+    signature: str
+
+
+class LicenseMetadata(BaseModel):
+    """Cached license metadata stored in Redis."""
+
+    tenant_id: str
+    organization_name: str | None = None
+    seats: int
+    used_seats: int
+    plan_type: PlanType
+    issued_at: datetime
+    expires_at: datetime
+    grace_period_end: datetime | None = None
+    status: ApplicationStatus
+    source: LicenseSource | None = None
+    stripe_subscription_id: str | None = None
+
+
+class LicenseStatusResponse(BaseModel):
+    """Response for license status API."""
+
+    has_license: bool
+    seats: int = 0
+    used_seats: int = 0
+    plan_type: PlanType | None = None
+    issued_at: datetime | None = None
+    expires_at: datetime | None = None
+    grace_period_end: datetime | None = None
+    status: ApplicationStatus | None = None
+    source: LicenseSource | None = None
+
+
+class LicenseResponse(BaseModel):
+    """Response after license fetch/upload."""
+
+    success: bool
+    message: str | None = None
+    license: LicensePayload | None = None
+
+
+class LicenseUploadResponse(BaseModel):
+    """Response after license upload."""
+
+    success: bool
+    message: str | None = None
+
+
+class SeatUsageResponse(BaseModel):
+    """Response for seat usage API."""
+
+    total_seats: int
+    used_seats: int
+    available_seats: int

backend/ee/onyx/server/query_and_chat/chat_backend.py

@@ -8,12 +8,10 @@ from ee.onyx.server.query_and_chat.models import (
     BasicCreateChatMessageWithHistoryRequest,
 )
 from onyx.auth.users import current_user
-from onyx.chat.chat_utils import combine_message_thread
-from onyx.chat.chat_utils import create_chat_chain
+from onyx.chat.chat_utils import create_chat_history_chain
 from onyx.chat.models import ChatBasicResponse
 from onyx.chat.process_message import gather_stream
 from onyx.chat.process_message import stream_chat_message_objects
-from onyx.configs.chat_configs import CHAT_TARGET_CHUNK_PERCENTAGE
 from onyx.configs.constants import MessageType
 from onyx.context.search.models import OptionalSearchSetting
 from onyx.context.search.models import RetrievalDetails
@@ -22,9 +20,8 @@ from onyx.db.chat import create_new_chat_message
 from onyx.db.chat import get_or_create_root_message
 from onyx.db.engine.sql_engine import get_session
 from onyx.db.models import User
-from onyx.llm.factory import get_llms_for_persona
+from onyx.llm.factory import get_llm_for_persona
 from onyx.natural_language_processing.utils import get_tokenizer
-from onyx.secondary_llm_flows.query_expansion import thread_based_query_rephrase
 from onyx.server.query_and_chat.models import CreateChatMessageRequest
 from onyx.utils.logger import setup_logger
 
@@ -69,9 +66,9 @@ def handle_simplified_chat_message(
         chat_session_id = chat_message_req.chat_session_id
 
     try:
-        parent_message, _ = create_chat_chain(
+        parent_message = create_chat_history_chain(
             chat_session_id=chat_session_id, db_session=db_session
-        )
+        )[-1]
     except Exception:
         parent_message = get_or_create_root_message(
             chat_session_id=chat_session_id, db_session=db_session
@@ -103,7 +100,6 @@ def handle_simplified_chat_message(
         chunks_below=0,
         full_doc=chat_message_req.full_doc,
         structured_response_format=chat_message_req.structured_response_format,
-        use_agentic_search=chat_message_req.use_agentic_search,
     )
 
     packets = stream_chat_message_objects(
@@ -161,15 +157,13 @@ def handle_send_message_simple_with_history(
         persona_id=req.persona_id,
     )
 
-    llm, _ = get_llms_for_persona(persona=chat_session.persona, user=user)
+    llm = get_llm_for_persona(persona=chat_session.persona, user=user)
 
     llm_tokenizer = get_tokenizer(
         model_name=llm.config.model_name,
         provider_type=llm.config.model_provider,
     )
 
-    max_history_tokens = int(llm.config.max_input_tokens * CHAT_TARGET_CHUNK_PERCENTAGE)
-
     # Every chat Session begins with an empty root message
     root_message = get_or_create_root_message(
         chat_session_id=chat_session.id, db_session=db_session
@@ -188,17 +182,6 @@ def handle_send_message_simple_with_history(
         )
     db_session.commit()
 
-    history_str = combine_message_thread(
-        messages=msg_history,
-        max_tokens=max_history_tokens,
-        llm_tokenizer=llm_tokenizer,
-    )
-
-    rephrased_query = req.query_override or thread_based_query_rephrase(
-        user_query=query,
-        history_str=history_str,
-    )
-
     if req.retrieval_options is None and req.search_doc_ids is None:
         retrieval_options: RetrievalDetails | None = RetrievalDetails(
             run_search=OptionalSearchSetting.ALWAYS,
@@ -216,12 +199,11 @@ def handle_send_message_simple_with_history(
         retrieval_options=retrieval_options,
         # Simple API does not support reranking, hide complexity from user
         rerank_settings=None,
-        query_override=rephrased_query,
+        query_override=None,
         chunks_above=0,
         chunks_below=0,
         full_doc=req.full_doc,
         structured_response_format=req.structured_response_format,
-        use_agentic_search=req.use_agentic_search,
     )
 
     packets = stream_chat_message_objects(

backend/ee/onyx/server/query_and_chat/models.py

@@ -8,10 +8,29 @@ from pydantic import model_validator
 
 from onyx.chat.models import ThreadMessage
 from onyx.configs.constants import DocumentSource
+from onyx.context.search.models import BaseFilters
+from onyx.context.search.models import BasicChunkRequest
 from onyx.context.search.models import ChunkContext
+from onyx.context.search.models import InferenceChunk
 from onyx.context.search.models import RetrievalDetails
 from onyx.server.manage.models import StandardAnswer
-from onyx.server.query_and_chat.streaming_models import SubQuestionIdentifier
+
+
+class StandardAnswerRequest(BaseModel):
+    message: str
+    slack_bot_categories: list[str]
+
+
+class StandardAnswerResponse(BaseModel):
+    standard_answers: list[StandardAnswer] = Field(default_factory=list)
+
+
+class DocumentSearchRequest(BasicChunkRequest):
+    user_selected_filters: BaseFilters | None = None
+
+
+class DocumentSearchResponse(BaseModel):
+    top_documents: list[InferenceChunk]
 
 
 class BasicCreateChatMessageRequest(ChunkContext):
@@ -35,9 +54,6 @@ class BasicCreateChatMessageRequest(ChunkContext):
     # https://platform.openai.com/docs/guides/structured-outputs/introduction
     structured_response_format: dict | None = None
 
-    # If True, uses agentic search instead of basic search
-    use_agentic_search: bool = False
-
     @model_validator(mode="after")
     def validate_chat_session_or_persona(self) -> "BasicCreateChatMessageRequest":
         if self.chat_session_id is None and self.persona_id is None:
@@ -57,8 +73,6 @@ class BasicCreateChatMessageWithHistoryRequest(ChunkContext):
     # only works if using an OpenAI model. See the following for more details:
     # https://platform.openai.com/docs/guides/structured-outputs/introduction
     structured_response_format: dict | None = None
-    # If True, uses agentic search instead of basic search
-    use_agentic_search: bool = False
 
 
 class SimpleDoc(BaseModel):
@@ -71,17 +85,17 @@ class SimpleDoc(BaseModel):
     metadata: dict | None
 
 
-class AgentSubQuestion(SubQuestionIdentifier):
+class AgentSubQuestion(BaseModel):
     sub_question: str
     document_ids: list[str]
 
 
-class AgentAnswer(SubQuestionIdentifier):
+class AgentAnswer(BaseModel):
     answer: str
     answer_type: Literal["agent_sub_answer", "agent_level_answer"]
 
 
-class AgentSubQuery(SubQuestionIdentifier):
+class AgentSubQuery(BaseModel):
     sub_query: str
     query_id: int
 
@@ -127,12 +141,3 @@ class AgentSubQuery(SubQuestionIdentifier):
             sorted(level_question_dict.items(), key=lambda x: (x is None, x))
         )
         return sorted_dict
-
-
-class StandardAnswerRequest(BaseModel):
-    message: str
-    slack_bot_categories: list[str]
-
-
-class StandardAnswerResponse(BaseModel):
-    standard_answers: list[StandardAnswer] = Field(default_factory=list)

backend/ee/onyx/server/query_history/api.py

@@ -24,7 +24,7 @@ from onyx.auth.users import current_admin_user
 from onyx.auth.users import get_display_email
 from onyx.background.celery.versioned_apps.client import app as client_app
 from onyx.background.task_utils import construct_query_history_report_name
-from onyx.chat.chat_utils import create_chat_chain
+from onyx.chat.chat_utils import create_chat_history_chain
 from onyx.configs.app_configs import ONYX_QUERY_HISTORY_TYPE
 from onyx.configs.constants import FileOrigin
 from onyx.configs.constants import FileType
@@ -123,10 +123,9 @@ def snapshot_from_chat_session(
 ) -> ChatSessionSnapshot | None:
     try:
         # Older chats may not have the right structure
-        last_message, messages = create_chat_chain(
+        messages = create_chat_history_chain(
             chat_session_id=chat_session.id, db_session=db_session
         )
-        messages.append(last_message)
     except RuntimeError:
         return None
 

backend/ee/onyx/server/tenants/provisioning.py

@@ -37,17 +37,14 @@ from onyx.db.models import AvailableTenant
 from onyx.db.models import IndexModelStatus
 from onyx.db.models import SearchSettings
 from onyx.db.models import UserTenantMapping
-from onyx.llm.llm_provider_options import ANTHROPIC_PROVIDER_NAME
-from onyx.llm.llm_provider_options import ANTHROPIC_VISIBLE_MODEL_NAMES
+from onyx.llm.constants import LlmProviderNames
 from onyx.llm.llm_provider_options import get_anthropic_model_names
-from onyx.llm.llm_provider_options import OPEN_AI_MODEL_NAMES
-from onyx.llm.llm_provider_options import OPEN_AI_VISIBLE_MODEL_NAMES
-from onyx.llm.llm_provider_options import OPENAI_PROVIDER_NAME
+from onyx.llm.llm_provider_options import get_openai_model_names
 from onyx.server.manage.embedding.models import CloudEmbeddingProviderCreationRequest
 from onyx.server.manage.llm.models import LLMProviderUpsertRequest
 from onyx.server.manage.llm.models import ModelConfigurationUpsertRequest
 from onyx.setup import setup_onyx
-from onyx.utils.telemetry import create_milestone_and_report
+from onyx.utils.telemetry import mt_cloud_telemetry
 from shared_configs.configs import MULTI_TENANT
 from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA
 from shared_configs.configs import TENANT_ID_PREFIX
@@ -268,14 +265,13 @@ def configure_default_api_keys(db_session: Session) -> None:
     if ANTHROPIC_DEFAULT_API_KEY:
         anthropic_provider = LLMProviderUpsertRequest(
             name="Anthropic",
-            provider=ANTHROPIC_PROVIDER_NAME,
+            provider=LlmProviderNames.ANTHROPIC,
             api_key=ANTHROPIC_DEFAULT_API_KEY,
             default_model_name="claude-3-7-sonnet-20250219",
-            fast_default_model_name="claude-3-5-sonnet-20241022",
             model_configurations=[
                 ModelConfigurationUpsertRequest(
                     name=name,
-                    is_visible=name in ANTHROPIC_VISIBLE_MODEL_NAMES,
+                    is_visible=False,
                     max_input_tokens=None,
                 )
                 for name in get_anthropic_model_names()
@@ -295,17 +291,16 @@ def configure_default_api_keys(db_session: Session) -> None:
     if OPENAI_DEFAULT_API_KEY:
         openai_provider = LLMProviderUpsertRequest(
             name="OpenAI",
-            provider=OPENAI_PROVIDER_NAME,
+            provider=LlmProviderNames.OPENAI,
             api_key=OPENAI_DEFAULT_API_KEY,
             default_model_name="gpt-4o",
-            fast_default_model_name="gpt-4o-mini",
             model_configurations=[
                 ModelConfigurationUpsertRequest(
                     name=model_name,
-                    is_visible=model_name in OPEN_AI_VISIBLE_MODEL_NAMES,
+                    is_visible=False,
                     max_input_tokens=None,
                 )
-                for model_name in OPEN_AI_MODEL_NAMES
+                for model_name in get_openai_model_names()
             ],
             api_key_changed=True,
         )
@@ -564,17 +559,11 @@ async def assign_tenant_to_user(
     try:
         add_users_to_tenant([email], tenant_id)
 
-        # Create milestone record in the same transaction context as the tenant assignment
-        with get_session_with_tenant(tenant_id=tenant_id) as db_session:
-            create_milestone_and_report(
-                user=None,
-                distinct_id=tenant_id,
-                event_type=MilestoneRecordType.TENANT_CREATED,
-                properties={
-                    "email": email,
-                },
-                db_session=db_session,
-            )
+        mt_cloud_telemetry(
+            tenant_id=tenant_id,
+            distinct_id=email,
+            event=MilestoneRecordType.TENANT_CREATED,
+        )
     except Exception:
         logger.exception(f"Failed to assign tenant {tenant_id} to user {email}")
         raise Exception("Failed to assign tenant to user")

backend/ee/onyx/server/tenants/user_mapping.py

@@ -249,6 +249,17 @@ def accept_user_invite(email: str, tenant_id: str) -> None:
             )
             raise
 
+    # Remove from invited users list since they've accepted
+    token = CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id)
+    try:
+        invited_users = get_invited_users()
+        if email in invited_users:
+            invited_users.remove(email)
+            write_invited_users(invited_users)
+            logger.info(f"Removed {email} from invited users list after acceptance")
+    finally:
+        CURRENT_TENANT_ID_CONTEXTVAR.reset(token)
+
 
 def deny_user_invite(email: str, tenant_id: str) -> None:
     """

backend/ee/onyx/utils/license.py

@@ -0,0 +1,126 @@
+"""RSA-4096 license signature verification utilities."""
+
+import base64
+import json
+import os
+from datetime import datetime
+from datetime import timezone
+
+from cryptography.exceptions import InvalidSignature
+from cryptography.hazmat.primitives import hashes
+from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric import padding
+from cryptography.hazmat.primitives.asymmetric.rsa import RSAPublicKey
+
+from ee.onyx.server.license.models import LicenseData
+from ee.onyx.server.license.models import LicensePayload
+from onyx.server.settings.models import ApplicationStatus
+from onyx.utils.logger import setup_logger
+
+logger = setup_logger()
+
+
+# RSA-4096 Public Key for license verification
+# Load from environment variable - key is generated on the control plane
+# In production, inject via Kubernetes secrets or secrets manager
+LICENSE_PUBLIC_KEY_PEM = os.environ.get("LICENSE_PUBLIC_KEY_PEM", "")
+
+
+def _get_public_key() -> RSAPublicKey:
+    """Load the public key from environment variable."""
+    if not LICENSE_PUBLIC_KEY_PEM:
+        raise ValueError(
+            "LICENSE_PUBLIC_KEY_PEM environment variable not set. "
+            "License verification requires the control plane public key."
+        )
+    key = serialization.load_pem_public_key(LICENSE_PUBLIC_KEY_PEM.encode())
+    if not isinstance(key, RSAPublicKey):
+        raise ValueError("Expected RSA public key")
+    return key
+
+
+def verify_license_signature(license_data: str) -> LicensePayload:
+    """
+    Verify RSA-4096 signature and return payload if valid.
+
+    Args:
+        license_data: Base64-encoded JSON containing payload and signature
+
+    Returns:
+        LicensePayload if signature is valid
+
+    Raises:
+        ValueError: If license data is invalid or signature verification fails
+    """
+    try:
+        # Decode the license data
+        decoded = json.loads(base64.b64decode(license_data))
+        license_obj = LicenseData(**decoded)
+
+        payload_json = json.dumps(
+            license_obj.payload.model_dump(mode="json"), sort_keys=True
+        )
+        signature_bytes = base64.b64decode(license_obj.signature)
+
+        # Verify signature using PSS padding (modern standard)
+        public_key = _get_public_key()
+        public_key.verify(
+            signature_bytes,
+            payload_json.encode(),
+            padding.PSS(
+                mgf=padding.MGF1(hashes.SHA256()),
+                salt_length=padding.PSS.MAX_LENGTH,
+            ),
+            hashes.SHA256(),
+        )
+
+        return license_obj.payload
+
+    except InvalidSignature:
+        logger.error("License signature verification failed")
+        raise ValueError("Invalid license signature")
+    except json.JSONDecodeError:
+        logger.error("Failed to decode license JSON")
+        raise ValueError("Invalid license format: not valid JSON")
+    except (ValueError, KeyError, TypeError) as e:
+        logger.error(f"License data validation error: {type(e).__name__}")
+        raise ValueError(f"Invalid license format: {type(e).__name__}")
+    except Exception:
+        logger.exception("Unexpected error during license verification")
+        raise ValueError("License verification failed: unexpected error")
+
+
+def get_license_status(
+    payload: LicensePayload,
+    grace_period_end: datetime | None = None,
+) -> ApplicationStatus:
+    """
+    Determine current license status based on expiry.
+
+    Args:
+        payload: The verified license payload
+        grace_period_end: Optional grace period end datetime
+
+    Returns:
+        ApplicationStatus indicating current license state
+    """
+    now = datetime.now(timezone.utc)
+
+    # Check if grace period has expired
+    if grace_period_end and now > grace_period_end:
+        return ApplicationStatus.GATED_ACCESS
+
+    # Check if license has expired
+    if now > payload.expires_at:
+        if grace_period_end and now <= grace_period_end:
+            return ApplicationStatus.GRACE_PERIOD
+        return ApplicationStatus.GATED_ACCESS
+
+    # License is valid
+    return ApplicationStatus.ACTIVE
+
+
+def is_license_valid(payload: LicensePayload) -> bool:
+    """Check if a license is currently valid (not expired)."""
+    now = datetime.now(timezone.utc)
+    return now <= payload.expires_at

backend/model_server/constants.py

@@ -1,5 +1,4 @@
 MODEL_WARM_UP_STRING = "hi " * 512
-INFORMATION_CONTENT_MODEL_WARM_UP_STRING = "hi " * 16
 
 
 class GPUStatus:

backend/model_server/custom_models.py

@@ -1,562 +0,0 @@
-from typing import cast
-from typing import Optional
-from typing import TYPE_CHECKING
-
-import numpy as np
-import torch
-import torch.nn.functional as F
-from fastapi import APIRouter
-from huggingface_hub import snapshot_download  # type: ignore
-
-from model_server.constants import INFORMATION_CONTENT_MODEL_WARM_UP_STRING
-from model_server.constants import MODEL_WARM_UP_STRING
-from model_server.onyx_torch_model import ConnectorClassifier
-from model_server.onyx_torch_model import HybridClassifier
-from model_server.utils import simple_log_function_time
-from onyx.utils.logger import setup_logger
-from shared_configs.configs import CONNECTOR_CLASSIFIER_MODEL_REPO
-from shared_configs.configs import CONNECTOR_CLASSIFIER_MODEL_TAG
-from shared_configs.configs import (
-    INDEXING_INFORMATION_CONTENT_CLASSIFICATION_CUTOFF_LENGTH,
-)
-from shared_configs.configs import INDEXING_INFORMATION_CONTENT_CLASSIFICATION_MAX
-from shared_configs.configs import INDEXING_INFORMATION_CONTENT_CLASSIFICATION_MIN
-from shared_configs.configs import (
-    INDEXING_INFORMATION_CONTENT_CLASSIFICATION_TEMPERATURE,
-)
-from shared_configs.configs import INDEXING_ONLY
-from shared_configs.configs import INFORMATION_CONTENT_MODEL_TAG
-from shared_configs.configs import INFORMATION_CONTENT_MODEL_VERSION
-from shared_configs.configs import INTENT_MODEL_TAG
-from shared_configs.configs import INTENT_MODEL_VERSION
-from shared_configs.model_server_models import ConnectorClassificationRequest
-from shared_configs.model_server_models import ConnectorClassificationResponse
-from shared_configs.model_server_models import ContentClassificationPrediction
-from shared_configs.model_server_models import IntentRequest
-from shared_configs.model_server_models import IntentResponse
-
-if TYPE_CHECKING:
-    from setfit import SetFitModel  # type: ignore
-    from transformers import PreTrainedTokenizer, BatchEncoding  # type: ignore
-
-
-logger = setup_logger()
-
-router = APIRouter(prefix="/custom")
-
-_CONNECTOR_CLASSIFIER_TOKENIZER: Optional["PreTrainedTokenizer"] = None
-_CONNECTOR_CLASSIFIER_MODEL: ConnectorClassifier | None = None
-
-_INTENT_TOKENIZER: Optional["PreTrainedTokenizer"] = None
-_INTENT_MODEL: HybridClassifier | None = None
-
-_INFORMATION_CONTENT_MODEL: Optional["SetFitModel"] = None
-
-_INFORMATION_CONTENT_MODEL_PROMPT_PREFIX: str = ""  # spec to model version!
-
-
-def get_connector_classifier_tokenizer() -> "PreTrainedTokenizer":
-    global _CONNECTOR_CLASSIFIER_TOKENIZER
-    from transformers import AutoTokenizer, PreTrainedTokenizer
-
-    if _CONNECTOR_CLASSIFIER_TOKENIZER is None:
-        # The tokenizer details are not uploaded to the HF hub since it's just the
-        # unmodified distilbert tokenizer.
-        _CONNECTOR_CLASSIFIER_TOKENIZER = cast(
-            PreTrainedTokenizer,
-            AutoTokenizer.from_pretrained("distilbert-base-uncased"),
-        )
-    return _CONNECTOR_CLASSIFIER_TOKENIZER
-
-
-def get_local_connector_classifier(
-    model_name_or_path: str = CONNECTOR_CLASSIFIER_MODEL_REPO,
-    tag: str = CONNECTOR_CLASSIFIER_MODEL_TAG,
-) -> ConnectorClassifier:
-    global _CONNECTOR_CLASSIFIER_MODEL
-    if _CONNECTOR_CLASSIFIER_MODEL is None:
-        try:
-            # Calculate where the cache should be, then load from local if available
-            local_path = snapshot_download(
-                repo_id=model_name_or_path, revision=tag, local_files_only=True
-            )
-            _CONNECTOR_CLASSIFIER_MODEL = ConnectorClassifier.from_pretrained(
-                local_path
-            )
-        except Exception as e:
-            logger.warning(f"Failed to load model directly: {e}")
-            try:
-                # Attempt to download the model snapshot
-                logger.info(f"Downloading model snapshot for {model_name_or_path}")
-                local_path = snapshot_download(repo_id=model_name_or_path, revision=tag)
-                _CONNECTOR_CLASSIFIER_MODEL = ConnectorClassifier.from_pretrained(
-                    local_path
-                )
-            except Exception as e:
-                logger.error(
-                    f"Failed to load model even after attempted snapshot download: {e}"
-                )
-                raise
-    return _CONNECTOR_CLASSIFIER_MODEL
-
-
-def get_intent_model_tokenizer() -> "PreTrainedTokenizer":
-    from transformers import AutoTokenizer, PreTrainedTokenizer
-
-    global _INTENT_TOKENIZER
-    if _INTENT_TOKENIZER is None:
-        # The tokenizer details are not uploaded to the HF hub since it's just the
-        # unmodified distilbert tokenizer.
-        _INTENT_TOKENIZER = cast(
-            PreTrainedTokenizer,
-            AutoTokenizer.from_pretrained("distilbert-base-uncased"),
-        )
-    return _INTENT_TOKENIZER
-
-
-def get_local_intent_model(
-    model_name_or_path: str = INTENT_MODEL_VERSION,
-    tag: str | None = INTENT_MODEL_TAG,
-) -> HybridClassifier:
-    global _INTENT_MODEL
-    if _INTENT_MODEL is None:
-        try:
-            # Calculate where the cache should be, then load from local if available
-            logger.notice(f"Loading model from local cache: {model_name_or_path}")
-            local_path = snapshot_download(
-                repo_id=model_name_or_path, revision=tag, local_files_only=True
-            )
-            _INTENT_MODEL = HybridClassifier.from_pretrained(local_path)
-            logger.notice(f"Loaded model from local cache: {local_path}")
-        except Exception as e:
-            logger.warning(f"Failed to load model directly: {e}")
-            try:
-                # Attempt to download the model snapshot
-                logger.notice(f"Downloading model snapshot for {model_name_or_path}")
-                local_path = snapshot_download(
-                    repo_id=model_name_or_path, revision=tag, local_files_only=False
-                )
-                _INTENT_MODEL = HybridClassifier.from_pretrained(local_path)
-            except Exception as e:
-                logger.error(
-                    f"Failed to load model even after attempted snapshot download: {e}"
-                )
-                raise
-    return _INTENT_MODEL
-
-
-def get_local_information_content_model(
-    model_name_or_path: str = INFORMATION_CONTENT_MODEL_VERSION,
-    tag: str | None = INFORMATION_CONTENT_MODEL_TAG,
-) -> "SetFitModel":
-    from setfit import SetFitModel
-
-    global _INFORMATION_CONTENT_MODEL
-    if _INFORMATION_CONTENT_MODEL is None:
-        try:
-            # Calculate where the cache should be, then load from local if available
-            logger.notice(
-                f"Loading content information model from local cache: {model_name_or_path}"
-            )
-            local_path = snapshot_download(
-                repo_id=model_name_or_path, revision=tag, local_files_only=True
-            )
-            _INFORMATION_CONTENT_MODEL = SetFitModel.from_pretrained(local_path)
-            logger.notice(
-                f"Loaded content information model from local cache: {local_path}"
-            )
-        except Exception as e:
-            logger.warning(f"Failed to load content information model directly: {e}")
-            try:
-                # Attempt to download the model snapshot
-                logger.notice(
-                    f"Downloading content information model snapshot for {model_name_or_path}"
-                )
-                local_path = snapshot_download(
-                    repo_id=model_name_or_path, revision=tag, local_files_only=False
-                )
-                _INFORMATION_CONTENT_MODEL = SetFitModel.from_pretrained(local_path)
-            except Exception as e:
-                logger.error(
-                    f"Failed to load content information model even after attempted snapshot download: {e}"
-                )
-                raise
-
-    return _INFORMATION_CONTENT_MODEL
-
-
-def tokenize_connector_classification_query(
-    connectors: list[str],
-    query: str,
-    tokenizer: "PreTrainedTokenizer",
-    connector_token_end_id: int,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """
-    Tokenize the connectors & user query into one prompt for the forward pass of ConnectorClassifier models
-
-    The attention mask is just all 1s. The prompt is CLS + each connector name suffixed with the connector end
-    token and then the user query.
-    """
-
-    input_ids = torch.tensor([tokenizer.cls_token_id], dtype=torch.long)
-
-    for connector in connectors:
-        connector_token_ids = tokenizer(
-            connector,
-            add_special_tokens=False,
-            return_tensors="pt",
-        )
-
-        input_ids = torch.cat(
-            (
-                input_ids,
-                connector_token_ids["input_ids"].squeeze(dim=0),
-                torch.tensor([connector_token_end_id], dtype=torch.long),
-            ),
-            dim=-1,
-        )
-    query_token_ids = tokenizer(
-        query,
-        add_special_tokens=False,
-        return_tensors="pt",
-    )
-
-    input_ids = torch.cat(
-        (
-            input_ids,
-            query_token_ids["input_ids"].squeeze(dim=0),
-            torch.tensor([tokenizer.sep_token_id], dtype=torch.long),
-        ),
-        dim=-1,
-    )
-    attention_mask = torch.ones(input_ids.numel(), dtype=torch.long)
-
-    return input_ids.unsqueeze(0), attention_mask.unsqueeze(0)
-
-
-def warm_up_connector_classifier_model() -> None:
-    logger.info(
-        f"Warming up connector_classifier model {CONNECTOR_CLASSIFIER_MODEL_TAG}"
-    )
-    connector_classifier_tokenizer = get_connector_classifier_tokenizer()
-    connector_classifier = get_local_connector_classifier()
-
-    input_ids, attention_mask = tokenize_connector_classification_query(
-        ["GitHub"],
-        "onyx classifier query google doc",
-        connector_classifier_tokenizer,
-        connector_classifier.connector_end_token_id,
-    )
-    input_ids = input_ids.to(connector_classifier.device)
-    attention_mask = attention_mask.to(connector_classifier.device)
-
-    connector_classifier(input_ids, attention_mask)
-
-
-def warm_up_intent_model() -> None:
-    logger.notice(f"Warming up Intent Model: {INTENT_MODEL_VERSION}")
-    intent_tokenizer = get_intent_model_tokenizer()
-    tokens = intent_tokenizer(
-        MODEL_WARM_UP_STRING, return_tensors="pt", truncation=True, padding=True
-    )
-
-    intent_model = get_local_intent_model()
-    device = intent_model.device
-    intent_model(
-        query_ids=tokens["input_ids"].to(device),
-        query_mask=tokens["attention_mask"].to(device),
-    )
-
-
-def warm_up_information_content_model() -> None:
-    logger.notice("Warming up Content Model")  # TODO: add version if needed
-
-    information_content_model = get_local_information_content_model()
-    information_content_model(INFORMATION_CONTENT_MODEL_WARM_UP_STRING)
-
-
-@simple_log_function_time()
-def run_inference(tokens: "BatchEncoding") -> tuple[list[float], list[float]]:
-    intent_model = get_local_intent_model()
-    device = intent_model.device
-
-    outputs = intent_model(
-        query_ids=tokens["input_ids"].to(device),
-        query_mask=tokens["attention_mask"].to(device),
-    )
-
-    token_logits = outputs["token_logits"]
-    intent_logits = outputs["intent_logits"]
-
-    # Move tensors to CPU before applying softmax and converting to numpy
-    intent_probabilities = F.softmax(intent_logits.cpu(), dim=-1).numpy()[0]
-    token_probabilities = F.softmax(token_logits.cpu(), dim=-1).numpy()[0]
-
-    # Extract the probabilities for the positive class (index 1) for each token
-    token_positive_probs = token_probabilities[:, 1].tolist()
-
-    return intent_probabilities.tolist(), token_positive_probs
-
-
-@simple_log_function_time()
-def run_content_classification_inference(
-    text_inputs: list[str],
-) -> list[ContentClassificationPrediction]:
-    """
-    Assign a score to the segments in question. The model stored in get_local_information_content_model()
-    creates the 'model score' based on its training, and the scores are then converted to a 0.0-1.0 scale.
-    In the code outside of the model/inference model servers that score will be converted into the actual
-    boost factor.
-    """
-
-    def _prob_to_score(prob: float) -> float:
-        """
-        Conversion of base score to 0.0 - 1.0 score. Note that the min/max values depend on the model!
-        """
-        _MIN_BASE_SCORE = 0.25
-        _MAX_BASE_SCORE = 0.75
-        if prob < _MIN_BASE_SCORE:
-            raw_score = 0.0
-        elif prob < _MAX_BASE_SCORE:
-            raw_score = (prob - _MIN_BASE_SCORE) / (_MAX_BASE_SCORE - _MIN_BASE_SCORE)
-        else:
-            raw_score = 1.0
-        return (
-            INDEXING_INFORMATION_CONTENT_CLASSIFICATION_MIN
-            + (
-                INDEXING_INFORMATION_CONTENT_CLASSIFICATION_MAX
-                - INDEXING_INFORMATION_CONTENT_CLASSIFICATION_MIN
-            )
-            * raw_score
-        )
-
-    _BATCH_SIZE = 32
-    content_model = get_local_information_content_model()
-
-    # Process inputs in batches
-    all_output_classes: list[int] = []
-    all_base_output_probabilities: list[float] = []
-
-    for i in range(0, len(text_inputs), _BATCH_SIZE):
-        batch = text_inputs[i : i + _BATCH_SIZE]
-        batch_with_prefix = []
-        batch_indices = []
-
-        # Pre-allocate results for this batch
-        batch_output_classes: list[np.ndarray] = [np.array(1)] * len(batch)
-        batch_probabilities: list[np.ndarray] = [np.array(1.0)] * len(batch)
-
-        # Pre-process batch to handle long input exceptions
-        for j, text in enumerate(batch):
-            if len(text) == 0:
-                # if no input, treat as non-informative from the model's perspective
-                batch_output_classes[j] = np.array(0)
-                batch_probabilities[j] = np.array(0.0)
-                logger.warning("Input for Content Information Model is empty")
-
-            elif (
-                len(text.split())
-                <= INDEXING_INFORMATION_CONTENT_CLASSIFICATION_CUTOFF_LENGTH
-            ):
-                # if input is short, use the model
-                batch_with_prefix.append(
-                    _INFORMATION_CONTENT_MODEL_PROMPT_PREFIX + text
-                )
-                batch_indices.append(j)
-            else:
-                # if longer than cutoff, treat as informative (stay with default), but issue warning
-                logger.warning("Input for Content Information Model too long")
-
-        if batch_with_prefix:  # Only run model if we have valid inputs
-            # Get predictions for the batch
-            model_output_classes = content_model(batch_with_prefix)
-            model_output_probabilities = content_model.predict_proba(batch_with_prefix)
-
-            # Place results in the correct positions
-            for idx, batch_idx in enumerate(batch_indices):
-                batch_output_classes[batch_idx] = model_output_classes[idx].numpy()
-                batch_probabilities[batch_idx] = model_output_probabilities[idx][
-                    1
-                ].numpy()  # x[1] is prob of the positive class
-
-        all_output_classes.extend([int(x) for x in batch_output_classes])
-        all_base_output_probabilities.extend([float(x) for x in batch_probabilities])
-
-    logits = [
-        np.log(p / (1 - p)) if p != 0.0 and p != 1.0 else (100 if p == 1.0 else -100)
-        for p in all_base_output_probabilities
-    ]
-    scaled_logits = [
-        logit / INDEXING_INFORMATION_CONTENT_CLASSIFICATION_TEMPERATURE
-        for logit in logits
-    ]
-    output_probabilities_with_temp = [
-        np.exp(scaled_logit) / (1 + np.exp(scaled_logit))
-        for scaled_logit in scaled_logits
-    ]
-
-    prediction_scores = [
-        _prob_to_score(p_temp) for p_temp in output_probabilities_with_temp
-    ]
-
-    content_classification_predictions = [
-        ContentClassificationPrediction(
-            predicted_label=predicted_label, content_boost_factor=output_score
-        )
-        for predicted_label, output_score in zip(all_output_classes, prediction_scores)
-    ]
-
-    return content_classification_predictions
-
-
-def map_keywords(
-    input_ids: torch.Tensor, tokenizer: "PreTrainedTokenizer", is_keyword: list[bool]
-) -> list[str]:
-    tokens = tokenizer.convert_ids_to_tokens(input_ids)  # type: ignore
-
-    if not len(tokens) == len(is_keyword):
-        raise ValueError("Length of tokens and keyword predictions must match")
-
-    if input_ids[0] == tokenizer.cls_token_id:
-        tokens = tokens[1:]
-        is_keyword = is_keyword[1:]
-
-    if input_ids[-1] == tokenizer.sep_token_id:
-        tokens = tokens[:-1]
-        is_keyword = is_keyword[:-1]
-
-    unk_token = tokenizer.unk_token
-    if unk_token in tokens:
-        raise ValueError("Unknown token detected in the input")
-
-    keywords = []
-    current_keyword = ""
-
-    for ind, token in enumerate(tokens):
-        if is_keyword[ind]:
-            if token.startswith("##"):
-                current_keyword += token[2:]
-            else:
-                if current_keyword:
-                    keywords.append(current_keyword)
-                current_keyword = token
-        else:
-            # If mispredicted a later token of a keyword, add it to the current keyword
-            # to complete it
-            if current_keyword:
-                if len(current_keyword) > 2 and current_keyword.startswith("##"):
-                    current_keyword = current_keyword[2:]
-
-                else:
-                    keywords.append(current_keyword)
-                    current_keyword = ""
-
-    if current_keyword:
-        keywords.append(current_keyword)
-
-    return keywords
-
-
-def clean_keywords(keywords: list[str]) -> list[str]:
-    cleaned_words = []
-    for word in keywords:
-        word = word[:-2] if word.endswith("'s") else word
-        word = word.replace("/", " ")
-        word = word.replace("'", "").replace('"', "")
-        cleaned_words.extend([w for w in word.strip().split() if w and not w.isspace()])
-    return cleaned_words
-
-
-def run_connector_classification(req: ConnectorClassificationRequest) -> list[str]:
-    tokenizer = get_connector_classifier_tokenizer()
-    model = get_local_connector_classifier()
-
-    connector_names = req.available_connectors
-
-    input_ids, attention_mask = tokenize_connector_classification_query(
-        connector_names,
-        req.query,
-        tokenizer,
-        model.connector_end_token_id,
-    )
-    input_ids = input_ids.to(model.device)
-    attention_mask = attention_mask.to(model.device)
-
-    global_confidence, classifier_confidence = model(input_ids, attention_mask)
-
-    if global_confidence.item() < 0.5:
-        return []
-
-    passed_connectors = []
-
-    for i, connector_name in enumerate(connector_names):
-        if classifier_confidence.view(-1)[i].item() > 0.5:
-            passed_connectors.append(connector_name)
-
-    return passed_connectors
-
-
-def run_analysis(intent_req: IntentRequest) -> tuple[bool, list[str]]:
-    tokenizer = get_intent_model_tokenizer()
-    model_input = tokenizer(
-        intent_req.query, return_tensors="pt", truncation=False, padding=False
-    )
-
-    if len(model_input.input_ids[0]) > 512:
-        # If the user text is too long, assume it is semantic and keep all words
-        return True, intent_req.query.split()
-
-    intent_probs, token_probs = run_inference(model_input)
-
-    is_keyword_sequence = intent_probs[0] >= intent_req.keyword_percent_threshold
-
-    keyword_preds = [
-        token_prob >= intent_req.keyword_percent_threshold for token_prob in token_probs
-    ]
-
-    try:
-        keywords = map_keywords(model_input.input_ids[0], tokenizer, keyword_preds)
-    except Exception as e:
-        logger.warning(
-            f"Failed to extract keywords for query: {intent_req.query} due to {e}"
-        )
-        # Fallback to keeping all words
-        keywords = intent_req.query.split()
-
-    cleaned_keywords = clean_keywords(keywords)
-
-    return is_keyword_sequence, cleaned_keywords
-
-
-@router.post("/connector-classification")
-async def process_connector_classification_request(
-    classification_request: ConnectorClassificationRequest,
-) -> ConnectorClassificationResponse:
-    if INDEXING_ONLY:
-        raise RuntimeError(
-            "Indexing model server should not call connector classification endpoint"
-        )
-
-    if len(classification_request.available_connectors) == 0:
-        return ConnectorClassificationResponse(connectors=[])
-
-    connectors = run_connector_classification(classification_request)
-    return ConnectorClassificationResponse(connectors=connectors)
-
-
-@router.post("/query-analysis")
-async def process_analysis_request(
-    intent_request: IntentRequest,
-) -> IntentResponse:
-    if INDEXING_ONLY:
-        raise RuntimeError("Indexing model server should not call intent endpoint")
-
-    is_keyword, keywords = run_analysis(intent_request)
-    return IntentResponse(is_keyword=is_keyword, keywords=keywords)
-
-
-@router.post("/content-classification")
-async def process_content_classification_request(
-    content_classification_requests: list[str],
-) -> list[ContentClassificationPrediction]:
-    return run_content_classification_inference(content_classification_requests)

backend/model_server/encoders.py

@@ -1,7 +1,6 @@
 import asyncio
 import time
 from typing import Any
-from typing import Optional
 from typing import TYPE_CHECKING
 
 from fastapi import APIRouter
@@ -10,16 +9,13 @@ from fastapi import Request
 
 from model_server.utils import simple_log_function_time
 from onyx.utils.logger import setup_logger
-from shared_configs.configs import INDEXING_ONLY
 from shared_configs.enums import EmbedTextType
 from shared_configs.model_server_models import Embedding
 from shared_configs.model_server_models import EmbedRequest
 from shared_configs.model_server_models import EmbedResponse
-from shared_configs.model_server_models import RerankRequest
-from shared_configs.model_server_models import RerankResponse
 
 if TYPE_CHECKING:
-    from sentence_transformers import CrossEncoder, SentenceTransformer
+    from sentence_transformers import SentenceTransformer
 
 logger = setup_logger()
 
@@ -27,11 +23,6 @@ router = APIRouter(prefix="/encoder")
 
 
 _GLOBAL_MODELS_DICT: dict[str, "SentenceTransformer"] = {}
-_RERANK_MODEL: Optional["CrossEncoder"] = None
-
-# If we are not only indexing, dont want retry very long
-_RETRY_DELAY = 10 if INDEXING_ONLY else 0.1
-_RETRY_TRIES = 10 if INDEXING_ONLY else 2
 
 
 def get_embedding_model(
@@ -42,7 +33,7 @@ def get_embedding_model(
     Loads or returns a cached SentenceTransformer, sets max_seq_length, pins device,
     pre-warms rotary caches once, and wraps encode() with a lock to avoid cache races.
     """
-    from sentence_transformers import SentenceTransformer  # type: ignore
+    from sentence_transformers import SentenceTransformer
 
     def _prewarm_rope(st_model: "SentenceTransformer", target_len: int) -> None:
         """
@@ -87,19 +78,6 @@ def get_embedding_model(
     return _GLOBAL_MODELS_DICT[model_name]
 
 
-def get_local_reranking_model(
-    model_name: str,
-) -> "CrossEncoder":
-    global _RERANK_MODEL
-    from sentence_transformers import CrossEncoder  # type: ignore
-
-    if _RERANK_MODEL is None:
-        logger.notice(f"Loading {model_name}")
-        model = CrossEncoder(model_name)
-        _RERANK_MODEL = model
-    return _RERANK_MODEL
-
-
 ENCODING_RETRIES = 3
 ENCODING_RETRY_DELAY = 0.1
 
@@ -189,16 +167,6 @@ async def embed_text(
     return embeddings
 
 
-@simple_log_function_time()
-async def local_rerank(query: str, docs: list[str], model_name: str) -> list[float]:
-    cross_encoder = get_local_reranking_model(model_name)
-    # Run CPU-bound reranking in a thread pool
-    return await asyncio.get_event_loop().run_in_executor(
-        None,
-        lambda: cross_encoder.predict([(query, doc) for doc in docs]).tolist(),  # type: ignore
-    )
-
-
 @router.post("/bi-encoder-embed")
 async def route_bi_encoder_embed(
     request: Request,
@@ -254,39 +222,3 @@ async def process_embed_request(
         raise HTTPException(
             status_code=500, detail=f"Error during embedding process: {e}"
         )
-
-
-@router.post("/cross-encoder-scores")
-async def process_rerank_request(rerank_request: RerankRequest) -> RerankResponse:
-    """Cross encoders can be purely black box from the app perspective"""
-    # Only local models should use this endpoint - API providers should make direct API calls
-    if rerank_request.provider_type is not None:
-        raise ValueError(
-            f"Model server reranking endpoint should only be used for local models. "
-            f"API provider '{rerank_request.provider_type}' should make direct API calls instead."
-        )
-
-    if INDEXING_ONLY:
-        raise RuntimeError("Indexing model server should not call intent endpoint")
-
-    if not rerank_request.documents or not rerank_request.query:
-        raise HTTPException(
-            status_code=400, detail="Missing documents or query for reranking"
-        )
-    if not all(rerank_request.documents):
-        raise ValueError("Empty documents cannot be reranked.")
-
-    try:
-        # At this point, provider_type is None, so handle local reranking
-        sim_scores = await local_rerank(
-            query=rerank_request.query,
-            docs=rerank_request.documents,
-            model_name=rerank_request.model_name,
-        )
-        return RerankResponse(scores=sim_scores)
-
-    except Exception as e:
-        logger.exception(f"Error during reranking process:\n{str(e)}")
-        raise HTTPException(
-            status_code=500, detail="Failed to run Cross-Encoder reranking"
-        )

backend/model_server/legacy/README.md

@@ -0,0 +1,5 @@
+This directory contains code that was useful and may become useful again in the future.
+
+We stopped using rerankers because the state of the art rerankers are not significantly better than the biencoders and much worse than LLMs which are also capable of acting on a small set of documents for filtering, reranking, etc.
+
+We stopped using the internal query classifier as that's now offloaded to the LLM which does query expansion so we know ahead of time if it's a keyword or semantic query.

backend/model_server/legacy/custom_models.py

@@ -0,0 +1,573 @@
+# from typing import cast
+# from typing import Optional
+# from typing import TYPE_CHECKING
+
+# import numpy as np
+# import torch
+# import torch.nn.functional as F
+# from fastapi import APIRouter
+# from huggingface_hub import snapshot_download
+# from pydantic import BaseModel
+
+# from model_server.constants import MODEL_WARM_UP_STRING
+# from model_server.legacy.onyx_torch_model import ConnectorClassifier
+# from model_server.legacy.onyx_torch_model import HybridClassifier
+# from model_server.utils import simple_log_function_time
+# from onyx.utils.logger import setup_logger
+# from shared_configs.configs import CONNECTOR_CLASSIFIER_MODEL_REPO
+# from shared_configs.configs import CONNECTOR_CLASSIFIER_MODEL_TAG
+# from shared_configs.configs import INDEXING_ONLY
+# from shared_configs.configs import INTENT_MODEL_TAG
+# from shared_configs.configs import INTENT_MODEL_VERSION
+# from shared_configs.model_server_models import IntentRequest
+# from shared_configs.model_server_models import IntentResponse
+
+# if TYPE_CHECKING:
+#     from setfit import SetFitModel  # type: ignore[import-untyped]
+#     from transformers import PreTrainedTokenizer, BatchEncoding
+
+
+# INFORMATION_CONTENT_MODEL_WARM_UP_STRING = "hi" * 50
+
+# INDEXING_INFORMATION_CONTENT_CLASSIFICATION_MAX = 1.0
+# INDEXING_INFORMATION_CONTENT_CLASSIFICATION_MIN = 0.7
+# INDEXING_INFORMATION_CONTENT_CLASSIFICATION_TEMPERATURE = 4.0
+# INDEXING_INFORMATION_CONTENT_CLASSIFICATION_CUTOFF_LENGTH = 10
+# INFORMATION_CONTENT_MODEL_VERSION = "onyx-dot-app/information-content-model"
+# INFORMATION_CONTENT_MODEL_TAG: str | None = None
+
+
+# class ConnectorClassificationRequest(BaseModel):
+#     available_connectors: list[str]
+#     query: str
+
+
+# class ConnectorClassificationResponse(BaseModel):
+#     connectors: list[str]
+
+
+# class ContentClassificationPrediction(BaseModel):
+#     predicted_label: int
+#     content_boost_factor: float
+
+
+# logger = setup_logger()
+
+# router = APIRouter(prefix="/custom")
+
+# _CONNECTOR_CLASSIFIER_TOKENIZER: Optional["PreTrainedTokenizer"] = None
+# _CONNECTOR_CLASSIFIER_MODEL: ConnectorClassifier | None = None
+
+# _INTENT_TOKENIZER: Optional["PreTrainedTokenizer"] = None
+# _INTENT_MODEL: HybridClassifier | None = None
+
+# _INFORMATION_CONTENT_MODEL: Optional["SetFitModel"] = None
+
+# _INFORMATION_CONTENT_MODEL_PROMPT_PREFIX: str = ""  # spec to model version!
+
+
+# def get_connector_classifier_tokenizer() -> "PreTrainedTokenizer":
+#     global _CONNECTOR_CLASSIFIER_TOKENIZER
+#     from transformers import AutoTokenizer, PreTrainedTokenizer
+
+#     if _CONNECTOR_CLASSIFIER_TOKENIZER is None:
+#         # The tokenizer details are not uploaded to the HF hub since it's just the
+#         # unmodified distilbert tokenizer.
+#         _CONNECTOR_CLASSIFIER_TOKENIZER = cast(
+#             PreTrainedTokenizer,
+#             AutoTokenizer.from_pretrained("distilbert-base-uncased"),
+#         )
+#     return _CONNECTOR_CLASSIFIER_TOKENIZER
+
+
+# def get_local_connector_classifier(
+#     model_name_or_path: str = CONNECTOR_CLASSIFIER_MODEL_REPO,
+#     tag: str = CONNECTOR_CLASSIFIER_MODEL_TAG,
+# ) -> ConnectorClassifier:
+#     global _CONNECTOR_CLASSIFIER_MODEL
+#     if _CONNECTOR_CLASSIFIER_MODEL is None:
+#         try:
+#             # Calculate where the cache should be, then load from local if available
+#             local_path = snapshot_download(
+#                 repo_id=model_name_or_path, revision=tag, local_files_only=True
+#             )
+#             _CONNECTOR_CLASSIFIER_MODEL = ConnectorClassifier.from_pretrained(
+#                 local_path
+#             )
+#         except Exception as e:
+#             logger.warning(f"Failed to load model directly: {e}")
+#             try:
+#                 # Attempt to download the model snapshot
+#                 logger.info(f"Downloading model snapshot for {model_name_or_path}")
+#                 local_path = snapshot_download(repo_id=model_name_or_path, revision=tag)
+#                 _CONNECTOR_CLASSIFIER_MODEL = ConnectorClassifier.from_pretrained(
+#                     local_path
+#                 )
+#             except Exception as e:
+#                 logger.error(
+#                     f"Failed to load model even after attempted snapshot download: {e}"
+#                 )
+#                 raise
+#     return _CONNECTOR_CLASSIFIER_MODEL
+
+
+# def get_intent_model_tokenizer() -> "PreTrainedTokenizer":
+#     from transformers import AutoTokenizer, PreTrainedTokenizer
+
+#     global _INTENT_TOKENIZER
+#     if _INTENT_TOKENIZER is None:
+#         # The tokenizer details are not uploaded to the HF hub since it's just the
+#         # unmodified distilbert tokenizer.
+#         _INTENT_TOKENIZER = cast(
+#             PreTrainedTokenizer,
+#             AutoTokenizer.from_pretrained("distilbert-base-uncased"),
+#         )
+#     return _INTENT_TOKENIZER
+
+
+# def get_local_intent_model(
+#     model_name_or_path: str = INTENT_MODEL_VERSION,
+#     tag: str | None = INTENT_MODEL_TAG,
+# ) -> HybridClassifier:
+#     global _INTENT_MODEL
+#     if _INTENT_MODEL is None:
+#         try:
+#             # Calculate where the cache should be, then load from local if available
+#             logger.notice(f"Loading model from local cache: {model_name_or_path}")
+#             local_path = snapshot_download(
+#                 repo_id=model_name_or_path, revision=tag, local_files_only=True
+#             )
+#             _INTENT_MODEL = HybridClassifier.from_pretrained(local_path)
+#             logger.notice(f"Loaded model from local cache: {local_path}")
+#         except Exception as e:
+#             logger.warning(f"Failed to load model directly: {e}")
+#             try:
+#                 # Attempt to download the model snapshot
+#                 logger.notice(f"Downloading model snapshot for {model_name_or_path}")
+#                 local_path = snapshot_download(
+#                     repo_id=model_name_or_path, revision=tag, local_files_only=False
+#                 )
+#                 _INTENT_MODEL = HybridClassifier.from_pretrained(local_path)
+#             except Exception as e:
+#                 logger.error(
+#                     f"Failed to load model even after attempted snapshot download: {e}"
+#                 )
+#                 raise
+#     return _INTENT_MODEL
+
+
+# def get_local_information_content_model(
+#     model_name_or_path: str = INFORMATION_CONTENT_MODEL_VERSION,
+#     tag: str | None = INFORMATION_CONTENT_MODEL_TAG,
+# ) -> "SetFitModel":
+#     from setfit import SetFitModel
+
+#     global _INFORMATION_CONTENT_MODEL
+#     if _INFORMATION_CONTENT_MODEL is None:
+#         try:
+#             # Calculate where the cache should be, then load from local if available
+#             logger.notice(
+#                 f"Loading content information model from local cache: {model_name_or_path}"
+#             )
+#             local_path = snapshot_download(
+#                 repo_id=model_name_or_path, revision=tag, local_files_only=True
+#             )
+#             _INFORMATION_CONTENT_MODEL = SetFitModel.from_pretrained(local_path)
+#             logger.notice(
+#                 f"Loaded content information model from local cache: {local_path}"
+#             )
+#         except Exception as e:
+#             logger.warning(f"Failed to load content information model directly: {e}")
+#             try:
+#                 # Attempt to download the model snapshot
+#                 logger.notice(
+#                     f"Downloading content information model snapshot for {model_name_or_path}"
+#                 )
+#                 local_path = snapshot_download(
+#                     repo_id=model_name_or_path, revision=tag, local_files_only=False
+#                 )
+#                 _INFORMATION_CONTENT_MODEL = SetFitModel.from_pretrained(local_path)
+#             except Exception as e:
+#                 logger.error(
+#                     f"Failed to load content information model even after attempted snapshot download: {e}"
+#                 )
+#                 raise
+
+#     return _INFORMATION_CONTENT_MODEL
+
+
+# def tokenize_connector_classification_query(
+#     connectors: list[str],
+#     query: str,
+#     tokenizer: "PreTrainedTokenizer",
+#     connector_token_end_id: int,
+# ) -> tuple[torch.Tensor, torch.Tensor]:
+#     """
+#     Tokenize the connectors & user query into one prompt for the forward pass of ConnectorClassifier models
+
+#     The attention mask is just all 1s. The prompt is CLS + each connector name suffixed with the connector end
+#     token and then the user query.
+#     """
+
+#     input_ids = torch.tensor([tokenizer.cls_token_id], dtype=torch.long)
+
+#     for connector in connectors:
+#         connector_token_ids = tokenizer(
+#             connector,
+#             add_special_tokens=False,
+#             return_tensors="pt",
+#         )
+
+#         input_ids = torch.cat(
+#             (
+#                 input_ids,
+#                 connector_token_ids["input_ids"].squeeze(dim=0),
+#                 torch.tensor([connector_token_end_id], dtype=torch.long),
+#             ),
+#             dim=-1,
+#         )
+#     query_token_ids = tokenizer(
+#         query,
+#         add_special_tokens=False,
+#         return_tensors="pt",
+#     )
+
+#     input_ids = torch.cat(
+#         (
+#             input_ids,
+#             query_token_ids["input_ids"].squeeze(dim=0),
+#             torch.tensor([tokenizer.sep_token_id], dtype=torch.long),
+#         ),
+#         dim=-1,
+#     )
+#     attention_mask = torch.ones(input_ids.numel(), dtype=torch.long)
+
+#     return input_ids.unsqueeze(0), attention_mask.unsqueeze(0)
+
+
+# def warm_up_connector_classifier_model() -> None:
+#     logger.info(
+#         f"Warming up connector_classifier model {CONNECTOR_CLASSIFIER_MODEL_TAG}"
+#     )
+#     connector_classifier_tokenizer = get_connector_classifier_tokenizer()
+#     connector_classifier = get_local_connector_classifier()
+
+#     input_ids, attention_mask = tokenize_connector_classification_query(
+#         ["GitHub"],
+#         "onyx classifier query google doc",
+#         connector_classifier_tokenizer,
+#         connector_classifier.connector_end_token_id,
+#     )
+#     input_ids = input_ids.to(connector_classifier.device)
+#     attention_mask = attention_mask.to(connector_classifier.device)
+
+#     connector_classifier(input_ids, attention_mask)
+
+
+# def warm_up_intent_model() -> None:
+#     logger.notice(f"Warming up Intent Model: {INTENT_MODEL_VERSION}")
+#     intent_tokenizer = get_intent_model_tokenizer()
+#     tokens = intent_tokenizer(
+#         MODEL_WARM_UP_STRING, return_tensors="pt", truncation=True, padding=True
+#     )
+
+#     intent_model = get_local_intent_model()
+#     device = intent_model.device
+#     intent_model(
+#         query_ids=tokens["input_ids"].to(device),
+#         query_mask=tokens["attention_mask"].to(device),
+#     )
+
+
+# def warm_up_information_content_model() -> None:
+#     logger.notice("Warming up Content Model")  # TODO: add version if needed
+
+#     information_content_model = get_local_information_content_model()
+#     information_content_model(INFORMATION_CONTENT_MODEL_WARM_UP_STRING)
+
+
+# @simple_log_function_time()
+# def run_inference(tokens: "BatchEncoding") -> tuple[list[float], list[float]]:
+#     intent_model = get_local_intent_model()
+#     device = intent_model.device
+
+#     outputs = intent_model(
+#         query_ids=tokens["input_ids"].to(device),
+#         query_mask=tokens["attention_mask"].to(device),
+#     )
+
+#     token_logits = outputs["token_logits"]
+#     intent_logits = outputs["intent_logits"]
+
+#     # Move tensors to CPU before applying softmax and converting to numpy
+#     intent_probabilities = F.softmax(intent_logits.cpu(), dim=-1).numpy()[0]
+#     token_probabilities = F.softmax(token_logits.cpu(), dim=-1).numpy()[0]
+
+#     # Extract the probabilities for the positive class (index 1) for each token
+#     token_positive_probs = token_probabilities[:, 1].tolist()
+
+#     return intent_probabilities.tolist(), token_positive_probs
+
+
+# @simple_log_function_time()
+# def run_content_classification_inference(
+#     text_inputs: list[str],
+# ) -> list[ContentClassificationPrediction]:
+#     """
+#     Assign a score to the segments in question. The model stored in get_local_information_content_model()
+#     creates the 'model score' based on its training, and the scores are then converted to a 0.0-1.0 scale.
+#     In the code outside of the model/inference model servers that score will be converted into the actual
+#     boost factor.
+#     """
+
+#     def _prob_to_score(prob: float) -> float:
+#         """
+#         Conversion of base score to 0.0 - 1.0 score. Note that the min/max values depend on the model!
+#         """
+#         _MIN_BASE_SCORE = 0.25
+#         _MAX_BASE_SCORE = 0.75
+#         if prob < _MIN_BASE_SCORE:
+#             raw_score = 0.0
+#         elif prob < _MAX_BASE_SCORE:
+#             raw_score = (prob - _MIN_BASE_SCORE) / (_MAX_BASE_SCORE - _MIN_BASE_SCORE)
+#         else:
+#             raw_score = 1.0
+#         return (
+#             INDEXING_INFORMATION_CONTENT_CLASSIFICATION_MIN
+#             + (
+#                 INDEXING_INFORMATION_CONTENT_CLASSIFICATION_MAX
+#                 - INDEXING_INFORMATION_CONTENT_CLASSIFICATION_MIN
+#             )
+#             * raw_score
+#         )
+
+#     _BATCH_SIZE = 32
+#     content_model = get_local_information_content_model()
+
+#     # Process inputs in batches
+#     all_output_classes: list[int] = []
+#     all_base_output_probabilities: list[float] = []
+
+#     for i in range(0, len(text_inputs), _BATCH_SIZE):
+#         batch = text_inputs[i : i + _BATCH_SIZE]
+#         batch_with_prefix = []
+#         batch_indices = []
+
+#         # Pre-allocate results for this batch
+#         batch_output_classes: list[np.ndarray] = [np.array(1)] * len(batch)
+#         batch_probabilities: list[np.ndarray] = [np.array(1.0)] * len(batch)
+
+#         # Pre-process batch to handle long input exceptions
+#         for j, text in enumerate(batch):
+#             if len(text) == 0:
+#                 # if no input, treat as non-informative from the model's perspective
+#                 batch_output_classes[j] = np.array(0)
+#                 batch_probabilities[j] = np.array(0.0)
+#                 logger.warning("Input for Content Information Model is empty")
+
+#             elif (
+#                 len(text.split())
+#                 <= INDEXING_INFORMATION_CONTENT_CLASSIFICATION_CUTOFF_LENGTH
+#             ):
+#                 # if input is short, use the model
+#                 batch_with_prefix.append(
+#                     _INFORMATION_CONTENT_MODEL_PROMPT_PREFIX + text
+#                 )
+#                 batch_indices.append(j)
+#             else:
+#                 # if longer than cutoff, treat as informative (stay with default), but issue warning
+#                 logger.warning("Input for Content Information Model too long")
+
+#         if batch_with_prefix:  # Only run model if we have valid inputs
+#             # Get predictions for the batch
+#             model_output_classes = content_model(batch_with_prefix)
+#             model_output_probabilities = content_model.predict_proba(batch_with_prefix)
+
+#             # Place results in the correct positions
+#             for idx, batch_idx in enumerate(batch_indices):
+#                 batch_output_classes[batch_idx] = model_output_classes[idx].numpy()
+#                 batch_probabilities[batch_idx] = model_output_probabilities[idx][
+#                     1
+#                 ].numpy()  # x[1] is prob of the positive class
+
+#         all_output_classes.extend([int(x) for x in batch_output_classes])
+#         all_base_output_probabilities.extend([float(x) for x in batch_probabilities])
+
+#     logits = [
+#         np.log(p / (1 - p)) if p != 0.0 and p != 1.0 else (100 if p == 1.0 else -100)
+#         for p in all_base_output_probabilities
+#     ]
+#     scaled_logits = [
+#         logit / INDEXING_INFORMATION_CONTENT_CLASSIFICATION_TEMPERATURE
+#         for logit in logits
+#     ]
+#     output_probabilities_with_temp = [
+#         np.exp(scaled_logit) / (1 + np.exp(scaled_logit))
+#         for scaled_logit in scaled_logits
+#     ]
+
+#     prediction_scores = [
+#         _prob_to_score(p_temp) for p_temp in output_probabilities_with_temp
+#     ]
+
+#     content_classification_predictions = [
+#         ContentClassificationPrediction(
+#             predicted_label=predicted_label, content_boost_factor=output_score
+#         )
+#         for predicted_label, output_score in zip(all_output_classes, prediction_scores)
+#     ]
+
+#     return content_classification_predictions
+
+
+# def map_keywords(
+#     input_ids: torch.Tensor, tokenizer: "PreTrainedTokenizer", is_keyword: list[bool]
+# ) -> list[str]:
+#     tokens = tokenizer.convert_ids_to_tokens(input_ids)  # type: ignore
+
+#     if not len(tokens) == len(is_keyword):
+#         raise ValueError("Length of tokens and keyword predictions must match")
+
+#     if input_ids[0] == tokenizer.cls_token_id:
+#         tokens = tokens[1:]
+#         is_keyword = is_keyword[1:]
+
+#     if input_ids[-1] == tokenizer.sep_token_id:
+#         tokens = tokens[:-1]
+#         is_keyword = is_keyword[:-1]
+
+#     unk_token = tokenizer.unk_token
+#     if unk_token in tokens:
+#         raise ValueError("Unknown token detected in the input")
+
+#     keywords = []
+#     current_keyword = ""
+
+#     for ind, token in enumerate(tokens):
+#         if is_keyword[ind]:
+#             if token.startswith("##"):
+#                 current_keyword += token[2:]
+#             else:
+#                 if current_keyword:
+#                     keywords.append(current_keyword)
+#                 current_keyword = token
+#         else:
+#             # If mispredicted a later token of a keyword, add it to the current keyword
+#             # to complete it
+#             if current_keyword:
+#                 if len(current_keyword) > 2 and current_keyword.startswith("##"):
+#                     current_keyword = current_keyword[2:]
+
+#                 else:
+#                     keywords.append(current_keyword)
+#                     current_keyword = ""
+
+#     if current_keyword:
+#         keywords.append(current_keyword)
+
+#     return keywords
+
+
+# def clean_keywords(keywords: list[str]) -> list[str]:
+#     cleaned_words = []
+#     for word in keywords:
+#         word = word[:-2] if word.endswith("'s") else word
+#         word = word.replace("/", " ")
+#         word = word.replace("'", "").replace('"', "")
+#         cleaned_words.extend([w for w in word.strip().split() if w and not w.isspace()])
+#     return cleaned_words
+
+
+# def run_connector_classification(req: ConnectorClassificationRequest) -> list[str]:
+#     tokenizer = get_connector_classifier_tokenizer()
+#     model = get_local_connector_classifier()
+
+#     connector_names = req.available_connectors
+
+#     input_ids, attention_mask = tokenize_connector_classification_query(
+#         connector_names,
+#         req.query,
+#         tokenizer,
+#         model.connector_end_token_id,
+#     )
+#     input_ids = input_ids.to(model.device)
+#     attention_mask = attention_mask.to(model.device)
+
+#     global_confidence, classifier_confidence = model(input_ids, attention_mask)
+
+#     if global_confidence.item() < 0.5:
+#         return []
+
+#     passed_connectors = []
+
+#     for i, connector_name in enumerate(connector_names):
+#         if classifier_confidence.view(-1)[i].item() > 0.5:
+#             passed_connectors.append(connector_name)
+
+#     return passed_connectors
+
+
+# def run_analysis(intent_req: IntentRequest) -> tuple[bool, list[str]]:
+#     tokenizer = get_intent_model_tokenizer()
+#     model_input = tokenizer(
+#         intent_req.query, return_tensors="pt", truncation=False, padding=False
+#     )
+
+#     if len(model_input.input_ids[0]) > 512:
+#         # If the user text is too long, assume it is semantic and keep all words
+#         return True, intent_req.query.split()
+
+#     intent_probs, token_probs = run_inference(model_input)
+
+#     is_keyword_sequence = intent_probs[0] >= intent_req.keyword_percent_threshold
+
+#     keyword_preds = [
+#         token_prob >= intent_req.keyword_percent_threshold for token_prob in token_probs
+#     ]
+
+#     try:
+#         keywords = map_keywords(model_input.input_ids[0], tokenizer, keyword_preds)
+#     except Exception as e:
+#         logger.warning(
+#             f"Failed to extract keywords for query: {intent_req.query} due to {e}"
+#         )
+#         # Fallback to keeping all words
+#         keywords = intent_req.query.split()
+
+#     cleaned_keywords = clean_keywords(keywords)
+
+#     return is_keyword_sequence, cleaned_keywords
+
+
+# @router.post("/connector-classification")
+# async def process_connector_classification_request(
+#     classification_request: ConnectorClassificationRequest,
+# ) -> ConnectorClassificationResponse:
+#     if INDEXING_ONLY:
+#         raise RuntimeError(
+#             "Indexing model server should not call connector classification endpoint"
+#         )
+
+#     if len(classification_request.available_connectors) == 0:
+#         return ConnectorClassificationResponse(connectors=[])
+
+#     connectors = run_connector_classification(classification_request)
+#     return ConnectorClassificationResponse(connectors=connectors)
+
+
+# @router.post("/query-analysis")
+# async def process_analysis_request(
+#     intent_request: IntentRequest,
+# ) -> IntentResponse:
+#     if INDEXING_ONLY:
+#         raise RuntimeError("Indexing model server should not call intent endpoint")
+
+#     is_keyword, keywords = run_analysis(intent_request)
+#     return IntentResponse(is_keyword=is_keyword, keywords=keywords)
+
+
+# @router.post("/content-classification")
+# async def process_content_classification_request(
+#     content_classification_requests: list[str],
+# ) -> list[ContentClassificationPrediction]:
+#     return run_content_classification_inference(content_classification_requests)

backend/model_server/legacy/onyx_torch_model.py

@@ -0,0 +1,154 @@
+# import json
+# import os
+# from typing import cast
+# from typing import TYPE_CHECKING
+
+# import torch
+# import torch.nn as nn
+
+
+# if TYPE_CHECKING:
+#     from transformers import DistilBertConfig
+
+
+# class HybridClassifier(nn.Module):
+#     def __init__(self) -> None:
+#         from transformers import DistilBertConfig, DistilBertModel
+
+#         super().__init__()
+#         config = DistilBertConfig()
+#         self.distilbert = DistilBertModel(config)
+#         config = self.distilbert.config  # type: ignore
+
+#         # Keyword tokenwise binary classification layer
+#         self.keyword_classifier = nn.Linear(config.dim, 2)
+
+#         # Intent Classifier layers
+#         self.pre_classifier = nn.Linear(config.dim, config.dim)
+#         self.intent_classifier = nn.Linear(config.dim, 2)
+
+#         self.device = torch.device("cpu")
+
+#     def forward(
+#         self,
+#         query_ids: torch.Tensor,
+#         query_mask: torch.Tensor,
+#     ) -> dict[str, torch.Tensor]:
+#         outputs = self.distilbert(input_ids=query_ids, attention_mask=query_mask)
+#         sequence_output = outputs.last_hidden_state
+
+#         # Intent classification on the CLS token
+#         cls_token_state = sequence_output[:, 0, :]
+#         pre_classifier_out = self.pre_classifier(cls_token_state)
+#         intent_logits = self.intent_classifier(pre_classifier_out)
+
+#         # Keyword classification on all tokens
+#         token_logits = self.keyword_classifier(sequence_output)
+
+#         return {"intent_logits": intent_logits, "token_logits": token_logits}
+
+#     @classmethod
+#     def from_pretrained(cls, load_directory: str) -> "HybridClassifier":
+#         model_path = os.path.join(load_directory, "pytorch_model.bin")
+#         config_path = os.path.join(load_directory, "config.json")
+
+#         with open(config_path, "r") as f:
+#             config = json.load(f)
+#         model = cls(**config)
+
+#         if torch.backends.mps.is_available():
+#             # Apple silicon GPU
+#             device = torch.device("mps")
+#         elif torch.cuda.is_available():
+#             device = torch.device("cuda")
+#         else:
+#             device = torch.device("cpu")
+
+#         model.load_state_dict(torch.load(model_path, map_location=device))
+#         model = model.to(device)
+
+#         model.device = device
+
+#         model.eval()
+#         # Eval doesn't set requires_grad to False, do it manually to save memory and have faster inference
+#         for param in model.parameters():
+#             param.requires_grad = False
+
+#         return model
+
+
+# class ConnectorClassifier(nn.Module):
+#     def __init__(self, config: "DistilBertConfig") -> None:
+#         from transformers import DistilBertTokenizer, DistilBertModel
+
+#         super().__init__()
+
+#         self.config = config
+#         self.distilbert = DistilBertModel(config)
+#         config = self.distilbert.config  # type: ignore
+#         self.connector_global_classifier = nn.Linear(config.dim, 1)
+#         self.connector_match_classifier = nn.Linear(config.dim, 1)
+#         self.tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-uncased")
+
+#         # Token indicating end of connector name, and on which classifier is used
+#         self.connector_end_token_id = self.tokenizer.get_vocab()[
+#             self.config.connector_end_token
+#         ]
+
+#         self.device = torch.device("cpu")
+
+#     def forward(
+#         self,
+#         input_ids: torch.Tensor,
+#         attention_mask: torch.Tensor,
+#     ) -> tuple[torch.Tensor, torch.Tensor]:
+#         hidden_states = self.distilbert(
+#             input_ids=input_ids, attention_mask=attention_mask
+#         ).last_hidden_state
+
+#         cls_hidden_states = hidden_states[
+#             :, 0, :
+#         ]  # Take leap of faith that first token is always [CLS]
+#         global_logits = self.connector_global_classifier(cls_hidden_states).view(-1)
+#         global_confidence = torch.sigmoid(global_logits).view(-1)
+
+#         connector_end_position_ids = input_ids == self.connector_end_token_id
+#         connector_end_hidden_states = hidden_states[connector_end_position_ids]
+#         classifier_output = self.connector_match_classifier(connector_end_hidden_states)
+#         classifier_confidence = torch.nn.functional.sigmoid(classifier_output).view(-1)
+
+#         return global_confidence, classifier_confidence
+
+#     @classmethod
+#     def from_pretrained(cls, repo_dir: str) -> "ConnectorClassifier":
+#         from transformers import DistilBertConfig
+
+#         config = cast(
+#             DistilBertConfig,
+#             DistilBertConfig.from_pretrained(os.path.join(repo_dir, "config.json")),
+#         )
+#         device = (
+#             torch.device("cuda")
+#             if torch.cuda.is_available()
+#             else (
+#                 torch.device("mps")
+#                 if torch.backends.mps.is_available()
+#                 else torch.device("cpu")
+#             )
+#         )
+#         state_dict = torch.load(
+#             os.path.join(repo_dir, "pytorch_model.pt"),
+#             map_location=device,
+#             weights_only=True,
+#         )
+
+#         model = cls(config)
+#         model.load_state_dict(state_dict)
+#         model.to(device)
+#         model.device = device
+#         model.eval()
+
+#         for param in model.parameters():
+#             param.requires_grad = False
+
+#         return model

backend/model_server/legacy/reranker.py

@@ -0,0 +1,80 @@
+# import asyncio
+# from typing import Optional
+# from typing import TYPE_CHECKING
+
+# from fastapi import APIRouter
+# from fastapi import HTTPException
+
+# from model_server.utils import simple_log_function_time
+# from onyx.utils.logger import setup_logger
+# from shared_configs.configs import INDEXING_ONLY
+# from shared_configs.model_server_models import RerankRequest
+# from shared_configs.model_server_models import RerankResponse
+
+# if TYPE_CHECKING:
+#     from sentence_transformers import CrossEncoder
+
+# logger = setup_logger()
+
+# router = APIRouter(prefix="/encoder")
+
+# _RERANK_MODEL: Optional["CrossEncoder"] = None
+
+
+# def get_local_reranking_model(
+#     model_name: str,
+# ) -> "CrossEncoder":
+#     global _RERANK_MODEL
+#     from sentence_transformers import CrossEncoder
+
+#     if _RERANK_MODEL is None:
+#         logger.notice(f"Loading {model_name}")
+#         model = CrossEncoder(model_name)
+#         _RERANK_MODEL = model
+#     return _RERANK_MODEL
+
+
+# @simple_log_function_time()
+# async def local_rerank(query: str, docs: list[str], model_name: str) -> list[float]:
+#     cross_encoder = get_local_reranking_model(model_name)
+#     # Run CPU-bound reranking in a thread pool
+#     return await asyncio.get_event_loop().run_in_executor(
+#         None,
+#         lambda: cross_encoder.predict([(query, doc) for doc in docs]).tolist(),
+#     )
+
+
+# @router.post("/cross-encoder-scores")
+# async def process_rerank_request(rerank_request: RerankRequest) -> RerankResponse:
+#     """Cross encoders can be purely black box from the app perspective"""
+#     # Only local models should use this endpoint - API providers should make direct API calls
+#     if rerank_request.provider_type is not None:
+#         raise ValueError(
+#             f"Model server reranking endpoint should only be used for local models. "
+#             f"API provider '{rerank_request.provider_type}' should make direct API calls instead."
+#         )
+
+#     if INDEXING_ONLY:
+#         raise RuntimeError("Indexing model server should not call reranking endpoint")
+
+#     if not rerank_request.documents or not rerank_request.query:
+#         raise HTTPException(
+#             status_code=400, detail="Missing documents or query for reranking"
+#         )
+#     if not all(rerank_request.documents):
+#         raise ValueError("Empty documents cannot be reranked.")
+
+#     try:
+#         # At this point, provider_type is None, so handle local reranking
+#         sim_scores = await local_rerank(
+#             query=rerank_request.query,
+#             docs=rerank_request.documents,
+#             model_name=rerank_request.model_name,
+#         )
+#         return RerankResponse(scores=sim_scores)
+
+#     except Exception as e:
+#         logger.exception(f"Error during reranking process:\n{str(e)}")
+#         raise HTTPException(
+#             status_code=500, detail="Failed to run Cross-Encoder reranking"
+#         )

backend/model_server/main.py

@@ -12,11 +12,8 @@ from fastapi import FastAPI
 from prometheus_fastapi_instrumentator import Instrumentator
 from sentry_sdk.integrations.fastapi import FastApiIntegration
 from sentry_sdk.integrations.starlette import StarletteIntegration
-from transformers import logging as transformer_logging  # type:ignore
+from transformers import logging as transformer_logging
 
-from model_server.custom_models import router as custom_models_router
-from model_server.custom_models import warm_up_information_content_model
-from model_server.custom_models import warm_up_intent_model
 from model_server.encoders import router as encoders_router
 from model_server.management_endpoints import router as management_router
 from model_server.utils import get_gpu_type
@@ -30,7 +27,6 @@ from shared_configs.configs import MIN_THREADS_ML_MODELS
 from shared_configs.configs import MODEL_SERVER_ALLOWED_HOST
 from shared_configs.configs import MODEL_SERVER_PORT
 from shared_configs.configs import SENTRY_DSN
-from shared_configs.configs import SKIP_WARM_UP
 
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 os.environ["HF_HUB_DISABLE_TELEMETRY"] = "1"
@@ -92,18 +88,6 @@ async def lifespan(app: FastAPI) -> AsyncGenerator:
     torch.set_num_threads(max(MIN_THREADS_ML_MODELS, torch.get_num_threads()))
     logger.notice(f"Torch Threads: {torch.get_num_threads()}")
 
-    if not SKIP_WARM_UP:
-        if not INDEXING_ONLY:
-            logger.notice("Warming up intent model for inference model server")
-            warm_up_intent_model()
-        else:
-            logger.notice(
-                "Warming up content information model for indexing model server"
-            )
-            warm_up_information_content_model()
-    else:
-        logger.notice("Skipping model warmup due to SKIP_WARM_UP=true")
-
     yield
 
 
@@ -123,7 +107,6 @@ def get_model_app() -> FastAPI:
 
     application.include_router(management_router)
     application.include_router(encoders_router)
-    application.include_router(custom_models_router)
 
     request_id_prefix = "INF"
     if INDEXING_ONLY:

backend/model_server/onyx_torch_model.py

@@ -1,154 +0,0 @@
-import json
-import os
-from typing import cast
-from typing import TYPE_CHECKING
-
-import torch
-import torch.nn as nn
-
-
-if TYPE_CHECKING:
-    from transformers import DistilBertConfig  # type: ignore
-
-
-class HybridClassifier(nn.Module):
-    def __init__(self) -> None:
-        from transformers import DistilBertConfig, DistilBertModel
-
-        super().__init__()
-        config = DistilBertConfig()
-        self.distilbert = DistilBertModel(config)
-        config = self.distilbert.config  # type: ignore
-
-        # Keyword tokenwise binary classification layer
-        self.keyword_classifier = nn.Linear(config.dim, 2)
-
-        # Intent Classifier layers
-        self.pre_classifier = nn.Linear(config.dim, config.dim)
-        self.intent_classifier = nn.Linear(config.dim, 2)
-
-        self.device = torch.device("cpu")
-
-    def forward(
-        self,
-        query_ids: torch.Tensor,
-        query_mask: torch.Tensor,
-    ) -> dict[str, torch.Tensor]:
-        outputs = self.distilbert(input_ids=query_ids, attention_mask=query_mask)  # type: ignore
-        sequence_output = outputs.last_hidden_state
-
-        # Intent classification on the CLS token
-        cls_token_state = sequence_output[:, 0, :]
-        pre_classifier_out = self.pre_classifier(cls_token_state)
-        intent_logits = self.intent_classifier(pre_classifier_out)
-
-        # Keyword classification on all tokens
-        token_logits = self.keyword_classifier(sequence_output)
-
-        return {"intent_logits": intent_logits, "token_logits": token_logits}
-
-    @classmethod
-    def from_pretrained(cls, load_directory: str) -> "HybridClassifier":
-        model_path = os.path.join(load_directory, "pytorch_model.bin")
-        config_path = os.path.join(load_directory, "config.json")
-
-        with open(config_path, "r") as f:
-            config = json.load(f)
-        model = cls(**config)
-
-        if torch.backends.mps.is_available():
-            # Apple silicon GPU
-            device = torch.device("mps")
-        elif torch.cuda.is_available():
-            device = torch.device("cuda")
-        else:
-            device = torch.device("cpu")
-
-        model.load_state_dict(torch.load(model_path, map_location=device))
-        model = model.to(device)
-
-        model.device = device
-
-        model.eval()
-        # Eval doesn't set requires_grad to False, do it manually to save memory and have faster inference
-        for param in model.parameters():
-            param.requires_grad = False
-
-        return model
-
-
-class ConnectorClassifier(nn.Module):
-    def __init__(self, config: "DistilBertConfig") -> None:
-        from transformers import DistilBertTokenizer, DistilBertModel
-
-        super().__init__()
-
-        self.config = config
-        self.distilbert = DistilBertModel(config)
-        config = self.distilbert.config  # type: ignore
-        self.connector_global_classifier = nn.Linear(config.dim, 1)
-        self.connector_match_classifier = nn.Linear(config.dim, 1)
-        self.tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-uncased")
-
-        # Token indicating end of connector name, and on which classifier is used
-        self.connector_end_token_id = self.tokenizer.get_vocab()[
-            self.config.connector_end_token
-        ]
-
-        self.device = torch.device("cpu")
-
-    def forward(
-        self,
-        input_ids: torch.Tensor,
-        attention_mask: torch.Tensor,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        hidden_states = self.distilbert(  # type: ignore
-            input_ids=input_ids, attention_mask=attention_mask
-        ).last_hidden_state
-
-        cls_hidden_states = hidden_states[
-            :, 0, :
-        ]  # Take leap of faith that first token is always [CLS]
-        global_logits = self.connector_global_classifier(cls_hidden_states).view(-1)
-        global_confidence = torch.sigmoid(global_logits).view(-1)
-
-        connector_end_position_ids = input_ids == self.connector_end_token_id
-        connector_end_hidden_states = hidden_states[connector_end_position_ids]
-        classifier_output = self.connector_match_classifier(connector_end_hidden_states)
-        classifier_confidence = torch.nn.functional.sigmoid(classifier_output).view(-1)
-
-        return global_confidence, classifier_confidence
-
-    @classmethod
-    def from_pretrained(cls, repo_dir: str) -> "ConnectorClassifier":
-        from transformers import DistilBertConfig
-
-        config = cast(
-            DistilBertConfig,
-            DistilBertConfig.from_pretrained(os.path.join(repo_dir, "config.json")),
-        )
-        device = (
-            torch.device("cuda")
-            if torch.cuda.is_available()
-            else (
-                torch.device("mps")
-                if torch.backends.mps.is_available()
-                else torch.device("cpu")
-            )
-        )
-        state_dict = torch.load(
-            os.path.join(repo_dir, "pytorch_model.pt"),
-            map_location=device,
-            weights_only=True,
-        )
-
-        model = cls(config)
-        model.load_state_dict(state_dict)
-        model.to(device)
-        model.device = device
-        model.eval()
-
-        for param in model.parameters():
-            param.requires_grad = False
-
-        return model

backend/onyx/access/access.py

@@ -43,7 +43,7 @@ def get_access_for_document(
     versioned_get_access_for_document_fn = fetch_versioned_implementation(
         "onyx.access.access", "_get_access_for_document"
     )
-    return versioned_get_access_for_document_fn(document_id, db_session)  # type: ignore
+    return versioned_get_access_for_document_fn(document_id, db_session)
 
 
 def get_null_document_access() -> DocumentAccess:
@@ -93,9 +93,7 @@ def get_access_for_documents(
     versioned_get_access_for_documents_fn = fetch_versioned_implementation(
         "onyx.access.access", "_get_access_for_documents"
     )
-    return versioned_get_access_for_documents_fn(
-        document_ids, db_session
-    )  # type: ignore
+    return versioned_get_access_for_documents_fn(document_ids, db_session)
 
 
 def _get_acl_for_user(user: User | None, db_session: Session) -> set[str]:
@@ -113,7 +111,7 @@ def get_acl_for_user(user: User | None, db_session: Session | None = None) -> se
     versioned_acl_for_user_fn = fetch_versioned_implementation(
         "onyx.access.access", "_get_acl_for_user"
     )
-    return versioned_acl_for_user_fn(user, db_session)  # type: ignore
+    return versioned_acl_for_user_fn(user, db_session)
 
 
 def source_should_fetch_permissions_during_indexing(source: DocumentSource) -> bool:

backend/onyx/agents/agent_framework/message_format.py

@@ -1,73 +0,0 @@
-import json
-from collections.abc import Sequence
-from typing import cast
-
-from langchain_core.messages import AIMessage
-from langchain_core.messages import BaseMessage
-from langchain_core.messages import FunctionMessage
-
-from onyx.llm.message_types import AssistantMessage
-from onyx.llm.message_types import ChatCompletionMessage
-from onyx.llm.message_types import FunctionCall
-from onyx.llm.message_types import SystemMessage
-from onyx.llm.message_types import ToolCall
-from onyx.llm.message_types import ToolMessage
-from onyx.llm.message_types import UserMessageWithText
-
-
-HUMAN = "human"
-SYSTEM = "system"
-AI = "ai"
-FUNCTION = "function"
-
-
-def base_messages_to_chat_completion_msgs(
-    msgs: Sequence[BaseMessage],
-) -> list[ChatCompletionMessage]:
-    return [_base_message_to_chat_completion_msg(msg) for msg in msgs]
-
-
-def _base_message_to_chat_completion_msg(
-    msg: BaseMessage,
-) -> ChatCompletionMessage:
-    if msg.type == HUMAN:
-        content = msg.content if isinstance(msg.content, str) else str(msg.content)
-        user_msg: UserMessageWithText = {"role": "user", "content": content}
-        return user_msg
-    if msg.type == SYSTEM:
-        content = msg.content if isinstance(msg.content, str) else str(msg.content)
-        system_msg: SystemMessage = {"role": "system", "content": content}
-        return system_msg
-    if msg.type == AI:
-        content = msg.content if isinstance(msg.content, str) else str(msg.content)
-        assistant_msg: AssistantMessage = {
-            "role": "assistant",
-            "content": content,
-        }
-        if isinstance(msg, AIMessage) and msg.tool_calls:
-            assistant_msg["tool_calls"] = [
-                ToolCall(
-                    id=tool_call.get("id") or "",
-                    type="function",
-                    function=FunctionCall(
-                        name=tool_call["name"],
-                        arguments=json.dumps(tool_call["args"]),
-                    ),
-                )
-                for tool_call in msg.tool_calls
-            ]
-        return assistant_msg
-    if msg.type == FUNCTION:
-        function_message = cast(FunctionMessage, msg)
-        content = (
-            function_message.content
-            if isinstance(function_message.content, str)
-            else str(function_message.content)
-        )
-        tool_msg: ToolMessage = {
-            "role": "tool",
-            "content": content,
-            "tool_call_id": function_message.name or "",
-        }
-        return tool_msg
-    raise ValueError(f"Unexpected message type: {msg.type}")

backend/onyx/agents/agent_framework/models.py

@@ -1,47 +0,0 @@
-from typing import Any
-from typing import Literal
-from typing import TypeAlias
-
-from pydantic import BaseModel
-
-from onyx.llm.model_response import ModelResponseStream
-
-
-class ToolCallStreamItem(BaseModel):
-    call_id: str | None = None
-
-    id: str | None = None
-
-    name: str | None = None
-
-    arguments: str | None = None
-
-    type: Literal["function_call"] = "function_call"
-
-    index: int | None = None
-
-
-class ToolCallOutputStreamItem(BaseModel):
-    call_id: str | None = None
-
-    output: Any
-
-    type: Literal["function_call_output"] = "function_call_output"
-
-
-RunItemStreamEventDetails: TypeAlias = ToolCallStreamItem | ToolCallOutputStreamItem
-
-
-class RunItemStreamEvent(BaseModel):
-    type: Literal[
-        "message_start",
-        "message_done",
-        "reasoning_start",
-        "reasoning_done",
-        "tool_call",
-        "tool_call_output",
-    ]
-    details: RunItemStreamEventDetails | None = None
-
-
-StreamEvent: TypeAlias = ModelResponseStream | RunItemStreamEvent

backend/onyx/agents/agent_framework/query.py

@@ -1,365 +0,0 @@
-import json
-from collections.abc import Callable
-from collections.abc import Iterator
-from collections.abc import Sequence
-from dataclasses import dataclass
-from typing import Any
-
-import onyx.tracing.framework._error_tracing as _error_tracing
-from onyx.agents.agent_framework.models import RunItemStreamEvent
-from onyx.agents.agent_framework.models import StreamEvent
-from onyx.agents.agent_framework.models import ToolCallOutputStreamItem
-from onyx.agents.agent_framework.models import ToolCallStreamItem
-from onyx.llm.interfaces import LanguageModelInput
-from onyx.llm.interfaces import LLM
-from onyx.llm.interfaces import ToolChoiceOptions
-from onyx.llm.message_types import ChatCompletionMessage
-from onyx.llm.message_types import ToolCall
-from onyx.llm.model_response import ModelResponseStream
-from onyx.tools.tool import RunContextWrapper
-from onyx.tools.tool import Tool
-from onyx.tracing.framework.create import agent_span
-from onyx.tracing.framework.create import function_span
-from onyx.tracing.framework.create import generation_span
-from onyx.tracing.framework.spans import SpanError
-
-
-@dataclass
-class QueryResult:
-    stream: Iterator[StreamEvent]
-    new_messages_stateful: list[ChatCompletionMessage]
-
-
-def _serialize_tool_output(output: Any) -> str:
-    if isinstance(output, str):
-        return output
-    try:
-        return json.dumps(output)
-    except TypeError:
-        return str(output)
-
-
-def _parse_tool_calls_from_message_content(
-    content: str,
-) -> list[dict[str, Any]]:
-    """Parse JSON content that represents tool call instructions."""
-    try:
-        parsed_content = json.loads(content)
-    except json.JSONDecodeError:
-        return []
-
-    if isinstance(parsed_content, dict):
-        candidates = [parsed_content]
-    elif isinstance(parsed_content, list):
-        candidates = [item for item in parsed_content if isinstance(item, dict)]
-    else:
-        return []
-
-    tool_calls: list[dict[str, Any]] = []
-
-    for candidate in candidates:
-        name = candidate.get("name")
-        arguments = candidate.get("arguments")
-
-        if not isinstance(name, str) or arguments is None:
-            continue
-
-        if not isinstance(arguments, dict):
-            continue
-
-        call_id = candidate.get("id")
-        arguments_str = json.dumps(arguments)
-        tool_calls.append(
-            {
-                "id": call_id,
-                "name": name,
-                "arguments": arguments_str,
-            }
-        )
-
-    return tool_calls
-
-
-def _try_convert_content_to_tool_calls_for_non_tool_calling_llms(
-    tool_calls_in_progress: dict[int, dict[str, Any]],
-    content_parts: list[str],
-    structured_response_format: dict | None,
-    next_synthetic_tool_call_id: Callable[[], str],
-) -> None:
-    """Populate tool_calls_in_progress when a non-tool-calling LLM returns JSON content describing tool calls."""
-    if tool_calls_in_progress or not content_parts or structured_response_format:
-        return
-
-    tool_calls_from_content = _parse_tool_calls_from_message_content(
-        "".join(content_parts)
-    )
-
-    if not tool_calls_from_content:
-        return
-
-    content_parts.clear()
-
-    for index, tool_call_data in enumerate(tool_calls_from_content):
-        call_id = tool_call_data["id"] or next_synthetic_tool_call_id()
-        tool_calls_in_progress[index] = {
-            "id": call_id,
-            "name": tool_call_data["name"],
-            "arguments": tool_call_data["arguments"],
-        }
-
-
-def _update_tool_call_with_delta(
-    tool_calls_in_progress: dict[int, dict[str, Any]],
-    tool_call_delta: Any,
-) -> None:
-    index = tool_call_delta.index
-
-    if index not in tool_calls_in_progress:
-        tool_calls_in_progress[index] = {
-            "id": None,
-            "name": None,
-            "arguments": "",
-        }
-
-    if tool_call_delta.id:
-        tool_calls_in_progress[index]["id"] = tool_call_delta.id
-
-    if tool_call_delta.function:
-        if tool_call_delta.function.name:
-            tool_calls_in_progress[index]["name"] = tool_call_delta.function.name
-
-        if tool_call_delta.function.arguments:
-            tool_calls_in_progress[index][
-                "arguments"
-            ] += tool_call_delta.function.arguments
-
-
-def query(
-    llm_with_default_settings: LLM,
-    messages: LanguageModelInput,
-    tools: Sequence[Tool],
-    context: Any,
-    tool_choice: ToolChoiceOptions | None = None,
-    structured_response_format: dict | None = None,
-) -> QueryResult:
-    tool_definitions = [tool.tool_definition() for tool in tools]
-    tools_by_name = {tool.name: tool for tool in tools}
-
-    new_messages_stateful: list[ChatCompletionMessage] = []
-
-    current_span = agent_span(
-        name="agent_framework_query",
-        output_type="dict" if structured_response_format else "str",
-    )
-    current_span.start(mark_as_current=True)
-    current_span.span_data.tools = [t.name for t in tools]
-
-    def stream_generator() -> Iterator[StreamEvent]:
-        message_started = False
-        reasoning_started = False
-
-        tool_calls_in_progress: dict[int, dict[str, Any]] = {}
-
-        content_parts: list[str] = []
-
-        synthetic_tool_call_counter = 0
-
-        def _next_synthetic_tool_call_id() -> str:
-            nonlocal synthetic_tool_call_counter
-            call_id = f"synthetic_tool_call_{synthetic_tool_call_counter}"
-            synthetic_tool_call_counter += 1
-            return call_id
-
-        with generation_span(  # type: ignore[misc]
-            model=llm_with_default_settings.config.model_name,
-            model_config={
-                "base_url": str(llm_with_default_settings.config.api_base or ""),
-                "model_impl": "litellm",
-            },
-        ) as span_generation:
-            # Only set input if messages is a sequence (not a string)
-            # ChatCompletionMessage TypedDicts are compatible with Mapping[str, Any] at runtime
-            if isinstance(messages, Sequence) and not isinstance(messages, str):
-                # Convert ChatCompletionMessage sequence to Sequence[Mapping[str, Any]]
-                span_generation.span_data.input = [dict(msg) for msg in messages]  # type: ignore[assignment]
-            for chunk in llm_with_default_settings.stream(
-                prompt=messages,
-                tools=tool_definitions,
-                tool_choice=tool_choice,
-                structured_response_format=structured_response_format,
-            ):
-                assert isinstance(chunk, ModelResponseStream)
-                usage = getattr(chunk, "usage", None)
-                if usage:
-                    span_generation.span_data.usage = {
-                        "input_tokens": usage.prompt_tokens,
-                        "output_tokens": usage.completion_tokens,
-                        "cache_read_input_tokens": usage.cache_read_input_tokens,
-                        "cache_creation_input_tokens": usage.cache_creation_input_tokens,
-                    }
-
-                delta = chunk.choice.delta
-                finish_reason = chunk.choice.finish_reason
-
-                if delta.reasoning_content:
-                    if not reasoning_started:
-                        yield RunItemStreamEvent(type="reasoning_start")
-                        reasoning_started = True
-
-                if delta.content:
-                    if reasoning_started:
-                        yield RunItemStreamEvent(type="reasoning_done")
-                        reasoning_started = False
-                    content_parts.append(delta.content)
-                    if not message_started:
-                        yield RunItemStreamEvent(type="message_start")
-                        message_started = True
-
-                if delta.tool_calls:
-                    if reasoning_started:
-                        yield RunItemStreamEvent(type="reasoning_done")
-                        reasoning_started = False
-                    if message_started:
-                        yield RunItemStreamEvent(type="message_done")
-                        message_started = False
-
-                    for tool_call_delta in delta.tool_calls:
-                        _update_tool_call_with_delta(
-                            tool_calls_in_progress, tool_call_delta
-                        )
-
-                yield chunk
-
-                if not finish_reason:
-                    continue
-
-                if reasoning_started:
-                    yield RunItemStreamEvent(type="reasoning_done")
-                    reasoning_started = False
-                if message_started:
-                    yield RunItemStreamEvent(type="message_done")
-                    message_started = False
-
-                if tool_choice != "none":
-                    _try_convert_content_to_tool_calls_for_non_tool_calling_llms(
-                        tool_calls_in_progress,
-                        content_parts,
-                        structured_response_format,
-                        _next_synthetic_tool_call_id,
-                    )
-
-                if content_parts:
-                    new_messages_stateful.append(
-                        {
-                            "role": "assistant",
-                            "content": "".join(content_parts),
-                        }
-                    )
-            span_generation.span_data.output = new_messages_stateful
-
-        # Execute tool calls outside of the stream loop and generation_span
-        if tool_calls_in_progress:
-            sorted_tool_calls = sorted(tool_calls_in_progress.items())
-
-            # Build tool calls for the message and execute tools
-            assistant_tool_calls: list[ToolCall] = []
-            tool_outputs: dict[str, str] = {}
-
-            for _, tool_call_data in sorted_tool_calls:
-                call_id = tool_call_data["id"]
-                name = tool_call_data["name"]
-                arguments_str = tool_call_data["arguments"]
-
-                if call_id is None or name is None:
-                    continue
-
-                assistant_tool_calls.append(
-                    {
-                        "id": call_id,
-                        "type": "function",
-                        "function": {
-                            "name": name,
-                            "arguments": arguments_str,
-                        },
-                    }
-                )
-
-                yield RunItemStreamEvent(
-                    type="tool_call",
-                    details=ToolCallStreamItem(
-                        call_id=call_id,
-                        name=name,
-                        arguments=arguments_str,
-                    ),
-                )
-
-                if name in tools_by_name:
-                    tool = tools_by_name[name]
-                    arguments = json.loads(arguments_str)
-
-                    run_context = RunContextWrapper(context=context)
-
-                    # TODO: Instead of executing sequentially, execute in parallel
-                    # In practice, it's not a must right now since we don't use parallel
-                    # tool calls, so kicking the can down the road for now.
-                    with function_span(tool.name) as span_fn:
-                        span_fn.span_data.input = arguments
-                        try:
-                            output = tool.run_v2(run_context, **arguments)
-                            tool_outputs[call_id] = _serialize_tool_output(output)
-                            span_fn.span_data.output = output
-                        except Exception as e:
-                            _error_tracing.attach_error_to_current_span(
-                                SpanError(
-                                    message="Error running tool",
-                                    data={"tool_name": tool.name, "error": str(e)},
-                                )
-                            )
-                            # Treat the error as the tool output so the framework can continue
-                            error_output = f"Error: {str(e)}"
-                            tool_outputs[call_id] = error_output
-                            output = error_output
-
-                    yield RunItemStreamEvent(
-                        type="tool_call_output",
-                        details=ToolCallOutputStreamItem(
-                            call_id=call_id,
-                            output=output,
-                        ),
-                    )
-                else:
-                    not_found_output = f"Tool {name} not found"
-                    tool_outputs[call_id] = _serialize_tool_output(not_found_output)
-                    yield RunItemStreamEvent(
-                        type="tool_call_output",
-                        details=ToolCallOutputStreamItem(
-                            call_id=call_id,
-                            output=not_found_output,
-                        ),
-                    )
-
-            new_messages_stateful.append(
-                {
-                    "role": "assistant",
-                    "content": None,
-                    "tool_calls": assistant_tool_calls,
-                }
-            )
-
-            for _, tool_call_data in sorted_tool_calls:
-                call_id = tool_call_data["id"]
-
-                if call_id in tool_outputs:
-                    new_messages_stateful.append(
-                        {
-                            "role": "tool",
-                            "content": tool_outputs[call_id],
-                            "tool_call_id": call_id,
-                        }
-                    )
-        current_span.finish(reset_current=True)
-
-    return QueryResult(
-        stream=stream_generator(),
-        new_messages_stateful=new_messages_stateful,
-    )

backend/onyx/agents/agent_sdk/message_format.py

@@ -1,167 +0,0 @@
-from collections.abc import Sequence
-
-from langchain.schema.messages import BaseMessage
-
-from onyx.agents.agent_sdk.message_types import AgentSDKMessage
-from onyx.agents.agent_sdk.message_types import AssistantMessageWithContent
-from onyx.agents.agent_sdk.message_types import ImageContent
-from onyx.agents.agent_sdk.message_types import InputTextContent
-from onyx.agents.agent_sdk.message_types import SystemMessage
-from onyx.agents.agent_sdk.message_types import UserMessage
-
-
-# TODO: Currently, we only support native API input for images. For other
-# files, we process the content and share it as text in the message. In
-# the future, we might support native file uploads for other types of files.
-def base_messages_to_agent_sdk_msgs(
-    msgs: Sequence[BaseMessage],
-    is_responses_api: bool,
-) -> list[AgentSDKMessage]:
-    return [_base_message_to_agent_sdk_msg(msg, is_responses_api) for msg in msgs]
-
-
-def _base_message_to_agent_sdk_msg(
-    msg: BaseMessage, is_responses_api: bool
-) -> AgentSDKMessage:
-    message_type_to_agent_sdk_role = {
-        "human": "user",
-        "system": "system",
-        "ai": "assistant",
-    }
-    role = message_type_to_agent_sdk_role[msg.type]
-
-    # Convert content to Agent SDK format
-    content = msg.content
-
-    if isinstance(content, str):
-        # For system/user/assistant messages, use InputTextContent
-        if role in ("system", "user"):
-            input_text_content: list[InputTextContent | ImageContent] = [
-                InputTextContent(type="input_text", text=content)
-            ]
-            if role == "system":
-                # SystemMessage only accepts InputTextContent
-                system_msg: SystemMessage = {
-                    "role": "system",
-                    "content": [InputTextContent(type="input_text", text=content)],
-                }
-                return system_msg
-            else:  # user
-                user_msg: UserMessage = {
-                    "role": "user",
-                    "content": input_text_content,
-                }
-                return user_msg
-        else:  # assistant
-            assistant_msg: AssistantMessageWithContent
-            if is_responses_api:
-                from onyx.agents.agent_sdk.message_types import OutputTextContent
-
-                assistant_msg = {
-                    "role": "assistant",
-                    "content": [OutputTextContent(type="output_text", text=content)],
-                }
-            else:
-                assistant_msg = {
-                    "role": "assistant",
-                    "content": [InputTextContent(type="input_text", text=content)],
-                }
-            return assistant_msg
-    elif isinstance(content, list):
-        # For lists, we need to process based on the role
-        if role == "assistant":
-            # For responses API, use OutputTextContent; otherwise use InputTextContent
-            assistant_content: list[InputTextContent | OutputTextContent] = []
-
-            if is_responses_api:
-                from onyx.agents.agent_sdk.message_types import OutputTextContent
-
-                for item in content:
-                    if isinstance(item, str):
-                        assistant_content.append(
-                            OutputTextContent(type="output_text", text=item)
-                        )
-                    elif isinstance(item, dict) and item.get("type") == "text":
-                        assistant_content.append(
-                            OutputTextContent(
-                                type="output_text", text=item.get("text", "")
-                            )
-                        )
-                    else:
-                        raise ValueError(
-                            f"Unexpected item type for assistant message: {type(item)}. Item: {item}"
-                        )
-            else:
-                for item in content:
-                    if isinstance(item, str):
-                        assistant_content.append(
-                            InputTextContent(type="input_text", text=item)
-                        )
-                    elif isinstance(item, dict) and item.get("type") == "text":
-                        assistant_content.append(
-                            InputTextContent(
-                                type="input_text", text=item.get("text", "")
-                            )
-                        )
-                    else:
-                        raise ValueError(
-                            f"Unexpected item type for assistant message: {type(item)}. Item: {item}"
-                        )
-
-            assistant_msg_list: AssistantMessageWithContent = {
-                "role": "assistant",
-                "content": assistant_content,
-            }
-            return assistant_msg_list
-        else:  # system or user - use InputTextContent
-            input_content: list[InputTextContent | ImageContent] = []
-            for item in content:
-                if isinstance(item, str):
-                    input_content.append(InputTextContent(type="input_text", text=item))
-                elif isinstance(item, dict):
-                    item_type = item.get("type")
-                    if item_type == "text":
-                        input_content.append(
-                            InputTextContent(
-                                type="input_text", text=item.get("text", "")
-                            )
-                        )
-                    elif item_type == "image_url":
-                        # Convert image_url to input_image format
-                        image_url = item.get("image_url", {})
-                        if isinstance(image_url, dict):
-                            url = image_url.get("url", "")
-                        else:
-                            url = image_url
-                        input_content.append(
-                            ImageContent(
-                                type="input_image", image_url=url, detail="auto"
-                            )
-                        )
-                    else:
-                        raise ValueError(f"Unexpected item type: {item_type}")
-                else:
-                    raise ValueError(
-                        f"Unexpected item type: {type(item)}. Item: {item}"
-                    )
-
-            if role == "system":
-                # SystemMessage only accepts InputTextContent (no images)
-                text_only_content = [
-                    c for c in input_content if c["type"] == "input_text"
-                ]
-                system_msg_list: SystemMessage = {
-                    "role": "system",
-                    "content": text_only_content,  # type: ignore[typeddict-item]
-                }
-                return system_msg_list
-            else:  # user
-                user_msg_list: UserMessage = {
-                    "role": "user",
-                    "content": input_content,
-                }
-                return user_msg_list
-    else:
-        raise ValueError(
-            f"Unexpected content type: {type(content)}. Content: {content}"
-        )

backend/onyx/agents/agent_sdk/message_types.py

@@ -1,125 +0,0 @@
-"""Strongly typed message structures for Agent SDK messages."""
-
-from typing import Literal
-from typing import NotRequired
-
-from typing_extensions import TypedDict
-
-
-class InputTextContent(TypedDict):
-    type: Literal["input_text"]
-    text: str
-
-
-class OutputTextContent(TypedDict):
-    type: Literal["output_text"]
-    text: str
-
-
-TextContent = InputTextContent | OutputTextContent
-
-
-class ImageContent(TypedDict):
-    type: Literal["input_image"]
-    image_url: str
-    detail: str
-
-
-# Tool call structures
-class ToolCallFunction(TypedDict):
-    name: str
-    arguments: str
-
-
-class ToolCall(TypedDict):
-    id: str
-    type: Literal["function"]
-    function: ToolCallFunction
-
-
-# Message types
-class SystemMessage(TypedDict):
-    role: Literal["system"]
-    content: list[InputTextContent]  # System messages use input text
-
-
-class UserMessage(TypedDict):
-    role: Literal["user"]
-    content: list[
-        InputTextContent | ImageContent
-    ]  # User messages use input text or images
-
-
-class AssistantMessageWithContent(TypedDict):
-    role: Literal["assistant"]
-    content: list[
-        InputTextContent | OutputTextContent
-    ]  # Assistant messages use output_text for responses API compatibility
-
-
-class AssistantMessageWithToolCalls(TypedDict):
-    role: Literal["assistant"]
-    tool_calls: list[ToolCall]
-
-
-class AssistantMessageDuringAgentRun(TypedDict):
-    role: Literal["assistant"]
-    id: str
-    content: (
-        list[InputTextContent | OutputTextContent] | list[ToolCall]
-    )  # Assistant runtime messages receive output_text from agents SDK for responses API compatibility
-    status: Literal["completed", "failed", "in_progress"]
-    type: Literal["message"]
-
-
-class ToolMessage(TypedDict):
-    role: Literal["tool"]
-    content: str
-    tool_call_id: str
-
-
-class FunctionCallMessage(TypedDict):
-    """Agent SDK function call message format."""
-
-    type: Literal["function_call"]
-    id: NotRequired[str]
-    call_id: str
-    name: str
-    arguments: str
-
-
-class FunctionCallOutputMessage(TypedDict):
-    """Agent SDK function call output message format."""
-
-    type: Literal["function_call_output"]
-    call_id: str
-    output: str
-
-
-class SummaryText(TypedDict):
-    """Summary text item in reasoning messages."""
-
-    text: str
-    type: Literal["summary_text"]
-
-
-class ReasoningMessage(TypedDict):
-    """Agent SDK reasoning message format."""
-
-    id: str
-    type: Literal["reasoning"]
-    summary: list[SummaryText]
-
-
-# Union type for all Agent SDK messages
-AgentSDKMessage = (
-    SystemMessage
-    | UserMessage
-    | AssistantMessageWithContent
-    | AssistantMessageWithToolCalls
-    | AssistantMessageDuringAgentRun
-    | ToolMessage
-    | FunctionCallMessage
-    | FunctionCallOutputMessage
-    | ReasoningMessage
-)

backend/onyx/agents/agent_sdk/monkey_patches.py

@@ -1,36 +0,0 @@
-from typing import Any
-
-from agents.models.openai_responses import Converter as OpenAIResponsesConverter
-
-
-# TODO: I am very sad that I have to monkey patch this :(
-# Basically, OpenAI agents sdk doesn't convert the tool choice correctly
-# when they have a built-in tool in their framework, like they do for web_search
-# and image_generation.
-# Going to open up a thread with OpenAI agents team to see what they recommend
-# or what we can fix.
-# A discussion is warranted, but we likely want to just write our own LitellmModel for
-# the OpenAI agents SDK since they probably don't really care about Litellm and will
-# prioritize functionality for their own models.
-def monkey_patch_convert_tool_choice_to_ignore_openai_hosted_web_search() -> None:
-    if (
-        getattr(OpenAIResponsesConverter.convert_tool_choice, "__name__", "")
-        == "_patched_convert_tool_choice"
-    ):
-        return
-
-    orig_func = OpenAIResponsesConverter.convert_tool_choice.__func__  # type: ignore[attr-defined]
-
-    def _patched_convert_tool_choice(cls: type, tool_choice: Any) -> Any:
-        # Handle OpenAI hosted tools that we have custom implementations for
-        # Without this patch, the library uses special formatting that breaks our custom tools
-        # See: https://platform.openai.com/docs/api-reference/responses/create#responses_create-tool_choice-hosted_tool-type
-        if tool_choice == "web_search":
-            return {"type": "function", "name": "web_search"}
-        if tool_choice == "image_generation":
-            return {"type": "function", "name": "image_generation"}
-        return orig_func(cls, tool_choice)
-
-    OpenAIResponsesConverter.convert_tool_choice = classmethod(  # type: ignore[method-assign, assignment]
-        _patched_convert_tool_choice
-    )

backend/onyx/agents/agent_sdk/sync_agent_stream_adapter.py

@@ -1,178 +0,0 @@
-import asyncio
-import queue
-import threading
-from collections.abc import Iterator
-from collections.abc import Sequence
-from typing import Generic
-from typing import Optional
-from typing import TypeVar
-
-from agents import Agent
-from agents import RunResultStreaming
-from agents import TContext
-from agents.run import Runner
-
-from onyx.agents.agent_sdk.message_types import AgentSDKMessage
-from onyx.utils.threadpool_concurrency import run_in_background
-
-T = TypeVar("T")
-
-
-class SyncAgentStream(Generic[T]):
-    """
-    Convert an async streamed run into a sync iterator with cooperative cancellation.
-    Runs the Agent in a background thread.
-
-    Usage:
-        adapter = SyncStreamAdapter(
-            agent=agent,
-            input=input,
-            context=context,
-            max_turns=100,
-            queue_maxsize=0,         # optional backpressure
-        )
-        for ev in adapter:          # sync iteration
-            ...
-        # or cancel from elsewhere:
-        adapter.cancel()
-    """
-
-    _SENTINEL = object()
-
-    def __init__(
-        self,
-        *,
-        agent: Agent,
-        input: Sequence[AgentSDKMessage],
-        context: TContext | None = None,
-        max_turns: int = 100,
-        queue_maxsize: int = 0,
-    ) -> None:
-        self._agent = agent
-        self._input = input
-        self._context = context
-        self._max_turns = max_turns
-
-        self._q: "queue.Queue[object]" = queue.Queue(maxsize=queue_maxsize)
-        self._loop: Optional[asyncio.AbstractEventLoop] = None
-        self._thread: Optional[threading.Thread] = None
-        self.streamed: RunResultStreaming | None = None
-        self._exc: Optional[BaseException] = None
-        self._cancel_requested = threading.Event()
-        self._started = threading.Event()
-        self._done = threading.Event()
-
-        self._start_thread()
-
-    # ---------- public sync API ----------
-
-    def __iter__(self) -> Iterator[T]:
-        try:
-            while True:
-                item = self._q.get()
-                if item is self._SENTINEL:
-                    # If the consumer thread raised, surface it now
-                    if self._exc is not None:
-                        raise self._exc
-                    # Normal completion
-                    return
-                yield item  # type: ignore[misc,return-value]
-        finally:
-            # Ensure we fully clean up whether we exited due to exception,
-            # StopIteration, or external cancel.
-            self.close()
-
-    def cancel(self) -> bool:
-        """
-        Cooperatively cancel the underlying streamed run and shut down.
-        Safe to call multiple times and from any thread.
-        """
-        self._cancel_requested.set()
-        loop = self._loop
-        streamed = self.streamed
-        if loop is not None and streamed is not None and not self._done.is_set():
-            loop.call_soon_threadsafe(streamed.cancel)
-            return True
-        return False
-
-    def close(self, *, wait: bool = True) -> None:
-        """Idempotent shutdown."""
-        self.cancel()
-        # ask the loop to stop if it's still running
-        loop = self._loop
-        if loop is not None and loop.is_running():
-            try:
-                loop.call_soon_threadsafe(loop.stop)
-            except Exception:
-                pass
-        # join the thread
-        if wait and self._thread is not None and self._thread.is_alive():
-            self._thread.join(timeout=5.0)
-
-    # ---------- internals ----------
-
-    def _start_thread(self) -> None:
-        t = run_in_background(self._thread_main)
-        self._thread = t
-        # Optionally wait until the loop/worker is started so .cancel() is safe soon after init
-        self._started.wait(timeout=1.0)
-
-    def _thread_main(self) -> None:
-        loop = asyncio.new_event_loop()
-        self._loop = loop
-        asyncio.set_event_loop(loop)
-
-        async def worker() -> None:
-            try:
-                # Start the streamed run inside the loop thread
-                self.streamed = Runner.run_streamed(
-                    self._agent,
-                    self._input,  # type: ignore[arg-type]
-                    context=self._context,
-                    max_turns=self._max_turns,
-                )
-
-                # If cancel was requested before we created _streamed, honor it now
-                if self._cancel_requested.is_set():
-                    await self.streamed.cancel()  # type: ignore[func-returns-value]
-
-                # Consume async events and forward into the thread-safe queue
-                async for ev in self.streamed.stream_events():
-                    # Early exit if a late cancel arrives
-                    if self._cancel_requested.is_set():
-                        # Try to cancel gracefully; don't break until cancel takes effect
-                        try:
-                            await self.streamed.cancel()  # type: ignore[func-returns-value]
-                        except Exception:
-                            pass
-                        break
-                    # This put() may block if queue_maxsize > 0 (backpressure)
-                    self._q.put(ev)
-
-            except BaseException as e:
-                # Save exception to surface on the sync iterator side
-                self._exc = e
-            finally:
-                # Signal end-of-stream
-                self._q.put(self._SENTINEL)
-                self._done.set()
-
-        # Mark started and run the worker to completion
-        self._started.set()
-        try:
-            loop.run_until_complete(worker())
-        finally:
-            try:
-                # Drain pending tasks/callbacks safely
-                pending = asyncio.all_tasks(loop=loop)
-                for task in pending:
-                    task.cancel()
-                if pending:
-                    loop.run_until_complete(
-                        asyncio.gather(*pending, return_exceptions=True)
-                    )
-            except Exception:
-                pass
-            finally:
-                loop.close()
-                self._loop = None

backend/onyx/agents/agent_search/core_state.py

@@ -1,21 +0,0 @@
-from operator import add
-from typing import Annotated
-
-from pydantic import BaseModel
-
-
-class CoreState(BaseModel):
-    """
-    This is the core state that is shared across all subgraphs.
-    """
-
-    log_messages: Annotated[list[str], add] = []
-    current_step_nr: int = 1
-
-
-class SubgraphCoreState(BaseModel):
-    """
-    This is the core state that is shared across all subgraphs.
-    """
-
-    log_messages: Annotated[list[str], add] = []

backend/onyx/agents/agent_search/dc_search_analysis/edges.py

@@ -1,62 +0,0 @@
-from collections.abc import Hashable
-from typing import cast
-
-from langchain_core.runnables.config import RunnableConfig
-from langgraph.types import Send
-
-from onyx.agents.agent_search.dc_search_analysis.states import ObjectInformationInput
-from onyx.agents.agent_search.dc_search_analysis.states import (
-    ObjectResearchInformationUpdate,
-)
-from onyx.agents.agent_search.dc_search_analysis.states import ObjectSourceInput
-from onyx.agents.agent_search.dc_search_analysis.states import (
-    SearchSourcesObjectsUpdate,
-)
-from onyx.agents.agent_search.models import GraphConfig
-
-
-def parallel_object_source_research_edge(
-    state: SearchSourcesObjectsUpdate, config: RunnableConfig
-) -> list[Send | Hashable]:
-    """
-    LangGraph edge to parallelize the research for an individual object and source
-    """
-
-    search_objects = state.analysis_objects
-    search_sources = state.analysis_sources
-
-    object_source_combinations = [
-        (object, source) for object in search_objects for source in search_sources
-    ]
-
-    return [
-        Send(
-            "research_object_source",
-            ObjectSourceInput(
-                object_source_combination=object_source_combination,
-                log_messages=[],
-            ),
-        )
-        for object_source_combination in object_source_combinations
-    ]
-
-
-def parallel_object_research_consolidation_edge(
-    state: ObjectResearchInformationUpdate, config: RunnableConfig
-) -> list[Send | Hashable]:
-    """
-    LangGraph edge to parallelize the research for an individual object and source
-    """
-    cast(GraphConfig, config["metadata"]["config"])
-    object_research_information_results = state.object_research_information_results
-
-    return [
-        Send(
-            "consolidate_object_research",
-            ObjectInformationInput(
-                object_information=object_information,
-                log_messages=[],
-            ),
-        )
-        for object_information in object_research_information_results
-    ]

backend/onyx/agents/agent_search/dc_search_analysis/graph_builder.py

@@ -1,103 +0,0 @@
-from langgraph.graph import END
-from langgraph.graph import START
-from langgraph.graph import StateGraph
-
-from onyx.agents.agent_search.dc_search_analysis.edges import (
-    parallel_object_research_consolidation_edge,
-)
-from onyx.agents.agent_search.dc_search_analysis.edges import (
-    parallel_object_source_research_edge,
-)
-from onyx.agents.agent_search.dc_search_analysis.nodes.a1_search_objects import (
-    search_objects,
-)
-from onyx.agents.agent_search.dc_search_analysis.nodes.a2_research_object_source import (
-    research_object_source,
-)
-from onyx.agents.agent_search.dc_search_analysis.nodes.a3_structure_research_by_object import (
-    structure_research_by_object,
-)
-from onyx.agents.agent_search.dc_search_analysis.nodes.a4_consolidate_object_research import (
-    consolidate_object_research,
-)
-from onyx.agents.agent_search.dc_search_analysis.nodes.a5_consolidate_research import (
-    consolidate_research,
-)
-from onyx.agents.agent_search.dc_search_analysis.states import MainInput
-from onyx.agents.agent_search.dc_search_analysis.states import MainState
-from onyx.utils.logger import setup_logger
-
-logger = setup_logger()
-
-test_mode = False
-
-
-def divide_and_conquer_graph_builder(test_mode: bool = False) -> StateGraph:
-    """
-    LangGraph graph builder for the knowledge graph  search process.
-    """
-
-    graph = StateGraph(
-        state_schema=MainState,
-        input=MainInput,
-    )
-
-    ### Add nodes ###
-
-    graph.add_node(
-        "search_objects",
-        search_objects,
-    )
-
-    graph.add_node(
-        "structure_research_by_source",
-        structure_research_by_object,
-    )
-
-    graph.add_node(
-        "research_object_source",
-        research_object_source,
-    )
-
-    graph.add_node(
-        "consolidate_object_research",
-        consolidate_object_research,
-    )
-
-    graph.add_node(
-        "consolidate_research",
-        consolidate_research,
-    )
-
-    ### Add edges ###
-
-    graph.add_edge(start_key=START, end_key="search_objects")
-
-    graph.add_conditional_edges(
-        source="search_objects",
-        path=parallel_object_source_research_edge,
-        path_map=["research_object_source"],
-    )
-
-    graph.add_edge(
-        start_key="research_object_source",
-        end_key="structure_research_by_source",
-    )
-
-    graph.add_conditional_edges(
-        source="structure_research_by_source",
-        path=parallel_object_research_consolidation_edge,
-        path_map=["consolidate_object_research"],
-    )
-
-    graph.add_edge(
-        start_key="consolidate_object_research",
-        end_key="consolidate_research",
-    )
-
-    graph.add_edge(
-        start_key="consolidate_research",
-        end_key=END,
-    )
-
-    return graph

backend/onyx/agents/agent_search/dc_search_analysis/nodes/a1_search_objects.py

@@ -1,146 +0,0 @@
-from typing import cast
-
-from langchain_core.messages import HumanMessage
-from langchain_core.runnables import RunnableConfig
-from langgraph.types import StreamWriter
-
-from onyx.agents.agent_search.dc_search_analysis.ops import extract_section
-from onyx.agents.agent_search.dc_search_analysis.ops import research
-from onyx.agents.agent_search.dc_search_analysis.states import MainState
-from onyx.agents.agent_search.dc_search_analysis.states import (
-    SearchSourcesObjectsUpdate,
-)
-from onyx.agents.agent_search.models import GraphConfig
-from onyx.agents.agent_search.shared_graph_utils.agent_prompt_ops import (
-    trim_prompt_piece,
-)
-from onyx.prompts.agents.dc_prompts import DC_OBJECT_NO_BASE_DATA_EXTRACTION_PROMPT
-from onyx.prompts.agents.dc_prompts import DC_OBJECT_SEPARATOR
-from onyx.prompts.agents.dc_prompts import DC_OBJECT_WITH_BASE_DATA_EXTRACTION_PROMPT
-from onyx.secondary_llm_flows.source_filter import strings_to_document_sources
-from onyx.utils.logger import setup_logger
-from onyx.utils.threadpool_concurrency import run_with_timeout
-
-logger = setup_logger()
-
-
-def search_objects(
-    state: MainState, config: RunnableConfig, writer: StreamWriter = lambda _: None
-) -> SearchSourcesObjectsUpdate:
-    """
-    LangGraph node to start the agentic search process.
-    """
-
-    graph_config = cast(GraphConfig, config["metadata"]["config"])
-    question = graph_config.inputs.prompt_builder.raw_user_query
-    search_tool = graph_config.tooling.search_tool
-
-    if search_tool is None or graph_config.inputs.persona is None:
-        raise ValueError("Search tool and persona must be provided for DivCon search")
-
-    try:
-        instructions = graph_config.inputs.persona.system_prompt or ""
-
-        agent_1_instructions = extract_section(
-            instructions, "Agent Step 1:", "Agent Step 2:"
-        )
-        if agent_1_instructions is None:
-            raise ValueError("Agent 1 instructions not found")
-
-        agent_1_base_data = extract_section(instructions, "|Start Data|", "|End Data|")
-
-        agent_1_task = extract_section(
-            agent_1_instructions, "Task:", "Independent Research Sources:"
-        )
-        if agent_1_task is None:
-            raise ValueError("Agent 1 task not found")
-
-        agent_1_independent_sources_str = extract_section(
-            agent_1_instructions, "Independent Research Sources:", "Output Objective:"
-        )
-        if agent_1_independent_sources_str is None:
-            raise ValueError("Agent 1 Independent Research Sources not found")
-
-        document_sources = strings_to_document_sources(
-            [
-                x.strip().lower()
-                for x in agent_1_independent_sources_str.split(DC_OBJECT_SEPARATOR)
-            ]
-        )
-
-        agent_1_output_objective = extract_section(
-            agent_1_instructions, "Output Objective:"
-        )
-        if agent_1_output_objective is None:
-            raise ValueError("Agent 1 output objective not found")
-
-    except Exception as e:
-        raise ValueError(
-            f"Agent 1 instructions not found or not formatted correctly: {e}"
-        )
-
-    # Extract objects
-
-    if agent_1_base_data is None:
-        # Retrieve chunks for objects
-
-        retrieved_docs = research(question, search_tool)[:10]
-
-        document_texts_list = []
-        for doc_num, doc in enumerate(retrieved_docs):
-            chunk_text = "Document " + str(doc_num) + ":\n" + doc.content
-            document_texts_list.append(chunk_text)
-
-        document_texts = "\n\n".join(document_texts_list)
-
-        dc_object_extraction_prompt = DC_OBJECT_NO_BASE_DATA_EXTRACTION_PROMPT.format(
-            question=question,
-            task=agent_1_task,
-            document_text=document_texts,
-            objects_of_interest=agent_1_output_objective,
-        )
-    else:
-        dc_object_extraction_prompt = DC_OBJECT_WITH_BASE_DATA_EXTRACTION_PROMPT.format(
-            question=question,
-            task=agent_1_task,
-            base_data=agent_1_base_data,
-            objects_of_interest=agent_1_output_objective,
-        )
-
-    msg = [
-        HumanMessage(
-            content=trim_prompt_piece(
-                config=graph_config.tooling.primary_llm.config,
-                prompt_piece=dc_object_extraction_prompt,
-                reserved_str="",
-            ),
-        )
-    ]
-    primary_llm = graph_config.tooling.primary_llm
-    # Grader
-    try:
-        llm_response = run_with_timeout(
-            30,
-            primary_llm.invoke_langchain,
-            prompt=msg,
-            timeout_override=30,
-            max_tokens=300,
-        )
-
-        cleaned_response = (
-            str(llm_response.content)
-            .replace("```json\n", "")
-            .replace("\n```", "")
-            .replace("\n", "")
-        )
-        cleaned_response = cleaned_response.split("OBJECTS:")[1]
-        object_list = [x.strip() for x in cleaned_response.split(";")]
-
-    except Exception as e:
-        raise ValueError(f"Error in search_objects: {e}")
-
-    return SearchSourcesObjectsUpdate(
-        analysis_objects=object_list,
-        analysis_sources=document_sources,
-        log_messages=["Agent 1 Task done"],
-    )

backend/onyx/agents/agent_search/dc_search_analysis/nodes/a2_research_object_source.py

@@ -1,180 +0,0 @@
-from datetime import datetime
-from datetime import timedelta
-from datetime import timezone
-from typing import cast
-
-from langchain_core.messages import HumanMessage
-from langchain_core.runnables import RunnableConfig
-from langgraph.types import StreamWriter
-
-from onyx.agents.agent_search.dc_search_analysis.ops import extract_section
-from onyx.agents.agent_search.dc_search_analysis.ops import research
-from onyx.agents.agent_search.dc_search_analysis.states import ObjectSourceInput
-from onyx.agents.agent_search.dc_search_analysis.states import (
-    ObjectSourceResearchUpdate,
-)
-from onyx.agents.agent_search.models import GraphConfig
-from onyx.agents.agent_search.shared_graph_utils.agent_prompt_ops import (
-    trim_prompt_piece,
-)
-from onyx.prompts.agents.dc_prompts import DC_OBJECT_SOURCE_RESEARCH_PROMPT
-from onyx.utils.logger import setup_logger
-from onyx.utils.threadpool_concurrency import run_with_timeout
-
-logger = setup_logger()
-
-
-def research_object_source(
-    state: ObjectSourceInput,
-    config: RunnableConfig,
-    writer: StreamWriter = lambda _: None,
-) -> ObjectSourceResearchUpdate:
-    """
-    LangGraph node to start the agentic search process.
-    """
-    datetime.now()
-
-    graph_config = cast(GraphConfig, config["metadata"]["config"])
-    search_tool = graph_config.tooling.search_tool
-    question = graph_config.inputs.prompt_builder.raw_user_query
-    object, document_source = state.object_source_combination
-
-    if search_tool is None or graph_config.inputs.persona is None:
-        raise ValueError("Search tool and persona must be provided for DivCon search")
-
-    try:
-        instructions = graph_config.inputs.persona.system_prompt or ""
-
-        agent_2_instructions = extract_section(
-            instructions, "Agent Step 2:", "Agent Step 3:"
-        )
-        if agent_2_instructions is None:
-            raise ValueError("Agent 2 instructions not found")
-
-        agent_2_task = extract_section(
-            agent_2_instructions, "Task:", "Independent Research Sources:"
-        )
-        if agent_2_task is None:
-            raise ValueError("Agent 2 task not found")
-
-        agent_2_time_cutoff = extract_section(
-            agent_2_instructions, "Time Cutoff:", "Research Topics:"
-        )
-
-        agent_2_research_topics = extract_section(
-            agent_2_instructions, "Research Topics:", "Output Objective"
-        )
-
-        agent_2_output_objective = extract_section(
-            agent_2_instructions, "Output Objective:"
-        )
-        if agent_2_output_objective is None:
-            raise ValueError("Agent 2 output objective not found")
-
-    except Exception:
-        raise ValueError(
-            "Agent 1 instructions not found or not formatted correctly: {e}"
-        )
-
-    # Populate prompt
-
-    # Retrieve chunks for objects
-
-    if agent_2_time_cutoff is not None and agent_2_time_cutoff.strip() != "":
-        if agent_2_time_cutoff.strip().endswith("d"):
-            try:
-                days = int(agent_2_time_cutoff.strip()[:-1])
-                agent_2_source_start_time = datetime.now(timezone.utc) - timedelta(
-                    days=days
-                )
-            except ValueError:
-                raise ValueError(
-                    f"Invalid time cutoff format: {agent_2_time_cutoff}. Expected format: '<number>d'"
-                )
-        else:
-            raise ValueError(
-                f"Invalid time cutoff format: {agent_2_time_cutoff}. Expected format: '<number>d'"
-            )
-    else:
-        agent_2_source_start_time = None
-
-    document_sources = [document_source] if document_source else None
-
-    if len(question.strip()) > 0:
-        research_area = f"{question} for {object}"
-    elif agent_2_research_topics and len(agent_2_research_topics.strip()) > 0:
-        research_area = f"{agent_2_research_topics} for {object}"
-    else:
-        research_area = object
-
-    retrieved_docs = research(
-        question=research_area,
-        search_tool=search_tool,
-        document_sources=document_sources,
-        time_cutoff=agent_2_source_start_time,
-    )
-
-    # Generate document text
-
-    document_texts_list = []
-    for doc_num, doc in enumerate(retrieved_docs):
-        chunk_text = "Document " + str(doc_num) + ":\n" + doc.content
-        document_texts_list.append(chunk_text)
-
-    document_texts = "\n\n".join(document_texts_list)
-
-    # Built prompt
-
-    today = datetime.now().strftime("%A, %Y-%m-%d")
-
-    dc_object_source_research_prompt = (
-        DC_OBJECT_SOURCE_RESEARCH_PROMPT.format(
-            today=today,
-            question=question,
-            task=agent_2_task,
-            document_text=document_texts,
-            format=agent_2_output_objective,
-        )
-        .replace("---object---", object)
-        .replace("---source---", document_source.value)
-    )
-
-    # Run LLM
-
-    msg = [
-        HumanMessage(
-            content=trim_prompt_piece(
-                config=graph_config.tooling.primary_llm.config,
-                prompt_piece=dc_object_source_research_prompt,
-                reserved_str="",
-            ),
-        )
-    ]
-    primary_llm = graph_config.tooling.primary_llm
-    # Grader
-    try:
-        llm_response = run_with_timeout(
-            30,
-            primary_llm.invoke_langchain,
-            prompt=msg,
-            timeout_override=30,
-            max_tokens=300,
-        )
-
-        cleaned_response = str(llm_response.content).replace("```json\n", "")
-        cleaned_response = cleaned_response.split("RESEARCH RESULTS:")[1]
-        object_research_results = {
-            "object": object,
-            "source": document_source.value,
-            "research_result": cleaned_response,
-        }
-
-    except Exception as e:
-        raise ValueError(f"Error in research_object_source: {e}")
-
-    logger.debug("DivCon Step A2 - Object Source Research - completed for an object")
-
-    return ObjectSourceResearchUpdate(
-        object_source_research_results=[object_research_results],
-        log_messages=["Agent Step 2 done for one object"],
-    )