review

.github/workflows/release-cli.yml

@@ -26,7 +26,7 @@ jobs:
       - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6
         with:
           persist-credentials: false
-      - uses: astral-sh/setup-uv@5a095e7a2014a4212f075830d4f7277575a9d098 # ratchet:astral-sh/setup-uv@v7
+      - uses: astral-sh/setup-uv@61cb8a9741eeb8a550a1b8544337180c0fc8476b # ratchet:astral-sh/setup-uv@v7
         with:
           enable-cache: false
           version: "0.9.9"
@@ -37,178 +37,3 @@ jobs:
         working-directory: cli
       - run: uv publish
         working-directory: cli
-
-  docker-amd64:
-    runs-on:
-      - runs-on
-      - runner=2cpu-linux-x64
-      - run-id=${{ github.run_id }}-cli-amd64
-      - extras=ecr-cache
-    environment: deploy
-    permissions:
-      id-token: write
-    timeout-minutes: 30
-    outputs:
-      digest: ${{ steps.build.outputs.digest }}
-    env:
-      REGISTRY_IMAGE: onyxdotapp/onyx-cli
-    steps:
-      - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2
-
-      - name: Checkout
-        uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6
-        with:
-          persist-credentials: false
-
-      - name: Configure AWS credentials
-        uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7 # ratchet:aws-actions/configure-aws-credentials@v6.0.0
-        with:
-          role-to-assume: ${{ secrets.AWS_OIDC_ROLE_ARN }}
-          aws-region: us-east-2
-
-      - name: Get AWS Secrets
-        uses: aws-actions/aws-secretsmanager-get-secrets@a9a7eb4e2f2871d30dc5b892576fde60a2ecc802 # ratchet:aws-actions/aws-secretsmanager-get-secrets@v2.0.10
-        with:
-          secret-ids: |
-            DOCKER_USERNAME, deploy/docker-username
-            DOCKER_TOKEN, deploy/docker-token
-          parse-json-secrets: true
-
-      - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@4d04d5d9486b7bd6fa91e7baf45bbb4f8b9deedd # ratchet:docker/setup-buildx-action@v4
-
-      - name: Login to Docker Hub
-        uses: docker/login-action@b45d80f862d83dbcd57f89517bcf500b2ab88fb2 # ratchet:docker/login-action@v4
-        with:
-          username: ${{ env.DOCKER_USERNAME }}
-          password: ${{ env.DOCKER_TOKEN }}
-
-      - name: Build and push AMD64
-        id: build
-        uses: docker/build-push-action@d08e5c354a6adb9ed34480a06d141179aa583294 # ratchet:docker/build-push-action@v7
-        with:
-          context: ./cli
-          file: ./cli/Dockerfile
-          platforms: linux/amd64
-          cache-from: type=registry,ref=${{ env.REGISTRY_IMAGE }}:latest
-          cache-to: type=inline
-          outputs: type=image,name=${{ env.REGISTRY_IMAGE }},push-by-digest=true,name-canonical=true,push=true
-
-  docker-arm64:
-    runs-on:
-      - runs-on
-      - runner=2cpu-linux-arm64
-      - run-id=${{ github.run_id }}-cli-arm64
-      - extras=ecr-cache
-    environment: deploy
-    permissions:
-      id-token: write
-    timeout-minutes: 30
-    outputs:
-      digest: ${{ steps.build.outputs.digest }}
-    env:
-      REGISTRY_IMAGE: onyxdotapp/onyx-cli
-    steps:
-      - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2
-
-      - name: Checkout
-        uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6
-        with:
-          persist-credentials: false
-
-      - name: Configure AWS credentials
-        uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7 # ratchet:aws-actions/configure-aws-credentials@v6.0.0
-        with:
-          role-to-assume: ${{ secrets.AWS_OIDC_ROLE_ARN }}
-          aws-region: us-east-2
-
-      - name: Get AWS Secrets
-        uses: aws-actions/aws-secretsmanager-get-secrets@a9a7eb4e2f2871d30dc5b892576fde60a2ecc802 # ratchet:aws-actions/aws-secretsmanager-get-secrets@v2.0.10
-        with:
-          secret-ids: |
-            DOCKER_USERNAME, deploy/docker-username
-            DOCKER_TOKEN, deploy/docker-token
-          parse-json-secrets: true
-
-      - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@4d04d5d9486b7bd6fa91e7baf45bbb4f8b9deedd # ratchet:docker/setup-buildx-action@v4
-
-      - name: Login to Docker Hub
-        uses: docker/login-action@b45d80f862d83dbcd57f89517bcf500b2ab88fb2 # ratchet:docker/login-action@v4
-        with:
-          username: ${{ env.DOCKER_USERNAME }}
-          password: ${{ env.DOCKER_TOKEN }}
-
-      - name: Build and push ARM64
-        id: build
-        uses: docker/build-push-action@d08e5c354a6adb9ed34480a06d141179aa583294 # ratchet:docker/build-push-action@v7
-        with:
-          context: ./cli
-          file: ./cli/Dockerfile
-          platforms: linux/arm64
-          cache-from: type=registry,ref=${{ env.REGISTRY_IMAGE }}:latest
-          cache-to: type=inline
-          outputs: type=image,name=${{ env.REGISTRY_IMAGE }},push-by-digest=true,name-canonical=true,push=true
-
-  merge-docker:
-    needs:
-      - docker-amd64
-      - docker-arm64
-    runs-on:
-      - runs-on
-      - runner=2cpu-linux-x64
-      - run-id=${{ github.run_id }}-cli-merge
-    environment: deploy
-    permissions:
-      id-token: write
-    timeout-minutes: 10
-    env:
-      REGISTRY_IMAGE: onyxdotapp/onyx-cli
-    steps:
-      - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2
-
-      - name: Configure AWS credentials
-        uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7 # ratchet:aws-actions/configure-aws-credentials@v6.0.0
-        with:
-          role-to-assume: ${{ secrets.AWS_OIDC_ROLE_ARN }}
-          aws-region: us-east-2
-
-      - name: Get AWS Secrets
-        uses: aws-actions/aws-secretsmanager-get-secrets@a9a7eb4e2f2871d30dc5b892576fde60a2ecc802 # ratchet:aws-actions/aws-secretsmanager-get-secrets@v2.0.10
-        with:
-          secret-ids: |
-            DOCKER_USERNAME, deploy/docker-username
-            DOCKER_TOKEN, deploy/docker-token
-          parse-json-secrets: true
-
-      - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@4d04d5d9486b7bd6fa91e7baf45bbb4f8b9deedd # ratchet:docker/setup-buildx-action@v4
-
-      - name: Login to Docker Hub
-        uses: docker/login-action@b45d80f862d83dbcd57f89517bcf500b2ab88fb2 # ratchet:docker/login-action@v4
-        with:
-          username: ${{ env.DOCKER_USERNAME }}
-          password: ${{ env.DOCKER_TOKEN }}
-
-      - name: Create and push manifest
-        env:
-          AMD64_DIGEST: ${{ needs.docker-amd64.outputs.digest }}
-          ARM64_DIGEST: ${{ needs.docker-arm64.outputs.digest }}
-          TAG: ${{ github.ref_name }}
-        run: |
-          SANITIZED_TAG="${TAG#cli/}"
-          IMAGES=(
-            "${REGISTRY_IMAGE}@${AMD64_DIGEST}"
-            "${REGISTRY_IMAGE}@${ARM64_DIGEST}"
-          )
-
-          if [[ "$TAG" =~ ^cli/v[0-9]+\.[0-9]+\.[0-9]+$ ]]; then
-            docker buildx imagetools create \
-              -t "${REGISTRY_IMAGE}:${SANITIZED_TAG}" \
-              -t "${REGISTRY_IMAGE}:latest" \
-              "${IMAGES[@]}"
-          else
-            docker buildx imagetools create \
-              -t "${REGISTRY_IMAGE}:${SANITIZED_TAG}" \
-              "${IMAGES[@]}"
-          fi

AGENTS.md

@@ -544,8 +544,6 @@ To run them:
 npx playwright test <TEST_NAME>
 ```
 
-For shared fixtures, best practices, and detailed guidance, see `backend/tests/README.md`.
-
 ## Logs
 
 When (1) writing integration tests or (2) doing live tests (e.g. curl / playwright) you can get access

backend/Dockerfile

@@ -141,7 +141,6 @@ COPY --chown=onyx:onyx ./scripts/debugging /app/scripts/debugging
 COPY --chown=onyx:onyx ./scripts/force_delete_connector_by_id.py /app/scripts/force_delete_connector_by_id.py
 COPY --chown=onyx:onyx ./scripts/supervisord_entrypoint.sh /app/scripts/supervisord_entrypoint.sh
 COPY --chown=onyx:onyx ./scripts/setup_craft_templates.sh /app/scripts/setup_craft_templates.sh
-COPY --chown=onyx:onyx ./scripts/reencrypt_secrets.py /app/scripts/reencrypt_secrets.py
 RUN chmod +x /app/scripts/supervisord_entrypoint.sh /app/scripts/setup_craft_templates.sh
 
 # Run Craft template setup at build time when ENABLE_CRAFT=true

backend/ee/onyx/utils/encryption.py

@@ -14,91 +14,67 @@ from onyx.utils.variable_functionality import fetch_versioned_implementation
 logger = setup_logger()
 
 
-@lru_cache(maxsize=2)
+@lru_cache(maxsize=1)
 def _get_trimmed_key(key: str) -> bytes:
     encoded_key = key.encode()
     key_length = len(encoded_key)
     if key_length < 16:
         raise RuntimeError("Invalid ENCRYPTION_KEY_SECRET - too short")
+    elif key_length > 32:
+        key = key[:32]
+    elif key_length not in (16, 24, 32):
+        valid_lengths = [16, 24, 32]
+        key = key[: min(valid_lengths, key=lambda x: abs(x - key_length))]
 
-    # Trim to the largest valid AES key size that fits
-    valid_lengths = [32, 24, 16]
-    for size in valid_lengths:
-        if key_length >= size:
-            return encoded_key[:size]
-
-    raise AssertionError("unreachable")
+    return encoded_key
 
 
-def _encrypt_string(input_str: str, key: str | None = None) -> bytes:
-    effective_key = key if key is not None else ENCRYPTION_KEY_SECRET
-    if not effective_key:
+def _encrypt_string(input_str: str) -> bytes:
+    if not ENCRYPTION_KEY_SECRET:
         return input_str.encode()
 
-    trimmed = _get_trimmed_key(effective_key)
+    key = _get_trimmed_key(ENCRYPTION_KEY_SECRET)
     iv = urandom(16)
     padder = padding.PKCS7(algorithms.AES.block_size).padder()
     padded_data = padder.update(input_str.encode()) + padder.finalize()
 
-    cipher = Cipher(algorithms.AES(trimmed), modes.CBC(iv), backend=default_backend())
+    cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=default_backend())
     encryptor = cipher.encryptor()
     encrypted_data = encryptor.update(padded_data) + encryptor.finalize()
 
     return iv + encrypted_data
 
 
-def _decrypt_bytes(input_bytes: bytes, key: str | None = None) -> str:
-    effective_key = key if key is not None else ENCRYPTION_KEY_SECRET
-    if not effective_key:
+def _decrypt_bytes(input_bytes: bytes) -> str:
+    if not ENCRYPTION_KEY_SECRET:
         return input_bytes.decode()
 
-    trimmed = _get_trimmed_key(effective_key)
-    try:
-        iv = input_bytes[:16]
-        encrypted_data = input_bytes[16:]
+    key = _get_trimmed_key(ENCRYPTION_KEY_SECRET)
+    iv = input_bytes[:16]
+    encrypted_data = input_bytes[16:]
 
-        cipher = Cipher(
-            algorithms.AES(trimmed), modes.CBC(iv), backend=default_backend()
-        )
-        decryptor = cipher.decryptor()
-        decrypted_padded_data = decryptor.update(encrypted_data) + decryptor.finalize()
+    cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=default_backend())
+    decryptor = cipher.decryptor()
+    decrypted_padded_data = decryptor.update(encrypted_data) + decryptor.finalize()
 
-        unpadder = padding.PKCS7(algorithms.AES.block_size).unpadder()
-        decrypted_data = unpadder.update(decrypted_padded_data) + unpadder.finalize()
+    unpadder = padding.PKCS7(algorithms.AES.block_size).unpadder()
+    decrypted_data = unpadder.update(decrypted_padded_data) + unpadder.finalize()
 
-        return decrypted_data.decode()
-    except (ValueError, UnicodeDecodeError):
-        if key is not None:
-            # Explicit key was provided — don't fall back silently
-            raise
-        # Read path: attempt raw UTF-8 decode as a fallback for legacy data.
-        # Does NOT handle data encrypted with a different key — that
-        # ciphertext is not valid UTF-8 and will raise below.
-        logger.warning(
-            "AES decryption failed — falling back to raw decode. "
-            "Run the re-encrypt secrets script to rotate to the current key."
-        )
-        try:
-            return input_bytes.decode()
-        except UnicodeDecodeError:
-            raise ValueError(
-                "Data is not valid UTF-8 — likely encrypted with a different key. "
-                "Run the re-encrypt secrets script to rotate to the current key."
-            ) from None
+    return decrypted_data.decode()
 
 
-def encrypt_string_to_bytes(input_str: str, key: str | None = None) -> bytes:
+def encrypt_string_to_bytes(input_str: str) -> bytes:
     versioned_encryption_fn = fetch_versioned_implementation(
         "onyx.utils.encryption", "_encrypt_string"
     )
-    return versioned_encryption_fn(input_str, key=key)
+    return versioned_encryption_fn(input_str)
 
 
-def decrypt_bytes_to_string(input_bytes: bytes, key: str | None = None) -> str:
+def decrypt_bytes_to_string(input_bytes: bytes) -> str:
     versioned_decryption_fn = fetch_versioned_implementation(
         "onyx.utils.encryption", "_decrypt_bytes"
     )
-    return versioned_decryption_fn(input_bytes, key=key)
+    return versioned_decryption_fn(input_bytes)
 
 
 def test_encryption() -> None:

backend/onyx/chat/llm_step.py

@@ -15,7 +15,6 @@ from onyx.chat.citation_processor import DynamicCitationProcessor
 from onyx.chat.emitter import Emitter
 from onyx.chat.models import ChatMessageSimple
 from onyx.chat.models import LlmStepResult
-from onyx.chat.tool_call_args_streaming import maybe_emit_argument_delta
 from onyx.configs.app_configs import LOG_ONYX_MODEL_INTERACTIONS
 from onyx.configs.app_configs import PROMPT_CACHE_CHAT_HISTORY
 from onyx.configs.constants import MessageType
@@ -55,7 +54,6 @@ from onyx.server.query_and_chat.streaming_models import ReasoningStart
 from onyx.tools.models import ToolCallKickoff
 from onyx.tracing.framework.create import generation_span
 from onyx.utils.b64 import get_image_type_from_bytes
-from onyx.utils.jsonriver import Parser
 from onyx.utils.logger import setup_logger
 from onyx.utils.postgres_sanitization import sanitize_string
 from onyx.utils.text_processing import find_all_json_objects
@@ -1011,7 +1009,6 @@ def run_llm_step_pkt_generator(
         )
 
     id_to_tool_call_map: dict[int, dict[str, Any]] = {}
-    arg_parsers: dict[int, Parser] = {}
     reasoning_start = False
     answer_start = False
     accumulated_reasoning = ""
@@ -1218,14 +1215,7 @@ def run_llm_step_pkt_generator(
                 yield from _close_reasoning_if_active()
 
                 for tool_call_delta in delta.tool_calls:
-                    # maybe_emit depends and update being called first and attaching the delta
                     _update_tool_call_with_delta(id_to_tool_call_map, tool_call_delta)
-                    yield from maybe_emit_argument_delta(
-                        tool_calls_in_progress=id_to_tool_call_map,
-                        tool_call_delta=tool_call_delta,
-                        placement=_current_placement(),
-                        parsers=arg_parsers,
-                    )
 
         # Flush any tail text buffered while checking for split "<function_calls" markers.
         filtered_content_tail = xml_tool_call_content_filter.flush()

backend/onyx/chat/tool_call_args_streaming.py

@@ -1,77 +0,0 @@
-from collections.abc import Generator
-from collections.abc import Mapping
-from typing import Any
-from typing import Type
-
-from onyx.llm.model_response import ChatCompletionDeltaToolCall
-from onyx.server.query_and_chat.placement import Placement
-from onyx.server.query_and_chat.streaming_models import Packet
-from onyx.server.query_and_chat.streaming_models import ToolCallArgumentDelta
-from onyx.tools.built_in_tools import TOOL_NAME_TO_CLASS
-from onyx.tools.interface import Tool
-from onyx.utils.jsonriver import Parser
-
-
-def _get_tool_class(
-    tool_calls_in_progress: Mapping[int, Mapping[str, Any]],
-    tool_call_delta: ChatCompletionDeltaToolCall,
-) -> Type[Tool] | None:
-    """Look up the Tool subclass for a streaming tool call delta."""
-    tool_name = tool_calls_in_progress.get(tool_call_delta.index, {}).get("name")
-    if not tool_name:
-        return None
-    return TOOL_NAME_TO_CLASS.get(tool_name)
-
-
-def maybe_emit_argument_delta(
-    tool_calls_in_progress: Mapping[int, Mapping[str, Any]],
-    tool_call_delta: ChatCompletionDeltaToolCall,
-    placement: Placement,
-    parsers: dict[int, Parser],
-) -> Generator[Packet, None, None]:
-    """Emit decoded tool-call argument deltas to the frontend.
-
-    Uses a ``jsonriver.Parser`` per tool-call index to incrementally parse
-    the JSON argument string and extract only the newly-appended content
-    for each string-valued argument.
-
-    NOTE: Non-string arguments (numbers, booleans, null, arrays, objects)
-    are skipped — they are available in the final tool-call kickoff packet.
-
-    ``parsers`` is a mutable dict keyed by tool-call index. A new
-    ``Parser`` is created automatically for each new index.
-    """
-    tool_cls = _get_tool_class(tool_calls_in_progress, tool_call_delta)
-    if not tool_cls or not tool_cls.should_emit_argument_deltas():
-        return
-
-    fn = tool_call_delta.function
-    delta_fragment = fn.arguments if fn else None
-    if not delta_fragment:
-        return
-
-    idx = tool_call_delta.index
-    if idx not in parsers:
-        parsers[idx] = Parser()
-    parser = parsers[idx]
-
-    deltas = parser.feed(delta_fragment)
-
-    argument_deltas: dict[str, str] = {}
-    for delta in deltas:
-        if isinstance(delta, dict):
-            for key, value in delta.items():
-                if isinstance(value, str):
-                    argument_deltas[key] = argument_deltas.get(key, "") + value
-
-    if not argument_deltas:
-        return
-
-    tc_data = tool_calls_in_progress[tool_call_delta.index]
-    yield Packet(
-        placement=placement,
-        obj=ToolCallArgumentDelta(
-            tool_type=tc_data.get("name", ""),
-            argument_deltas=argument_deltas,
-        ),
-    )

backend/onyx/db/rotate_encryption_key.py

@@ -1,161 +0,0 @@
-"""Rotate encryption key for all encrypted columns.
-
-Dynamically discovers all columns using EncryptedString / EncryptedJson,
-decrypts each value with the old key, and re-encrypts with the current
-ENCRYPTION_KEY_SECRET.
-
-The operation is idempotent: rows already encrypted with the current key
-are skipped. Commits are made in batches so a crash mid-rotation can be
-safely resumed by re-running.
-"""
-
-import json
-from typing import Any
-
-from sqlalchemy import LargeBinary
-from sqlalchemy import select
-from sqlalchemy import update
-from sqlalchemy.orm import Session
-
-from onyx.configs.app_configs import ENCRYPTION_KEY_SECRET
-from onyx.db.models import Base
-from onyx.db.models import EncryptedJson
-from onyx.db.models import EncryptedString
-from onyx.utils.encryption import decrypt_bytes_to_string
-from onyx.utils.logger import setup_logger
-from onyx.utils.variable_functionality import global_version
-
-logger = setup_logger()
-
-_BATCH_SIZE = 500
-
-
-def _can_decrypt_with_current_key(data: bytes) -> bool:
-    """Check if data is already encrypted with the current key.
-
-    Passes the key explicitly so the fallback-to-raw-decode path in
-    _decrypt_bytes is NOT triggered — a clean success/failure signal.
-    """
-    try:
-        decrypt_bytes_to_string(data, key=ENCRYPTION_KEY_SECRET)
-        return True
-    except Exception:
-        return False
-
-
-def _discover_encrypted_columns() -> list[tuple[type, str, list[str], bool]]:
-    """Walk all ORM models and find columns using EncryptedString/EncryptedJson.
-
-    Returns list of (ModelClass, column_attr_name, [pk_attr_names], is_json).
-    """
-    results: list[tuple[type, str, list[str], bool]] = []
-
-    for mapper in Base.registry.mappers:
-        model_cls = mapper.class_
-        pk_names = [col.key for col in mapper.primary_key]
-
-        for prop in mapper.column_attrs:
-            for col in prop.columns:
-                if isinstance(col.type, EncryptedJson):
-                    results.append((model_cls, prop.key, pk_names, True))
-                elif isinstance(col.type, EncryptedString):
-                    results.append((model_cls, prop.key, pk_names, False))
-
-    return results
-
-
-def rotate_encryption_key(
-    db_session: Session,
-    old_key: str | None,
-    dry_run: bool = False,
-) -> dict[str, int]:
-    """Decrypt all encrypted columns with old_key and re-encrypt with the current key.
-
-    Args:
-        db_session: Active database session.
-        old_key: The previous encryption key. Pass None or "" if values were
-                 not previously encrypted with a key.
-        dry_run: If True, count rows that need rotation without modifying data.
-
-    Returns:
-        Dict of "table.column" -> number of rows re-encrypted (or would be).
-
-    Commits every _BATCH_SIZE rows so that locks are held briefly and progress
-    is preserved on crash. Already-rotated rows are detected and skipped,
-    making the operation safe to re-run.
-    """
-    if not global_version.is_ee_version():
-        raise RuntimeError("EE mode is not enabled — rotation requires EE encryption.")
-
-    if not ENCRYPTION_KEY_SECRET:
-        raise RuntimeError(
-            "ENCRYPTION_KEY_SECRET is not set — cannot rotate. "
-            "Set the target encryption key in the environment before running."
-        )
-
-    encrypted_columns = _discover_encrypted_columns()
-    totals: dict[str, int] = {}
-
-    for model_cls, col_name, pk_names, is_json in encrypted_columns:
-        table_name: str = model_cls.__tablename__  # type: ignore[attr-defined]
-        col_attr = getattr(model_cls, col_name)
-        pk_attrs = [getattr(model_cls, pk) for pk in pk_names]
-
-        # Read raw bytes directly, bypassing the TypeDecorator
-        raw_col = col_attr.property.columns[0]
-
-        stmt = select(*pk_attrs, raw_col.cast(LargeBinary)).where(col_attr.is_not(None))
-        rows = db_session.execute(stmt).all()
-
-        reencrypted = 0
-        batch_pending = 0
-        for row in rows:
-            raw_bytes: bytes | None = row[-1]
-            if raw_bytes is None:
-                continue
-
-            if _can_decrypt_with_current_key(raw_bytes):
-                continue
-
-            try:
-                if not old_key:
-                    decrypted_str = raw_bytes.decode("utf-8")
-                else:
-                    decrypted_str = decrypt_bytes_to_string(raw_bytes, key=old_key)
-
-                # For EncryptedJson, parse back to dict so the TypeDecorator
-                # can json.dumps() it cleanly (avoids double-encoding).
-                value: Any = json.loads(decrypted_str) if is_json else decrypted_str
-            except (ValueError, UnicodeDecodeError) as e:
-                pk_vals = [row[i] for i in range(len(pk_names))]
-                logger.warning(
-                    f"Could not decrypt/parse {table_name}.{col_name} "
-                    f"row {pk_vals} — skipping: {e}"
-                )
-                continue
-
-            if not dry_run:
-                pk_filters = [pk_attr == row[i] for i, pk_attr in enumerate(pk_attrs)]
-                update_stmt = (
-                    update(model_cls).where(*pk_filters).values({col_name: value})
-                )
-                db_session.execute(update_stmt)
-                batch_pending += 1
-
-                if batch_pending >= _BATCH_SIZE:
-                    db_session.commit()
-                    batch_pending = 0
-            reencrypted += 1
-
-        # Flush remaining rows in this column
-        if batch_pending > 0:
-            db_session.commit()
-
-        if reencrypted > 0:
-            totals[f"{table_name}.{col_name}"] = reencrypted
-            logger.info(
-                f"{'[DRY RUN] Would re-encrypt' if dry_run else 'Re-encrypted'} "
-                f"{reencrypted} value(s) in {table_name}.{col_name}"
-            )
-
-    return totals

backend/onyx/server/query_and_chat/streaming_models.py

@@ -41,7 +41,6 @@ class StreamingType(Enum):
     REASONING_DONE = "reasoning_done"
     CITATION_INFO = "citation_info"
     TOOL_CALL_DEBUG = "tool_call_debug"
-    TOOL_CALL_ARGUMENT_DELTA = "tool_call_argument_delta"
 
     MEMORY_TOOL_START = "memory_tool_start"
     MEMORY_TOOL_DELTA = "memory_tool_delta"
@@ -260,15 +259,6 @@ class CustomToolDelta(BaseObj):
     file_ids: list[str] | None = None
 
 
-class ToolCallArgumentDelta(BaseObj):
-    type: Literal["tool_call_argument_delta"] = (
-        StreamingType.TOOL_CALL_ARGUMENT_DELTA.value
-    )
-
-    tool_type: str
-    argument_deltas: dict[str, Any]
-
-
 ################################################
 # File Reader Packets
 ################################################
@@ -389,7 +379,6 @@ PacketObj = Union[
     # Citation Packets
     CitationInfo,
     ToolCallDebug,
-    ToolCallArgumentDelta,
     # Deep Research Packets
     DeepResearchPlanStart,
     DeepResearchPlanDelta,

backend/onyx/tools/built_in_tools.py

@@ -56,23 +56,3 @@ def get_built_in_tool_ids() -> list[str]:
 
 def get_built_in_tool_by_id(in_code_tool_id: str) -> Type[BUILT_IN_TOOL_TYPES]:
     return BUILT_IN_TOOL_MAP[in_code_tool_id]
-
-
-def _build_tool_name_to_class() -> dict[str, Type[BUILT_IN_TOOL_TYPES]]:
-    """Build a mapping from LLM-facing tool name to tool class."""
-    result: dict[str, Type[BUILT_IN_TOOL_TYPES]] = {}
-    for cls in BUILT_IN_TOOL_MAP.values():
-        name_attr = cls.__dict__.get("name")
-        if isinstance(name_attr, property) and name_attr.fget is not None:
-            tool_name = name_attr.fget(cls)
-        elif isinstance(name_attr, str):
-            tool_name = name_attr
-        else:
-            raise ValueError(
-                f"Built-in tool {cls.__name__} must define a valid LLM-facing tool name"
-            )
-        result[tool_name] = cls
-    return result
-
-
-TOOL_NAME_TO_CLASS: dict[str, Type[BUILT_IN_TOOL_TYPES]] = _build_tool_name_to_class()

backend/onyx/tools/interface.py

@@ -92,7 +92,3 @@ class Tool(abc.ABC, Generic[TOverride]):
         **llm_kwargs: Any,
     ) -> ToolResponse:
         raise NotImplementedError
-
-    @classmethod
-    def should_emit_argument_deltas(cls) -> bool:
-        return False

backend/onyx/tools/tool_implementations/python/python_tool.py

@@ -376,8 +376,3 @@ class PythonTool(Tool[PythonToolOverrideKwargs]):
                     rich_response=None,
                     llm_facing_response=llm_response,
                 )
-
-    @classmethod
-    @override
-    def should_emit_argument_deltas(cls) -> bool:
-        return True

backend/onyx/utils/encryption.py

@@ -11,20 +11,16 @@ logger = setup_logger()
 
 
 # IMPORTANT DO NOT DELETE, THIS IS USED BY fetch_versioned_implementation
-def _encrypt_string(input_str: str, key: str | None = None) -> bytes:  # noqa: ARG001
+def _encrypt_string(input_str: str) -> bytes:
     if ENCRYPTION_KEY_SECRET:
         logger.warning("MIT version of Onyx does not support encryption of secrets.")
-    elif key is not None:
-        logger.debug("MIT encrypt called with explicit key — key ignored.")
     return input_str.encode()
 
 
 # IMPORTANT DO NOT DELETE, THIS IS USED BY fetch_versioned_implementation
-def _decrypt_bytes(input_bytes: bytes, key: str | None = None) -> str:  # noqa: ARG001
-    if ENCRYPTION_KEY_SECRET:
-        logger.warning("MIT version of Onyx does not support decryption of secrets.")
-    elif key is not None:
-        logger.debug("MIT decrypt called with explicit key — key ignored.")
+def _decrypt_bytes(input_bytes: bytes) -> str:
+    # No need to double warn. If you wish to learn more about encryption features
+    # refer to the Onyx EE code
     return input_bytes.decode()
 
 
@@ -90,15 +86,15 @@ def _mask_list(items: list[Any]) -> list[Any]:
     return masked
 
 
-def encrypt_string_to_bytes(intput_str: str, key: str | None = None) -> bytes:
+def encrypt_string_to_bytes(intput_str: str) -> bytes:
     versioned_encryption_fn = fetch_versioned_implementation(
         "onyx.utils.encryption", "_encrypt_string"
     )
-    return versioned_encryption_fn(intput_str, key=key)
+    return versioned_encryption_fn(intput_str)
 
 
-def decrypt_bytes_to_string(intput_bytes: bytes, key: str | None = None) -> str:
+def decrypt_bytes_to_string(intput_bytes: bytes) -> str:
     versioned_decryption_fn = fetch_versioned_implementation(
         "onyx.utils.encryption", "_decrypt_bytes"
     )
-    return versioned_decryption_fn(intput_bytes, key=key)
+    return versioned_decryption_fn(intput_bytes)

backend/onyx/utils/variable_functionality.py

@@ -24,9 +24,6 @@ class OnyxVersion:
     def set_ee(self) -> None:
         self._is_ee = True
 
-    def unset_ee(self) -> None:
-        self._is_ee = False
-
     def is_ee_version(self) -> bool:
         return self._is_ee
 

backend/scripts/reencrypt_secrets.py

@@ -1,107 +0,0 @@
-"""Re-encrypt secrets under the current ENCRYPTION_KEY_SECRET.
-
-Decrypts all encrypted columns using the old key (or raw decode if the old key
-is empty), then re-encrypts them with the current ENCRYPTION_KEY_SECRET.
-
-Usage (docker):
-    docker exec -it onyx-api_server-1 \
-        python -m scripts.reencrypt_secrets --old-key "previous-key"
-
-Usage (kubernetes):
-    kubectl exec -it <pod> -- \
-        python -m scripts.reencrypt_secrets --old-key "previous-key"
-
-Omit --old-key (or pass "") if secrets were not previously encrypted.
-
-For multi-tenant deployments, pass --tenant-id to target a specific tenant,
-or --all-tenants to iterate every tenant.
-"""
-
-import argparse
-import os
-import sys
-
-parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
-sys.path.append(parent_dir)
-
-from onyx.db.rotate_encryption_key import rotate_encryption_key  # noqa: E402
-from onyx.db.engine.sql_engine import get_session_with_tenant  # noqa: E402
-from onyx.db.engine.sql_engine import SqlEngine  # noqa: E402
-from onyx.db.engine.tenant_utils import get_all_tenant_ids  # noqa: E402
-from onyx.utils.variable_functionality import global_version  # noqa: E402
-from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA  # noqa: E402
-
-
-def _run_for_tenant(tenant_id: str, old_key: str | None, dry_run: bool = False) -> None:
-    print(f"Re-encrypting secrets for tenant: {tenant_id}")
-    with get_session_with_tenant(tenant_id=tenant_id) as db_session:
-        results = rotate_encryption_key(db_session, old_key=old_key, dry_run=dry_run)
-
-    if results:
-        for col, count in results.items():
-            print(
-                f"  {col}: {count} row(s) {'would be ' if dry_run else ''}re-encrypted"
-            )
-    else:
-        print("No rows needed re-encryption.")
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser(
-        description="Re-encrypt secrets under the current encryption key."
-    )
-    parser.add_argument(
-        "--old-key",
-        default=None,
-        help="Previous encryption key. Omit or pass empty string if not applicable.",
-    )
-    parser.add_argument(
-        "--dry-run",
-        action="store_true",
-        help="Show what would be re-encrypted without making changes.",
-    )
-
-    tenant_group = parser.add_mutually_exclusive_group()
-    tenant_group.add_argument(
-        "--tenant-id",
-        default=None,
-        help="Target a specific tenant schema.",
-    )
-    tenant_group.add_argument(
-        "--all-tenants",
-        action="store_true",
-        help="Iterate all tenants.",
-    )
-
-    args = parser.parse_args()
-
-    old_key = args.old_key if args.old_key else None
-
-    global_version.set_ee()
-    SqlEngine.init_engine(pool_size=5, max_overflow=2)
-
-    if args.dry_run:
-        print("DRY RUN — no changes will be made")
-
-    if args.all_tenants:
-        tenant_ids = get_all_tenant_ids()
-        print(f"Found {len(tenant_ids)} tenant(s)")
-        failed_tenants: list[str] = []
-        for tid in tenant_ids:
-            try:
-                _run_for_tenant(tid, old_key, dry_run=args.dry_run)
-            except Exception as e:
-                print(f"  ERROR for tenant {tid}: {e}")
-                failed_tenants.append(tid)
-        if failed_tenants:
-            print(f"FAILED tenants ({len(failed_tenants)}): {failed_tenants}")
-            sys.exit(1)
-    else:
-        tenant_id = args.tenant_id or POSTGRES_DEFAULT_SCHEMA
-        _run_for_tenant(tenant_id, old_key, dry_run=args.dry_run)
-
-    print("Done.")
-
-
-if __name__ == "__main__":
-    main()

backend/tests/README.md

@@ -1,71 +0,0 @@
-# Backend Tests
-
-## Test Types
-
-There are four test categories, ordered by increasing scope:
-
-### Unit Tests (`tests/unit/`)
-
-No external services. Mock all I/O with `unittest.mock`. Use for complex, isolated
-logic (e.g. citation processing, encryption).
-
-```bash
-pytest -xv backend/tests/unit
-```
-
-### External Dependency Unit Tests (`tests/external_dependency_unit/`)
-
-External services (Postgres, Redis, Vespa, OpenAI, etc.) are running, but Onyx
-application containers are not. Tests call functions directly and can mock selectively.
-
-Use when you need a real database or real API calls but want control over setup.
-
-```bash
-python -m dotenv -f .vscode/.env run -- pytest backend/tests/external_dependency_unit
-```
-
-### Integration Tests (`tests/integration/`)
-
-Full Onyx deployment running. No mocking. Prefer this over other test types when possible.
-
-```bash
-python -m dotenv -f .vscode/.env run -- pytest backend/tests/integration
-```
-
-### Playwright / E2E Tests (`web/tests/e2e/`)
-
-Full stack including web server. Use for frontend-backend coordination.
-
-```bash
-npx playwright test <TEST_NAME>
-```
-
-## Shared Fixtures
-
-Shared fixtures live in `backend/tests/conftest.py`. Test subdirectories can define
-their own `conftest.py` for directory-scoped fixtures.
-
-## Best Practices
-
-### Use `enable_ee` fixture instead of inlining
-
-Enables EE mode for a test, with proper teardown and cache clearing.
-
-```python
-# Whole file (in a test module, NOT in conftest.py)
-pytestmark = pytest.mark.usefixtures("enable_ee")
-
-# Whole directory — add an autouse wrapper to the directory's conftest.py
-@pytest.fixture(autouse=True)
-def _enable_ee_for_directory(enable_ee: None) -> None:  # noqa: ARG001
-    """Wraps the shared enable_ee fixture with autouse for this directory."""
-
-# Single test
-def test_something(enable_ee: None) -> None: ...
-```
-
-**Note:** `pytestmark` in a `conftest.py` does NOT apply markers to tests in that
-directory — it only affects tests defined in the conftest itself (which is none).
-Use the autouse fixture wrapper pattern shown above instead.
-
-Do NOT inline `global_version.set_ee()` — always use the fixture.

backend/tests/conftest.py

@@ -1,24 +0,0 @@
-"""Root conftest — shared fixtures available to all test directories."""
-
-from collections.abc import Generator
-
-import pytest
-
-from onyx.utils.variable_functionality import fetch_versioned_implementation
-from onyx.utils.variable_functionality import global_version
-
-
-@pytest.fixture()
-def enable_ee() -> Generator[None, None, None]:
-    """Temporarily enable EE mode for a single test.
-
-    Restores the previous EE state and clears the versioned-implementation
-    cache on teardown so state doesn't leak between tests.
-    """
-    was_ee = global_version.is_ee_version()
-    global_version.set_ee()
-    fetch_versioned_implementation.cache_clear()
-    yield
-    if not was_ee:
-        global_version.unset_ee()
-    fetch_versioned_implementation.cache_clear()

backend/tests/daily/connectors/confluence/test_confluence_permissions_basic.py

@@ -45,7 +45,7 @@ def confluence_connector() -> ConfluenceConnector:
 def test_confluence_connector_permissions(
     mock_get_api_key: MagicMock,  # noqa: ARG001
     confluence_connector: ConfluenceConnector,
-    enable_ee: None,  # noqa: ARG001
+    set_ee_on: None,  # noqa: ARG001
 ) -> None:
     # Get all doc IDs from the full connector
     all_full_doc_ids = set()
@@ -93,7 +93,7 @@ def test_confluence_connector_permissions(
 def test_confluence_connector_restriction_handling(
     mock_get_api_key: MagicMock,  # noqa: ARG001
     mock_db_provider_class: MagicMock,
-    enable_ee: None,  # noqa: ARG001
+    set_ee_on: None,  # noqa: ARG001
 ) -> None:
     # Test space key
     test_space_key = "DailyPermS"

backend/tests/daily/connectors/conftest.py

@@ -4,6 +4,8 @@ from unittest.mock import patch
 
 import pytest
 
+from onyx.utils.variable_functionality import global_version
+
 
 @pytest.fixture
 def mock_get_unstructured_api_key() -> Generator[MagicMock, None, None]:
@@ -12,3 +14,14 @@ def mock_get_unstructured_api_key() -> Generator[MagicMock, None, None]:
         return_value=None,
     ) as mock:
         yield mock
+
+
+@pytest.fixture
+def set_ee_on() -> Generator[None, None, None]:
+    """Need EE to be enabled for these tests to work since
+    perm syncing is a an EE-only feature."""
+    global_version.set_ee()
+
+    yield
+
+    global_version._is_ee = False

backend/tests/daily/connectors/google_drive/test_drive_perm_sync.py

@@ -98,7 +98,7 @@ def _build_connector(
 
 def test_gdrive_perm_sync_with_real_data(
     google_drive_service_acct_connector_factory: Callable[..., GoogleDriveConnector],
-    enable_ee: None,  # noqa: ARG001
+    set_ee_on: None,  # noqa: ARG001
 ) -> None:
     """
     Test gdrive_doc_sync and gdrive_group_sync with real data from the test drive.

backend/tests/daily/connectors/slack/test_slack_perm_sync.py

@@ -1,10 +1,12 @@
 import time
+from collections.abc import Generator
 
 import pytest
 
 from onyx.connectors.models import HierarchyNode
 from onyx.connectors.models import SlimDocument
 from onyx.connectors.slack.connector import SlackConnector
+from onyx.utils.variable_functionality import global_version
 from tests.daily.connectors.utils import load_all_from_connector
 
 
@@ -17,7 +19,16 @@ PRIVATE_CHANNEL_USERS = [
     "test_user_2@onyx-test.com",
 ]
 
-pytestmark = pytest.mark.usefixtures("enable_ee")
+
+@pytest.fixture(autouse=True)
+def set_ee_on() -> Generator[None, None, None]:
+    """Need EE to be enabled for these tests to work since
+    perm syncing is a an EE-only feature."""
+    global_version.set_ee()
+
+    yield
+
+    global_version._is_ee = False
 
 
 @pytest.mark.parametrize(

backend/tests/daily/connectors/teams/test_teams_connector.py

@@ -1,11 +1,13 @@
 import os
 import time
+from collections.abc import Generator
 
 import pytest
 
 from onyx.access.models import ExternalAccess
 from onyx.connectors.models import HierarchyNode
 from onyx.connectors.teams.connector import TeamsConnector
+from onyx.utils.variable_functionality import global_version
 from tests.daily.connectors.teams.models import TeamsThread
 from tests.daily.connectors.utils import load_all_from_connector
 
@@ -166,9 +168,18 @@ def test_slim_docs_retrieval_from_teams_connector(
         _assert_is_valid_external_access(external_access=slim_doc.external_access)
 
 
+@pytest.fixture(autouse=False)
+def set_ee_on() -> Generator[None, None, None]:
+    """Need EE to be enabled for perm sync tests to work since
+    perm syncing is an EE-only feature."""
+    global_version.set_ee()
+    yield
+    global_version._is_ee = False
+
+
 def test_load_from_checkpoint_with_perm_sync(
     teams_connector: TeamsConnector,
-    enable_ee: None,  # noqa: ARG001
+    set_ee_on: None,  # noqa: ARG001
 ) -> None:
     """Test that load_from_checkpoint_with_perm_sync returns documents with external_access.
 

backend/tests/external_dependency_unit/connectors/jira/test_jira_doc_sync.py

@@ -1,6 +1,5 @@
 from typing import Any
 
-import pytest
 from pydantic import BaseModel
 from sqlalchemy.orm import Session
 
@@ -15,14 +14,13 @@ from onyx.db.models import ConnectorCredentialPair
 from onyx.db.models import Credential
 from onyx.db.utils import DocumentRow
 from onyx.db.utils import SortOrder
+from onyx.utils.variable_functionality import global_version
 
 
 # In order to get these tests to run, use the credentials from Bitwarden.
 # Search up "ENV vars for local and Github tests", and find the Jira relevant key-value pairs.
 # Required env vars: JIRA_USER_EMAIL, JIRA_API_TOKEN
 
-pytestmark = pytest.mark.usefixtures("enable_ee")
-
 
 class DocExternalAccessSet(BaseModel):
     """A version of DocExternalAccess that uses sets for comparison."""
@@ -54,6 +52,9 @@ def test_jira_doc_sync(
     This test uses the AS project which has applicationRole permission,
     meaning all documents should be marked as public.
     """
+    # NOTE: must set EE on or else the connector will skip the perm syncing
+    global_version.set_ee()
+
     try:
         # Use AS project specifically for this test
         connector_config = {
@@ -149,6 +150,9 @@ def test_jira_doc_sync_with_specific_permissions(
     This test uses a project that has specific user permissions to verify
     that specific users are correctly extracted.
     """
+    # NOTE: must set EE on or else the connector will skip the perm syncing
+    global_version.set_ee()
+
     try:
         # Use SUP project which has specific user permissions
         connector_config = {

backend/tests/external_dependency_unit/connectors/jira/test_jira_group_sync.py

@@ -1,6 +1,5 @@
 from typing import Any
 
-import pytest
 from sqlalchemy.orm import Session
 
 from ee.onyx.external_permissions.jira.group_sync import jira_group_sync
@@ -19,8 +18,6 @@ from tests.daily.connectors.confluence.models import ExternalUserGroupSet
 # Search up "ENV vars for local and Github tests", and find the Jira relevant key-value pairs.
 # Required env vars: JIRA_USER_EMAIL, JIRA_API_TOKEN
 
-pytestmark = pytest.mark.usefixtures("enable_ee")
-
 # Expected groups from the danswerai.atlassian.net Jira instance
 # Note: These groups are shared with Confluence since they're both Atlassian products
 # App accounts (bots, integrations) are filtered out

backend/tests/external_dependency_unit/db/test_rotate_encryption_key.py

@@ -1,305 +0,0 @@
-"""Tests for rotate_encryption_key against real Postgres.
-
-Uses real ORM models (Credential, InternetSearchProvider) and the actual
-Postgres database. Discovery is mocked in rotation tests to scope mutations
-to only the test rows — the real _discover_encrypted_columns walk is tested
-separately in TestDiscoverEncryptedColumns.
-
-Requires a running Postgres instance. Run with::
-
-    python -m dotenv -f .vscode/.env run -- pytest tests/external_dependency_unit/db/test_rotate_encryption_key.py
-"""
-
-import json
-from collections.abc import Generator
-from unittest.mock import patch
-
-import pytest
-from sqlalchemy import LargeBinary
-from sqlalchemy import select
-from sqlalchemy import text
-from sqlalchemy.orm import Session
-
-from ee.onyx.utils.encryption import _decrypt_bytes
-from ee.onyx.utils.encryption import _encrypt_string
-from ee.onyx.utils.encryption import _get_trimmed_key
-from onyx.configs.constants import DocumentSource
-from onyx.db.models import Credential
-from onyx.db.models import EncryptedJson
-from onyx.db.models import EncryptedString
-from onyx.db.models import InternetSearchProvider
-from onyx.db.rotate_encryption_key import _discover_encrypted_columns
-from onyx.db.rotate_encryption_key import rotate_encryption_key
-from onyx.utils.variable_functionality import fetch_versioned_implementation
-from onyx.utils.variable_functionality import global_version
-
-EE_MODULE = "ee.onyx.utils.encryption"
-ROTATE_MODULE = "onyx.db.rotate_encryption_key"
-
-OLD_KEY = "o" * 16
-NEW_KEY = "n" * 16
-
-
-@pytest.fixture(autouse=True)
-def _enable_ee() -> Generator[None, None, None]:
-    prev = global_version._is_ee
-    global_version.set_ee()
-    fetch_versioned_implementation.cache_clear()
-    yield
-    global_version._is_ee = prev
-    fetch_versioned_implementation.cache_clear()
-
-
-@pytest.fixture(autouse=True)
-def _clear_key_cache() -> None:
-    _get_trimmed_key.cache_clear()
-
-
-def _raw_credential_bytes(db_session: Session, credential_id: int) -> bytes | None:
-    """Read raw bytes from credential_json, bypassing the TypeDecorator."""
-    col = Credential.__table__.c.credential_json
-    stmt = select(col.cast(LargeBinary)).where(
-        Credential.__table__.c.id == credential_id
-    )
-    return db_session.execute(stmt).scalar()
-
-
-def _raw_isp_bytes(db_session: Session, isp_id: int) -> bytes | None:
-    """Read raw bytes from InternetSearchProvider.api_key."""
-    col = InternetSearchProvider.__table__.c.api_key
-    stmt = select(col.cast(LargeBinary)).where(
-        InternetSearchProvider.__table__.c.id == isp_id
-    )
-    return db_session.execute(stmt).scalar()
-
-
-class TestDiscoverEncryptedColumns:
-    """Verify _discover_encrypted_columns finds real production models."""
-
-    def test_discovers_credential_json(self) -> None:
-        results = _discover_encrypted_columns()
-        found = {
-            (model_cls.__tablename__, col_name, is_json)  # type: ignore[attr-defined]
-            for model_cls, col_name, _, is_json in results
-        }
-        assert ("credential", "credential_json", True) in found
-
-    def test_discovers_internet_search_provider_api_key(self) -> None:
-        results = _discover_encrypted_columns()
-        found = {
-            (model_cls.__tablename__, col_name, is_json)  # type: ignore[attr-defined]
-            for model_cls, col_name, _, is_json in results
-        }
-        assert ("internet_search_provider", "api_key", False) in found
-
-    def test_all_encrypted_string_columns_are_not_json(self) -> None:
-        results = _discover_encrypted_columns()
-        for model_cls, col_name, _, is_json in results:
-            col = getattr(model_cls, col_name).property.columns[0]
-            if isinstance(col.type, EncryptedString):
-                assert not is_json, (
-                    f"{model_cls.__tablename__}.{col_name} is EncryptedString "  # type: ignore[attr-defined]
-                    f"but is_json={is_json}"
-                )
-
-    def test_all_encrypted_json_columns_are_json(self) -> None:
-        results = _discover_encrypted_columns()
-        for model_cls, col_name, _, is_json in results:
-            col = getattr(model_cls, col_name).property.columns[0]
-            if isinstance(col.type, EncryptedJson):
-                assert is_json, (
-                    f"{model_cls.__tablename__}.{col_name} is EncryptedJson "  # type: ignore[attr-defined]
-                    f"but is_json={is_json}"
-                )
-
-
-class TestRotateCredential:
-    """Test rotation against the real Credential table (EncryptedJson).
-
-    Discovery is scoped to only the Credential model to avoid mutating
-    other tables in the test database.
-    """
-
-    @pytest.fixture(autouse=True)
-    def _limit_discovery(self) -> Generator[None, None, None]:
-        with patch(
-            f"{ROTATE_MODULE}._discover_encrypted_columns",
-            return_value=[(Credential, "credential_json", ["id"], True)],
-        ):
-            yield
-
-    @pytest.fixture()
-    def credential_id(
-        self, db_session: Session, tenant_context: None  # noqa: ARG002
-    ) -> Generator[int, None, None]:
-        """Insert a Credential row with raw encrypted bytes, clean up after."""
-        config = {"api_key": "sk-test-1234", "endpoint": "https://example.com"}
-        encrypted = _encrypt_string(json.dumps(config), key=OLD_KEY)
-
-        result = db_session.execute(
-            text(
-                "INSERT INTO credential "
-                "(source, credential_json, admin_public, curator_public) "
-                "VALUES (:source, :cred_json, true, false) "
-                "RETURNING id"
-            ),
-            {"source": DocumentSource.INGESTION_API.value, "cred_json": encrypted},
-        )
-        cred_id = result.scalar_one()
-        db_session.commit()
-
-        yield cred_id
-
-        db_session.execute(
-            text("DELETE FROM credential WHERE id = :id"), {"id": cred_id}
-        )
-        db_session.commit()
-
-    def test_rotates_credential_json(
-        self, db_session: Session, credential_id: int
-    ) -> None:
-        with (
-            patch(f"{ROTATE_MODULE}.ENCRYPTION_KEY_SECRET", NEW_KEY),
-            patch(f"{EE_MODULE}.ENCRYPTION_KEY_SECRET", NEW_KEY),
-        ):
-            totals = rotate_encryption_key(db_session, old_key=OLD_KEY)
-
-        assert totals.get("credential.credential_json", 0) >= 1
-
-        raw = _raw_credential_bytes(db_session, credential_id)
-        assert raw is not None
-        decrypted = json.loads(_decrypt_bytes(raw, key=NEW_KEY))
-        assert decrypted["api_key"] == "sk-test-1234"
-        assert decrypted["endpoint"] == "https://example.com"
-
-    def test_skips_already_rotated(
-        self, db_session: Session, credential_id: int
-    ) -> None:
-        with (
-            patch(f"{ROTATE_MODULE}.ENCRYPTION_KEY_SECRET", NEW_KEY),
-            patch(f"{EE_MODULE}.ENCRYPTION_KEY_SECRET", NEW_KEY),
-        ):
-            rotate_encryption_key(db_session, old_key=OLD_KEY)
-            _ = rotate_encryption_key(db_session, old_key=OLD_KEY)
-
-        raw = _raw_credential_bytes(db_session, credential_id)
-        assert raw is not None
-        decrypted = json.loads(_decrypt_bytes(raw, key=NEW_KEY))
-        assert decrypted["api_key"] == "sk-test-1234"
-
-    def test_dry_run_does_not_modify(
-        self, db_session: Session, credential_id: int
-    ) -> None:
-        original = _raw_credential_bytes(db_session, credential_id)
-
-        with (
-            patch(f"{ROTATE_MODULE}.ENCRYPTION_KEY_SECRET", NEW_KEY),
-            patch(f"{EE_MODULE}.ENCRYPTION_KEY_SECRET", NEW_KEY),
-        ):
-            totals = rotate_encryption_key(db_session, old_key=OLD_KEY, dry_run=True)
-
-        assert totals.get("credential.credential_json", 0) >= 1
-
-        raw_after = _raw_credential_bytes(db_session, credential_id)
-        assert raw_after == original
-
-
-class TestRotateInternetSearchProvider:
-    """Test rotation against the real InternetSearchProvider table (EncryptedString).
-
-    Discovery is scoped to only the InternetSearchProvider model to avoid
-    mutating other tables in the test database.
-    """
-
-    @pytest.fixture(autouse=True)
-    def _limit_discovery(self) -> Generator[None, None, None]:
-        with patch(
-            f"{ROTATE_MODULE}._discover_encrypted_columns",
-            return_value=[
-                (InternetSearchProvider, "api_key", ["id"], False),
-            ],
-        ):
-            yield
-
-    @pytest.fixture()
-    def isp_id(
-        self, db_session: Session, tenant_context: None  # noqa: ARG002
-    ) -> Generator[int, None, None]:
-        """Insert an InternetSearchProvider row with raw encrypted bytes."""
-        encrypted = _encrypt_string("sk-secret-api-key", key=OLD_KEY)
-
-        result = db_session.execute(
-            text(
-                "INSERT INTO internet_search_provider "
-                "(name, provider_type, api_key, is_active) "
-                "VALUES (:name, :ptype, :api_key, false) "
-                "RETURNING id"
-            ),
-            {
-                "name": f"test-rotation-{id(self)}",
-                "ptype": "test",
-                "api_key": encrypted,
-            },
-        )
-        isp_id = result.scalar_one()
-        db_session.commit()
-
-        yield isp_id
-
-        db_session.execute(
-            text("DELETE FROM internet_search_provider WHERE id = :id"),
-            {"id": isp_id},
-        )
-        db_session.commit()
-
-    def test_rotates_api_key(self, db_session: Session, isp_id: int) -> None:
-        with (
-            patch(f"{ROTATE_MODULE}.ENCRYPTION_KEY_SECRET", NEW_KEY),
-            patch(f"{EE_MODULE}.ENCRYPTION_KEY_SECRET", NEW_KEY),
-        ):
-            totals = rotate_encryption_key(db_session, old_key=OLD_KEY)
-
-        assert totals.get("internet_search_provider.api_key", 0) >= 1
-
-        raw = _raw_isp_bytes(db_session, isp_id)
-        assert raw is not None
-        assert _decrypt_bytes(raw, key=NEW_KEY) == "sk-secret-api-key"
-
-    def test_rotates_from_unencrypted(
-        self, db_session: Session, tenant_context: None  # noqa: ARG002
-    ) -> None:
-        """Test rotating data that was stored without any encryption key."""
-        result = db_session.execute(
-            text(
-                "INSERT INTO internet_search_provider "
-                "(name, provider_type, api_key, is_active) "
-                "VALUES (:name, :ptype, :api_key, false) "
-                "RETURNING id"
-            ),
-            {
-                "name": f"test-raw-{id(self)}",
-                "ptype": "test",
-                "api_key": b"raw-api-key",
-            },
-        )
-        isp_id = result.scalar_one()
-        db_session.commit()
-
-        try:
-            with (
-                patch(f"{ROTATE_MODULE}.ENCRYPTION_KEY_SECRET", NEW_KEY),
-                patch(f"{EE_MODULE}.ENCRYPTION_KEY_SECRET", NEW_KEY),
-            ):
-                totals = rotate_encryption_key(db_session, old_key=None)
-
-            assert totals.get("internet_search_provider.api_key", 0) >= 1
-
-            raw = _raw_isp_bytes(db_session, isp_id)
-            assert raw is not None
-            assert _decrypt_bytes(raw, key=NEW_KEY) == "raw-api-key"
-        finally:
-            db_session.execute(
-                text("DELETE FROM internet_search_provider WHERE id = :id"),
-                {"id": isp_id},
-            )
-            db_session.commit()

backend/tests/external_dependency_unit/tools/test_python_tool.py

@@ -950,7 +950,6 @@ from onyx.server.query_and_chat.streaming_models import Packet
 from onyx.server.query_and_chat.streaming_models import PythonToolDelta
 from onyx.server.query_and_chat.streaming_models import PythonToolStart
 from onyx.server.query_and_chat.streaming_models import SectionEnd
-from onyx.server.query_and_chat.streaming_models import ToolCallArgumentDelta
 from onyx.tools.tool_implementations.python.python_tool import PythonTool
 from tests.external_dependency_unit.answer.stream_test_builder import StreamTestBuilder
 from tests.external_dependency_unit.answer.stream_test_utils import create_chat_session
@@ -1292,21 +1291,12 @@ def test_code_interpreter_replay_packets_include_code_and_output(
                     tool_call_id="call_replay_test",
                     tool_call_argument_tokens=[json.dumps({"code": code})],
                 )
-            ).expect(
-                Packet(
-                    placement=create_placement(0),
-                    obj=ToolCallArgumentDelta(
-                        tool_type="python",
-                        argument_deltas={"code": code},
-                    ),
-                ),
-                forward=2,
             ).expect(
                 Packet(
                     placement=create_placement(0),
                     obj=PythonToolStart(code=code),
                 ),
-                forward=False,
+                forward=2,
             ).expect(
                 Packet(
                     placement=create_placement(0),

backend/tests/unit/ee/conftest.py

@@ -1,8 +0,0 @@
-"""Auto-enable EE mode for all tests under tests/unit/ee/."""
-
-import pytest
-
-
-@pytest.fixture(autouse=True)
-def _enable_ee_for_directory(enable_ee: None) -> None:  # noqa: ARG001
-    """Wraps the shared enable_ee fixture with autouse for this directory."""

backend/tests/unit/ee/onyx/utils/test_encryption.py

@@ -1,165 +0,0 @@
-"""Tests for EE AES-CBC encryption/decryption with explicit key support.
-
-With EE mode enabled (via conftest), fetch_versioned_implementation resolves
-to the EE implementations, so no patching of the MIT layer is needed.
-"""
-
-from unittest.mock import patch
-
-import pytest
-
-from ee.onyx.utils.encryption import _decrypt_bytes
-from ee.onyx.utils.encryption import _encrypt_string
-from ee.onyx.utils.encryption import _get_trimmed_key
-from ee.onyx.utils.encryption import decrypt_bytes_to_string
-from ee.onyx.utils.encryption import encrypt_string_to_bytes
-
-EE_MODULE = "ee.onyx.utils.encryption"
-
-# Keys must be exactly 16, 24, or 32 bytes for AES
-KEY_16 = "a" * 16
-KEY_16_ALT = "b" * 16
-KEY_24 = "d" * 24
-KEY_32 = "c" * 32
-
-
-@pytest.fixture(autouse=True)
-def _clear_key_cache() -> None:
-    _get_trimmed_key.cache_clear()
-
-
-class TestEncryptDecryptRoundTrip:
-    def test_roundtrip_with_env_key(self) -> None:
-        with patch(f"{EE_MODULE}.ENCRYPTION_KEY_SECRET", KEY_16):
-            encrypted = _encrypt_string("hello world")
-            assert encrypted != b"hello world"
-            assert _decrypt_bytes(encrypted) == "hello world"
-
-    def test_roundtrip_with_explicit_key(self) -> None:
-        encrypted = _encrypt_string("secret data", key=KEY_32)
-        assert encrypted != b"secret data"
-        assert _decrypt_bytes(encrypted, key=KEY_32) == "secret data"
-
-    def test_roundtrip_no_key(self) -> None:
-        """Without any key, data is raw-encoded (no encryption)."""
-        with patch(f"{EE_MODULE}.ENCRYPTION_KEY_SECRET", ""):
-            encrypted = _encrypt_string("plain text")
-            assert encrypted == b"plain text"
-            assert _decrypt_bytes(encrypted) == "plain text"
-
-    def test_explicit_key_overrides_env(self) -> None:
-        with patch(f"{EE_MODULE}.ENCRYPTION_KEY_SECRET", KEY_16):
-            encrypted = _encrypt_string("data", key=KEY_16_ALT)
-            with pytest.raises(ValueError):
-                _decrypt_bytes(encrypted, key=KEY_16)
-            assert _decrypt_bytes(encrypted, key=KEY_16_ALT) == "data"
-
-    def test_different_encryptions_produce_different_bytes(self) -> None:
-        """Each encryption uses a random IV, so results differ."""
-        a = _encrypt_string("same", key=KEY_16)
-        b = _encrypt_string("same", key=KEY_16)
-        assert a != b
-
-    def test_roundtrip_empty_string(self) -> None:
-        encrypted = _encrypt_string("", key=KEY_16)
-        assert encrypted != b""
-        assert _decrypt_bytes(encrypted, key=KEY_16) == ""
-
-    def test_roundtrip_unicode(self) -> None:
-        text = "日本語テスト 🔐 émojis"
-        encrypted = _encrypt_string(text, key=KEY_16)
-        assert _decrypt_bytes(encrypted, key=KEY_16) == text
-
-
-class TestDecryptFallbackBehavior:
-    def test_wrong_env_key_falls_back_to_raw_decode(self) -> None:
-        """Default key path: AES fails on non-AES data → fallback to raw decode."""
-        raw = "readable text".encode()
-        with patch(f"{EE_MODULE}.ENCRYPTION_KEY_SECRET", KEY_16):
-            assert _decrypt_bytes(raw) == "readable text"
-
-    def test_explicit_wrong_key_raises(self) -> None:
-        """Explicit key path: AES fails → raises, no fallback."""
-        encrypted = _encrypt_string("secret", key=KEY_16)
-        with pytest.raises(ValueError):
-            _decrypt_bytes(encrypted, key=KEY_16_ALT)
-
-    def test_explicit_none_key_with_no_env(self) -> None:
-        """key=None with empty env → raw decode."""
-        with patch(f"{EE_MODULE}.ENCRYPTION_KEY_SECRET", ""):
-            assert _decrypt_bytes(b"hello", key=None) == "hello"
-
-    def test_explicit_empty_string_key(self) -> None:
-        """key='' means no encryption."""
-        encrypted = _encrypt_string("test", key="")
-        assert encrypted == b"test"
-        assert _decrypt_bytes(encrypted, key="") == "test"
-
-
-class TestKeyValidation:
-    def test_key_too_short_raises(self) -> None:
-        with pytest.raises(RuntimeError, match="too short"):
-            _encrypt_string("data", key="short")
-
-    def test_16_byte_key(self) -> None:
-        encrypted = _encrypt_string("data", key=KEY_16)
-        assert _decrypt_bytes(encrypted, key=KEY_16) == "data"
-
-    def test_24_byte_key(self) -> None:
-        encrypted = _encrypt_string("data", key=KEY_24)
-        assert _decrypt_bytes(encrypted, key=KEY_24) == "data"
-
-    def test_32_byte_key(self) -> None:
-        encrypted = _encrypt_string("data", key=KEY_32)
-        assert _decrypt_bytes(encrypted, key=KEY_32) == "data"
-
-    def test_long_key_truncated_to_32(self) -> None:
-        """Keys longer than 32 bytes are truncated to 32."""
-        long_key = "e" * 64
-        encrypted = _encrypt_string("data", key=long_key)
-        assert _decrypt_bytes(encrypted, key=long_key) == "data"
-
-    def test_20_byte_key_trimmed_to_16(self) -> None:
-        """A 20-byte key is trimmed to the largest valid AES size that fits (16)."""
-        key_20 = "f" * 20
-        encrypted = _encrypt_string("data", key=key_20)
-        assert _decrypt_bytes(encrypted, key=key_20) == "data"
-
-        # Verify it was trimmed to 16 by checking that the first 16 bytes
-        # of the key can also decrypt it
-        key_16_same_prefix = "f" * 16
-        assert _decrypt_bytes(encrypted, key=key_16_same_prefix) == "data"
-
-    def test_25_byte_key_trimmed_to_24(self) -> None:
-        """A 25-byte key is trimmed to the largest valid AES size that fits (24)."""
-        key_25 = "g" * 25
-        encrypted = _encrypt_string("data", key=key_25)
-        assert _decrypt_bytes(encrypted, key=key_25) == "data"
-
-        key_24_same_prefix = "g" * 24
-        assert _decrypt_bytes(encrypted, key=key_24_same_prefix) == "data"
-
-    def test_30_byte_key_trimmed_to_24(self) -> None:
-        """A 30-byte key is trimmed to the largest valid AES size that fits (24)."""
-        key_30 = "h" * 30
-        encrypted = _encrypt_string("data", key=key_30)
-        assert _decrypt_bytes(encrypted, key=key_30) == "data"
-
-        key_24_same_prefix = "h" * 24
-        assert _decrypt_bytes(encrypted, key=key_24_same_prefix) == "data"
-
-
-class TestWrapperFunctions:
-    """Test encrypt_string_to_bytes / decrypt_bytes_to_string pass key through.
-
-    With EE mode enabled, the wrappers resolve to EE implementations automatically.
-    """
-
-    def test_wrapper_passes_key(self) -> None:
-        encrypted = encrypt_string_to_bytes("payload", key=KEY_16)
-        assert decrypt_bytes_to_string(encrypted, key=KEY_16) == "payload"
-
-    def test_wrapper_no_key_uses_env(self) -> None:
-        with patch(f"{EE_MODULE}.ENCRYPTION_KEY_SECRET", KEY_32):
-            encrypted = encrypt_string_to_bytes("payload")
-            assert decrypt_bytes_to_string(encrypted) == "payload"

backend/tests/unit/onyx/chat/test_argument_delta_streaming.py

@@ -1,630 +0,0 @@
-from typing import Any
-from unittest.mock import MagicMock
-from unittest.mock import patch
-
-from onyx.chat.tool_call_args_streaming import maybe_emit_argument_delta
-from onyx.server.query_and_chat.placement import Placement
-from onyx.server.query_and_chat.streaming_models import ToolCallArgumentDelta
-from onyx.utils.jsonriver import Parser
-
-
-def _make_tool_call_delta(
-    index: int = 0,
-    name: str | None = None,
-    arguments: str | None = None,
-    function_is_none: bool = False,
-) -> MagicMock:
-    """Create a mock tool_call_delta matching the LiteLLM streaming shape."""
-    delta = MagicMock()
-    delta.index = index
-    if function_is_none:
-        delta.function = None
-    else:
-        delta.function = MagicMock()
-        delta.function.name = name
-        delta.function.arguments = arguments
-    return delta
-
-
-def _make_placement() -> Placement:
-    return Placement(turn_index=0, tab_index=0)
-
-
-def _mock_tool_class(emit: bool = True) -> MagicMock:
-    cls = MagicMock()
-    cls.should_emit_argument_deltas.return_value = emit
-    return cls
-
-
-def _collect(
-    tc_map: dict[int, dict[str, Any]],
-    delta: MagicMock,
-    placement: Placement | None = None,
-    parsers: dict[int, Parser] | None = None,
-) -> list[Any]:
-    """Run maybe_emit_argument_delta and return the yielded packets."""
-    return list(
-        maybe_emit_argument_delta(
-            tc_map,
-            delta,
-            placement or _make_placement(),
-            parsers if parsers is not None else {},
-        )
-    )
-
-
-def _stream_fragments(
-    fragments: list[str],
-    tc_map: dict[int, dict[str, Any]],
-    placement: Placement | None = None,
-) -> list[str]:
-    """Feed fragments into maybe_emit_argument_delta one by one, returning
-    all emitted content values concatenated per-key as a flat list."""
-    pl = placement or _make_placement()
-    parsers: dict[int, Parser] = {}
-    emitted: list[str] = []
-    for frag in fragments:
-        tc_map[0]["arguments"] += frag
-        delta = _make_tool_call_delta(arguments=frag)
-        for packet in maybe_emit_argument_delta(tc_map, delta, pl, parsers=parsers):
-            obj = packet.obj
-            assert isinstance(obj, ToolCallArgumentDelta)
-            for value in obj.argument_deltas.values():
-                emitted.append(value)
-    return emitted
-
-
-class TestMaybeEmitArgumentDeltaGuards:
-    """Tests for conditions that cause no packet to be emitted."""
-
-    @patch("onyx.chat.tool_call_args_streaming._get_tool_class")
-    def test_no_emission_when_tool_does_not_opt_in(
-        self, mock_get_tool: MagicMock
-    ) -> None:
-        """Tools that return False from should_emit_argument_deltas emit nothing."""
-        mock_get_tool.return_value = _mock_tool_class(emit=False)
-
-        tc_map: dict[int, dict[str, Any]] = {
-            0: {"id": "tc_1", "name": "python", "arguments": '{"code": "x'}
-        }
-        assert _collect(tc_map, _make_tool_call_delta(arguments="x")) == []
-
-    @patch("onyx.chat.tool_call_args_streaming._get_tool_class")
-    def test_no_emission_when_tool_class_unknown(
-        self, mock_get_tool: MagicMock
-    ) -> None:
-        mock_get_tool.return_value = None
-
-        tc_map: dict[int, dict[str, Any]] = {
-            0: {"id": "tc_1", "name": "unknown", "arguments": '{"code": "x'}
-        }
-        assert _collect(tc_map, _make_tool_call_delta(arguments="x")) == []
-
-    @patch("onyx.chat.tool_call_args_streaming._get_tool_class")
-    def test_no_emission_when_no_argument_fragment(
-        self, mock_get_tool: MagicMock
-    ) -> None:
-        mock_get_tool.return_value = _mock_tool_class()
-
-        tc_map: dict[int, dict[str, Any]] = {
-            0: {"id": "tc_1", "name": "python", "arguments": '{"code": "x'}
-        }
-        assert _collect(tc_map, _make_tool_call_delta(arguments=None)) == []
-
-    @patch("onyx.chat.tool_call_args_streaming._get_tool_class")
-    def test_no_emission_when_key_value_not_started(
-        self, mock_get_tool: MagicMock
-    ) -> None:
-        """Key exists in JSON but its string value hasn't begun yet."""
-        mock_get_tool.return_value = _mock_tool_class()
-
-        tc_map: dict[int, dict[str, Any]] = {
-            0: {"id": "tc_1", "name": "python", "arguments": '{"code":'}
-        }
-        assert _collect(tc_map, _make_tool_call_delta(arguments=":")) == []
-
-    @patch("onyx.chat.tool_call_args_streaming._get_tool_class")
-    def test_no_emission_before_any_key(self, mock_get_tool: MagicMock) -> None:
-        """Only the opening brace has arrived — no key to stream yet."""
-        mock_get_tool.return_value = _mock_tool_class()
-
-        tc_map: dict[int, dict[str, Any]] = {
-            0: {"id": "tc_1", "name": "python", "arguments": "{"}
-        }
-        assert _collect(tc_map, _make_tool_call_delta(arguments="{")) == []
-
-
-class TestMaybeEmitArgumentDeltaBasic:
-    """Tests for correct packet content and incremental emission."""
-
-    @patch("onyx.chat.tool_call_args_streaming._get_tool_class")
-    def test_emits_packet_with_correct_fields(self, mock_get_tool: MagicMock) -> None:
-        mock_get_tool.return_value = _mock_tool_class()
-
-        tc_map: dict[int, dict[str, Any]] = {
-            0: {"id": "tc_1", "name": "python", "arguments": ""}
-        }
-        fragments = ['{"code": "', "print(1)", '"}']
-
-        pl = _make_placement()
-        parsers: dict[int, Parser] = {}
-        all_packets = []
-        for frag in fragments:
-            tc_map[0]["arguments"] += frag
-            packets = _collect(
-                tc_map, _make_tool_call_delta(arguments=frag), pl, parsers
-            )
-            all_packets.extend(packets)
-
-        assert len(all_packets) >= 1
-        # Verify packet structure
-        obj = all_packets[0].obj
-        assert isinstance(obj, ToolCallArgumentDelta)
-        assert obj.tool_type == "python"
-        # All emitted content should reconstruct the value
-        full_code = ""
-        for p in all_packets:
-            assert isinstance(p.obj, ToolCallArgumentDelta)
-            if "code" in p.obj.argument_deltas:
-                full_code += p.obj.argument_deltas["code"]
-        assert full_code == "print(1)"
-
-    @patch("onyx.chat.tool_call_args_streaming._get_tool_class")
-    def test_emits_only_new_content_on_subsequent_call(
-        self, mock_get_tool: MagicMock
-    ) -> None:
-        """After a first emission, subsequent calls emit only the diff."""
-        mock_get_tool.return_value = _mock_tool_class()
-
-        tc_map: dict[int, dict[str, Any]] = {
-            0: {"id": "tc_1", "name": "python", "arguments": ""}
-        }
-        parsers: dict[int, Parser] = {}
-        pl = _make_placement()
-
-        # First fragment opens the string
-        tc_map[0]["arguments"] = '{"code": "abc'
-        packets_1 = _collect(
-            tc_map, _make_tool_call_delta(arguments='{"code": "abc'), pl, parsers
-        )
-        code_1 = ""
-        for p in packets_1:
-            assert isinstance(p.obj, ToolCallArgumentDelta)
-            code_1 += p.obj.argument_deltas.get("code", "")
-        assert code_1 == "abc"
-
-        # Second fragment appends more
-        tc_map[0]["arguments"] = '{"code": "abcdef'
-        packets_2 = _collect(
-            tc_map, _make_tool_call_delta(arguments="def"), pl, parsers
-        )
-        code_2 = ""
-        for p in packets_2:
-            assert isinstance(p.obj, ToolCallArgumentDelta)
-            code_2 += p.obj.argument_deltas.get("code", "")
-        assert code_2 == "def"
-
-    @patch("onyx.chat.tool_call_args_streaming._get_tool_class")
-    def test_handles_multiple_keys_sequentially(self, mock_get_tool: MagicMock) -> None:
-        """When a second key starts, emissions switch to that key."""
-        mock_get_tool.return_value = _mock_tool_class()
-
-        tc_map: dict[int, dict[str, Any]] = {
-            0: {"id": "tc_1", "name": "python", "arguments": ""}
-        }
-        fragments = [
-            '{"code": "x',
-            '", "output": "hello',
-            '"}',
-        ]
-
-        emitted = _stream_fragments(fragments, tc_map)
-        full = "".join(emitted)
-        assert "x" in full
-        assert "hello" in full
-
-    @patch("onyx.chat.tool_call_args_streaming._get_tool_class")
-    def test_delta_spans_key_boundary(self, mock_get_tool: MagicMock) -> None:
-        """A single delta contains the end of one value and the start of the next key."""
-        mock_get_tool.return_value = _mock_tool_class()
-
-        tc_map: dict[int, dict[str, Any]] = {
-            0: {"id": "tc_1", "name": "python", "arguments": ""}
-        }
-        fragments = [
-            '{"code": "x',
-            'y", "lang": "py',
-            '"}',
-        ]
-
-        emitted = _stream_fragments(fragments, tc_map)
-        full = "".join(emitted)
-        assert "xy" in full
-        assert "py" in full
-
-    @patch("onyx.chat.tool_call_args_streaming._get_tool_class")
-    def test_empty_value_emits_nothing(self, mock_get_tool: MagicMock) -> None:
-        """An empty string value has nothing to emit."""
-        mock_get_tool.return_value = _mock_tool_class()
-
-        tc_map: dict[int, dict[str, Any]] = {
-            0: {"id": "tc_1", "name": "python", "arguments": ""}
-        }
-        # Opening quote just arrived, value is empty
-        tc_map[0]["arguments"] = '{"code": "'
-        packets = _collect(tc_map, _make_tool_call_delta(arguments='{"code": "'))
-        # No string content yet, so either no packet or empty deltas
-        for p in packets:
-            assert isinstance(p.obj, ToolCallArgumentDelta)
-            assert p.obj.argument_deltas.get("code", "") == ""
-
-
-class TestMaybeEmitArgumentDeltaDecoding:
-    """Tests verifying that JSON escape sequences are properly decoded."""
-
-    @patch("onyx.chat.tool_call_args_streaming._get_tool_class")
-    def test_decodes_newlines(self, mock_get_tool: MagicMock) -> None:
-        mock_get_tool.return_value = _mock_tool_class()
-
-        tc_map: dict[int, dict[str, Any]] = {
-            0: {"id": "tc_1", "name": "python", "arguments": ""}
-        }
-        fragments = ['{"code": "line1\\nline2"}']
-
-        emitted = _stream_fragments(fragments, tc_map)
-        assert "".join(emitted) == "line1\nline2"
-
-    @patch("onyx.chat.tool_call_args_streaming._get_tool_class")
-    def test_decodes_tabs(self, mock_get_tool: MagicMock) -> None:
-        mock_get_tool.return_value = _mock_tool_class()
-
-        tc_map: dict[int, dict[str, Any]] = {
-            0: {"id": "tc_1", "name": "python", "arguments": ""}
-        }
-        fragments = ['{"code": "\\tindented"}']
-
-        emitted = _stream_fragments(fragments, tc_map)
-        assert "".join(emitted) == "\tindented"
-
-    @patch("onyx.chat.tool_call_args_streaming._get_tool_class")
-    def test_decodes_escaped_quotes(self, mock_get_tool: MagicMock) -> None:
-        mock_get_tool.return_value = _mock_tool_class()
-
-        tc_map: dict[int, dict[str, Any]] = {
-            0: {"id": "tc_1", "name": "python", "arguments": ""}
-        }
-        fragments = ['{"code": "say \\"hi\\""}']
-
-        emitted = _stream_fragments(fragments, tc_map)
-        assert "".join(emitted) == 'say "hi"'
-
-    @patch("onyx.chat.tool_call_args_streaming._get_tool_class")
-    def test_decodes_escaped_backslashes(self, mock_get_tool: MagicMock) -> None:
-        mock_get_tool.return_value = _mock_tool_class()
-
-        tc_map: dict[int, dict[str, Any]] = {
-            0: {"id": "tc_1", "name": "python", "arguments": ""}
-        }
-        fragments = ['{"code": "path\\\\dir"}']
-
-        emitted = _stream_fragments(fragments, tc_map)
-        assert "".join(emitted) == "path\\dir"
-
-    @patch("onyx.chat.tool_call_args_streaming._get_tool_class")
-    def test_decodes_unicode_escape(self, mock_get_tool: MagicMock) -> None:
-        mock_get_tool.return_value = _mock_tool_class()
-
-        tc_map: dict[int, dict[str, Any]] = {
-            0: {"id": "tc_1", "name": "python", "arguments": ""}
-        }
-        fragments = ['{"code": "\\u0041"}']
-
-        emitted = _stream_fragments(fragments, tc_map)
-        assert "".join(emitted) == "A"
-
-    @patch("onyx.chat.tool_call_args_streaming._get_tool_class")
-    def test_incomplete_escape_at_end_decoded_on_next_chunk(
-        self, mock_get_tool: MagicMock
-    ) -> None:
-        """A trailing backslash (incomplete escape) is completed in the next chunk."""
-        mock_get_tool.return_value = _mock_tool_class()
-
-        tc_map: dict[int, dict[str, Any]] = {
-            0: {"id": "tc_1", "name": "python", "arguments": ""}
-        }
-        fragments = ['{"code": "hello\\', 'n"}']
-
-        emitted = _stream_fragments(fragments, tc_map)
-        assert "".join(emitted) == "hello\n"
-
-    @patch("onyx.chat.tool_call_args_streaming._get_tool_class")
-    def test_incomplete_unicode_escape_completed_on_next_chunk(
-        self, mock_get_tool: MagicMock
-    ) -> None:
-        """A partial \\uXX sequence is completed in the next chunk."""
-        mock_get_tool.return_value = _mock_tool_class()
-
-        tc_map: dict[int, dict[str, Any]] = {
-            0: {"id": "tc_1", "name": "python", "arguments": ""}
-        }
-        fragments = ['{"code": "hello\\u00', '41"}']
-
-        emitted = _stream_fragments(fragments, tc_map)
-        assert "".join(emitted) == "helloA"
-
-
-class TestArgumentDeltaStreamingE2E:
-    """Simulates realistic sequences of LLM argument deltas to verify
-    the full pipeline produces correct decoded output."""
-
-    @patch("onyx.chat.tool_call_args_streaming._get_tool_class")
-    def test_realistic_python_code_streaming(self, mock_get_tool: MagicMock) -> None:
-        """Streams: {"code": "print('hello')\\nprint('world')"}"""
-        mock_get_tool.return_value = _mock_tool_class()
-
-        tc_map: dict[int, dict[str, Any]] = {
-            0: {"id": "tc_1", "name": "python", "arguments": ""}
-        }
-        fragments = [
-            '{"',
-            "code",
-            '": "',
-            "print(",
-            "'hello')",
-            "\\n",
-            "print(",
-            "'world')",
-            '"}',
-        ]
-
-        full = "".join(_stream_fragments(fragments, tc_map))
-        assert full == "print('hello')\nprint('world')"
-
-    @patch("onyx.chat.tool_call_args_streaming._get_tool_class")
-    def test_streaming_with_tabs_and_newlines(self, mock_get_tool: MagicMock) -> None:
-        """Streams code with tabs and newlines."""
-        mock_get_tool.return_value = _mock_tool_class()
-
-        tc_map: dict[int, dict[str, Any]] = {
-            0: {"id": "tc_1", "name": "python", "arguments": ""}
-        }
-        fragments = [
-            '{"code": "',
-            "if True:",
-            "\\n",
-            "\\t",
-            "pass",
-            '"}',
-        ]
-
-        full = "".join(_stream_fragments(fragments, tc_map))
-        assert full == "if True:\n\tpass"
-
-    @patch("onyx.chat.tool_call_args_streaming._get_tool_class")
-    def test_split_escape_sequence(self, mock_get_tool: MagicMock) -> None:
-        """An escape sequence split across two fragments (backslash in one,
-        'n' in the next) should still decode correctly."""
-        mock_get_tool.return_value = _mock_tool_class()
-
-        tc_map: dict[int, dict[str, Any]] = {
-            0: {"id": "tc_1", "name": "python", "arguments": ""}
-        }
-        fragments = [
-            '{"code": "hello',
-            "\\",
-            "n",
-            'world"}',
-        ]
-
-        full = "".join(_stream_fragments(fragments, tc_map))
-        assert full == "hello\nworld"
-
-    @patch("onyx.chat.tool_call_args_streaming._get_tool_class")
-    def test_multiple_newlines_and_indentation(self, mock_get_tool: MagicMock) -> None:
-        """Streams a multi-line function with multiple escape sequences."""
-        mock_get_tool.return_value = _mock_tool_class()
-
-        tc_map: dict[int, dict[str, Any]] = {
-            0: {"id": "tc_1", "name": "python", "arguments": ""}
-        }
-        fragments = [
-            '{"code": "',
-            "def foo():",
-            "\\n",
-            "\\t",
-            "x = 1",
-            "\\n",
-            "\\t",
-            "return x",
-            '"}',
-        ]
-
-        full = "".join(_stream_fragments(fragments, tc_map))
-        assert full == "def foo():\n\tx = 1\n\treturn x"
-
-    @patch("onyx.chat.tool_call_args_streaming._get_tool_class")
-    def test_two_keys_streamed_sequentially(self, mock_get_tool: MagicMock) -> None:
-        """Streams code first, then a second key (language) — both decoded."""
-        mock_get_tool.return_value = _mock_tool_class()
-
-        tc_map: dict[int, dict[str, Any]] = {
-            0: {"id": "tc_1", "name": "python", "arguments": ""}
-        }
-        fragments = [
-            '{"code": "',
-            "x = 1",
-            '", "language": "',
-            "python",
-            '"}',
-        ]
-
-        emitted = _stream_fragments(fragments, tc_map)
-        # Should have emissions for both keys
-        full = "".join(emitted)
-        assert "x = 1" in full
-        assert "python" in full
-
-    @patch("onyx.chat.tool_call_args_streaming._get_tool_class")
-    def test_code_containing_dict_literal(self, mock_get_tool: MagicMock) -> None:
-        """Python code like `x = {"key": "val"}` contains JSON-like patterns.
-        The escaped quotes inside the *outer* JSON value should prevent the
-        inner `"key":` from being mistaken for a top-level JSON key."""
-        mock_get_tool.return_value = _mock_tool_class()
-
-        tc_map: dict[int, dict[str, Any]] = {
-            0: {"id": "tc_1", "name": "python", "arguments": ""}
-        }
-        # The LLM sends: {"code": "x = {\"key\": \"val\"}"}
-        # The inner quotes are escaped as \" in the JSON value.
-        fragments = [
-            '{"code": "',
-            "x = {",
-            '\\"key\\"',
-            ": ",
-            '\\"val\\"',
-            "}",
-            '"}',
-        ]
-
-        full = "".join(_stream_fragments(fragments, tc_map))
-        assert full == 'x = {"key": "val"}'
-
-    @patch("onyx.chat.tool_call_args_streaming._get_tool_class")
-    def test_code_with_colon_in_value(self, mock_get_tool: MagicMock) -> None:
-        """Colons inside the string value should not confuse key detection."""
-        mock_get_tool.return_value = _mock_tool_class()
-
-        tc_map: dict[int, dict[str, Any]] = {
-            0: {"id": "tc_1", "name": "python", "arguments": ""}
-        }
-        fragments = [
-            '{"code": "',
-            "url = ",
-            '\\"https://example.com\\"',
-            '"}',
-        ]
-
-        full = "".join(_stream_fragments(fragments, tc_map))
-        assert full == 'url = "https://example.com"'
-
-
-class TestMaybeEmitArgumentDeltaEdgeCases:
-    """Edge cases not covered by the standard test classes."""
-
-    @patch("onyx.chat.tool_call_args_streaming._get_tool_class")
-    def test_no_emission_when_function_is_none(self, mock_get_tool: MagicMock) -> None:
-        """Some delta chunks have function=None (e.g. role-only deltas)."""
-        mock_get_tool.return_value = _mock_tool_class()
-
-        tc_map: dict[int, dict[str, Any]] = {
-            0: {"id": "tc_1", "name": "python", "arguments": '{"code": "x'}
-        }
-        delta = _make_tool_call_delta(arguments=None, function_is_none=True)
-        assert _collect(tc_map, delta) == []
-
-    @patch("onyx.chat.tool_call_args_streaming._get_tool_class")
-    def test_multiple_concurrent_tool_calls(self, mock_get_tool: MagicMock) -> None:
-        """Two tool calls streaming at different indices in parallel."""
-        mock_get_tool.return_value = _mock_tool_class()
-
-        tc_map: dict[int, dict[str, Any]] = {
-            0: {"id": "tc_1", "name": "python", "arguments": ""},
-            1: {"id": "tc_2", "name": "python", "arguments": ""},
-        }
-
-        parsers: dict[int, Parser] = {}
-        pl = _make_placement()
-
-        # Feed full JSON to index 0
-        tc_map[0]["arguments"] = '{"code": "aaa"}'
-        packets_0 = _collect(
-            tc_map,
-            _make_tool_call_delta(index=0, arguments='{"code": "aaa"}'),
-            pl,
-            parsers,
-        )
-        code_0 = ""
-        for p in packets_0:
-            assert isinstance(p.obj, ToolCallArgumentDelta)
-            code_0 += p.obj.argument_deltas.get("code", "")
-        assert code_0 == "aaa"
-
-        # Feed full JSON to index 1
-        tc_map[1]["arguments"] = '{"code": "bbb"}'
-        packets_1 = _collect(
-            tc_map,
-            _make_tool_call_delta(index=1, arguments='{"code": "bbb"}'),
-            pl,
-            parsers,
-        )
-        code_1 = ""
-        for p in packets_1:
-            assert isinstance(p.obj, ToolCallArgumentDelta)
-            code_1 += p.obj.argument_deltas.get("code", "")
-        assert code_1 == "bbb"
-
-    @patch("onyx.chat.tool_call_args_streaming._get_tool_class")
-    def test_delta_with_four_arguments(self, mock_get_tool: MagicMock) -> None:
-        """A single delta contains four complete key-value pairs."""
-        mock_get_tool.return_value = _mock_tool_class()
-
-        full = '{"a": "one", "b": "two", "c": "three", "d": "four"}'
-        tc_map: dict[int, dict[str, Any]] = {
-            0: {"id": "tc_1", "name": "python", "arguments": ""}
-        }
-        tc_map[0]["arguments"] = full
-        parsers: dict[int, Parser] = {}
-        packets = _collect(
-            tc_map, _make_tool_call_delta(arguments=full), parsers=parsers
-        )
-
-        # Collect all argument deltas across packets
-        all_deltas: dict[str, str] = {}
-        for p in packets:
-            assert isinstance(p.obj, ToolCallArgumentDelta)
-            for k, v in p.obj.argument_deltas.items():
-                all_deltas[k] = all_deltas.get(k, "") + v
-
-        assert all_deltas == {
-            "a": "one",
-            "b": "two",
-            "c": "three",
-            "d": "four",
-        }
-
-    @patch("onyx.chat.tool_call_args_streaming._get_tool_class")
-    def test_delta_on_second_arg_after_first_complete(
-        self, mock_get_tool: MagicMock
-    ) -> None:
-        """First argument is fully complete; delta only adds to the second."""
-        mock_get_tool.return_value = _mock_tool_class()
-
-        tc_map: dict[int, dict[str, Any]] = {
-            0: {"id": "tc_1", "name": "python", "arguments": ""}
-        }
-
-        fragments = [
-            '{"code": "print(1)", "lang": "py',
-            '"}',
-        ]
-
-        emitted = _stream_fragments(fragments, tc_map)
-        full = "".join(emitted)
-        assert "print(1)" in full
-        assert "py" in full
-
-    @patch("onyx.chat.tool_call_args_streaming._get_tool_class")
-    def test_non_string_values_skipped(self, mock_get_tool: MagicMock) -> None:
-        """Non-string values (numbers, booleans, null) are skipped — they are
-        available in the final tool-call kickoff packet. String arguments
-        following them are still emitted."""
-        mock_get_tool.return_value = _mock_tool_class()
-
-        tc_map: dict[int, dict[str, Any]] = {
-            0: {"id": "tc_1", "name": "python", "arguments": ""}
-        }
-        fragments = ['{"timeout": 30, "code": "hello"}']
-
-        emitted = _stream_fragments(fragments, tc_map)
-        full = "".join(emitted)
-        assert full == "hello"

backend/tests/unit/onyx/connectors/jira/test_jira_permission_sync.py

@@ -9,8 +9,6 @@ from onyx.connectors.jira.utils import JIRA_SERVER_API_VERSION
 from onyx.db.models import ConnectorCredentialPair
 from onyx.utils.sensitive import make_mock_sensitive_value
 
-pytestmark = pytest.mark.usefixtures("enable_ee")
-
 
 @pytest.fixture
 def mock_jira_cc_pair(

cli/Dockerfile

@@ -1,22 +0,0 @@
-FROM golang:1.26-alpine@sha256:2389ebfa5b7f43eeafbd6be0c3700cc46690ef842ad962f6c5bd6be49ed82039 AS builder
-
-WORKDIR /app
-COPY ./ .
-
-ARG TARGETARCH
-RUN CGO_ENABLED=0 GOOS=linux GOARCH=${TARGETARCH} go build -ldflags="-s -w" -o onyx-cli .
-RUN mkdir -p /home/onyx/.config
-
-FROM scratch
-
-COPY --from=builder /etc/ssl/certs/ca-certificates.crt /etc/ssl/certs/
-COPY --from=builder --chown=65534:65534 /home/onyx /home/onyx
-
-COPY --from=builder /app/onyx-cli /onyx-cli
-
-ENV HOME=/home/onyx
-ENV XDG_CONFIG_HOME=/home/onyx/.config
-
-USER 65534:65534
-
-ENTRYPOINT ["/onyx-cli"]

cli/README.md

@@ -63,6 +63,31 @@ onyx-cli agents
 onyx-cli agents --json
 ```
 
+### Serve over SSH
+
+```shell
+# Start a public SSH endpoint for the CLI TUI
+onyx-cli serve --host 0.0.0.0 --port 2222
+
+# Connect as a client
+ssh your-host -p 2222
+```
+
+Clients can either:
+- paste an API key at the login prompt, or
+- skip the prompt by sending `ONYX_API_KEY` over SSH:
+
+```shell
+export ONYX_API_KEY=your-key
+ssh -o SendEnv=ONYX_API_KEY your-host -p 2222
+```
+
+Useful hardening flags:
+- `--idle-timeout` (default `15m`)
+- `--max-session-timeout` (default `8h`)
+- `--rate-limit-per-minute` (default `20`)
+- `--rate-limit-burst` (default `40`)
+
 ## Commands
 
 | Command | Description |
@@ -70,6 +95,7 @@ onyx-cli agents --json
 | `chat` | Launch the interactive chat TUI (default) |
 | `ask` | Ask a one-shot question (non-interactive) |
 | `agents` | List available agents |
+| `serve` | Serve the interactive chat TUI over SSH |
 | `configure` | Configure server URL and API key |
 | `validate-config` | Validate configuration and test connection |
 

cli/cmd/root.go

@@ -32,6 +32,7 @@ func Execute() error {
 	rootCmd.AddCommand(newAgentsCmd())
 	rootCmd.AddCommand(newConfigureCmd())
 	rootCmd.AddCommand(newValidateConfigCmd())
+	rootCmd.AddCommand(newServeCmd())
 
 	// Default command is chat
 	rootCmd.RunE = chatCmd.RunE

cli/cmd/serve.go

@@ -0,0 +1,401 @@
+package cmd
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"io"
+	"net"
+	"os"
+	"os/signal"
+	"strings"
+	"sync"
+	"syscall"
+	"time"
+
+	tea "github.com/charmbracelet/bubbletea"
+	"github.com/charmbracelet/log"
+	"github.com/charmbracelet/ssh"
+	"github.com/charmbracelet/wish"
+	"github.com/charmbracelet/wish/activeterm"
+	"github.com/charmbracelet/wish/bubbletea"
+	"github.com/charmbracelet/wish/logging"
+	"github.com/charmbracelet/wish/ratelimiter"
+	"github.com/onyx-dot-app/onyx/cli/internal/api"
+	"github.com/onyx-dot-app/onyx/cli/internal/config"
+	"github.com/onyx-dot-app/onyx/cli/internal/tui"
+	"github.com/spf13/cobra"
+	"golang.org/x/time/rate"
+)
+
+var sessionAPIKeys sync.Map
+
+const (
+	defaultServeIdleTimeout        = 15 * time.Minute
+	defaultServeMaxSessionTimeout  = 8 * time.Hour
+	defaultServeRateLimitPerMinute = 20
+	defaultServeRateLimitBurst     = 40
+	defaultServeRateLimitCacheSize = 4096
+	maxAPIKeyLength                = 256
+	apiKeyValidationTimeout        = 15 * time.Second
+)
+
+var errAPIKeyTooLong = fmt.Errorf("API key is too long (max %d characters)", maxAPIKeyLength)
+
+// sshWriter wraps an io.Writer and tracks the first write error, allowing
+// a sequence of writes to be checked once at the end.
+type sshWriter struct {
+	w   io.Writer
+	err error
+}
+
+func (w *sshWriter) print(s string) {
+	if w.err == nil {
+		_, w.err = fmt.Fprint(w.w, s)
+	}
+}
+
+func (w *sshWriter) printf(format string, args ...any) {
+	if w.err == nil {
+		_, w.err = fmt.Fprintf(w.w, format, args...)
+	}
+}
+
+func sessionEnv(s ssh.Session, key string) string {
+	prefix := key + "="
+	for _, env := range s.Environ() {
+		if strings.HasPrefix(env, prefix) {
+			return env[len(prefix):]
+		}
+	}
+	return ""
+}
+
+func validateAPIKey(serverURL string, apiKey string) error {
+	trimmedKey := strings.TrimSpace(apiKey)
+	if len(trimmedKey) > maxAPIKeyLength {
+		return errAPIKeyTooLong
+	}
+
+	cfg := config.OnyxCliConfig{
+		ServerURL: serverURL,
+		APIKey:    trimmedKey,
+	}
+	client := api.NewClient(cfg)
+	ctx, cancel := context.WithTimeout(context.Background(), apiKeyValidationTimeout)
+	defer cancel()
+	return client.TestConnection(ctx)
+}
+
+// readMasked reads input from the SSH session one chunk at a time, echoing '•'
+// for each printable character. Handles backspace and skips escape sequences.
+func readMasked(w *sshWriter, s ssh.Session) (string, error) {
+	var key []byte
+	buf := make([]byte, 4096)
+	inEscape := false
+
+	for {
+		n, err := s.Read(buf)
+		if err != nil {
+			return "", err
+		}
+		for i := 0; i < n; i++ {
+			b := buf[i]
+
+			if inEscape {
+				if (b >= 'A' && b <= 'Z') || (b >= 'a' && b <= 'z') || b == '~' {
+					inEscape = false
+				}
+				continue
+			}
+
+			switch {
+			case b == 27: // ESC
+				inEscape = true
+			case b == '\r' || b == '\n':
+				return string(key), nil
+			case b == 3: // Ctrl+C
+				return "", fmt.Errorf("interrupted")
+			case b == 4: // Ctrl+D
+				return "", fmt.Errorf("quit")
+			case b == 127 || b == 8: // Backspace / Delete
+				if len(key) > 0 {
+					key = key[:len(key)-1]
+					w.print("\b \b")
+				}
+			case b >= 32 && b < 127: // printable ASCII
+				if len(key) >= maxAPIKeyLength {
+					return "", errAPIKeyTooLong
+				}
+				key = append(key, b)
+				w.print("•")
+			}
+
+			if w.err != nil {
+				return "", w.err
+			}
+		}
+	}
+}
+
+// promptAPIKey shows a login screen and reads the API key interactively.
+// Validates the key before returning. Loops on failure so the user can retry.
+func promptAPIKey(s ssh.Session, serverURL string) (string, error) {
+	settingsURL := strings.TrimRight(serverURL, "/") + "/app/settings/accounts-access"
+	w := &sshWriter{w: s}
+
+	w.print("\r\n")
+	w.print("  \x1b[1;35mOnyx CLI\x1b[0m\r\n")
+	w.printf("  \x1b[90m%s\x1b[0m\r\n", serverURL)
+	w.print("\r\n")
+	w.print("  Generate an API key at:\r\n")
+	w.printf("  \x1b[4;34m%s\x1b[0m\r\n", settingsURL)
+	w.print("\r\n")
+	w.print("  \x1b[90mTip: skip this prompt by passing your key via SSH:\x1b[0m\r\n")
+	w.print("  \x1b[90m  export ONYX_API_KEY=<key>\x1b[0m\r\n")
+	w.print("  \x1b[90m  ssh -o SendEnv=ONYX_API_KEY <host> -p <port>\x1b[0m\r\n")
+	w.print("\r\n")
+
+	if w.err != nil {
+		return "", w.err
+	}
+
+	for {
+		w.print("  API Key: ")
+		if w.err != nil {
+			return "", w.err
+		}
+
+		key, err := readMasked(w, s)
+		if err != nil {
+			if errors.Is(err, errAPIKeyTooLong) {
+				w.print("\r\n")
+				w.printf("  \x1b[33m%s\x1b[0m\r\n\r\n", errAPIKeyTooLong.Error())
+				if w.err != nil {
+					return "", w.err
+				}
+				continue
+			}
+			w.print("\r\n")
+			return "", err
+		}
+		w.print("\r\n")
+
+		key = strings.TrimSpace(key)
+		if key == "" {
+			w.print("  \x1b[33mNo key entered.\x1b[0m\r\n\r\n")
+			if w.err != nil {
+				return "", w.err
+			}
+			continue
+		}
+
+		w.print("  \x1b[90mValidating…\x1b[0m")
+		if w.err != nil {
+			return "", w.err
+		}
+
+		err = validateAPIKey(serverURL, key)
+
+		w.print("\r\x1b[2K") // clear "Validating…" line
+
+		if err != nil {
+			w.printf("  \x1b[1;31m%s\x1b[0m\r\n\r\n", err.Error())
+			if w.err != nil {
+				return "", w.err
+			}
+			continue
+		}
+
+		w.print("  \x1b[32mAuthenticated.\x1b[0m\r\n\r\n")
+		return key, w.err
+	}
+}
+
+// authMiddleware prompts for an API key (or reads it from the session env)
+// before handing off to the next handler (bubbletea).
+func authMiddleware(serverCfg config.OnyxCliConfig) wish.Middleware {
+	return func(next ssh.Handler) ssh.Handler {
+		return func(s ssh.Session) {
+			w := &sshWriter{w: s}
+			apiKey := strings.TrimSpace(sessionEnv(s, config.EnvAPIKey))
+
+			if apiKey != "" {
+				if err := validateAPIKey(serverCfg.ServerURL, apiKey); err != nil {
+					w.print("\r\n")
+					w.print("  \x1b[33mThe ONYX_API_KEY provided via SSH environment is invalid.\x1b[0m\r\n")
+					w.printf("  \x1b[1;31m%s\x1b[0m\r\n\r\n", err.Error())
+					if w.err != nil {
+						return
+					}
+					apiKey = ""
+				}
+			}
+
+			if apiKey == "" {
+				var err error
+				apiKey, err = promptAPIKey(s, serverCfg.ServerURL)
+				if err != nil {
+					return
+				}
+			}
+
+			sessionAPIKeys.Store(s, apiKey)
+			defer sessionAPIKeys.Delete(s)
+			next(s)
+		}
+	}
+}
+
+func newServeCmd() *cobra.Command {
+	var (
+		host              string
+		port              int
+		keyPath           string
+		idleTimeout       time.Duration
+		maxSessionTimeout time.Duration
+		rateLimitPerMin   int
+		rateLimitBurst    int
+		rateLimitCache    int
+	)
+
+	cmd := &cobra.Command{
+		Use:   "serve",
+		Short: "Serve the Onyx TUI over SSH",
+		Long: `Start an SSH server that presents the interactive Onyx chat TUI to
+connecting clients. Each SSH session gets its own independent TUI instance.
+
+Clients are prompted for their Onyx API key on connect. The key can also be
+provided via the ONYX_API_KEY environment variable to skip the prompt:
+
+  ssh -o SendEnv=ONYX_API_KEY host -p port
+
+The server URL is taken from the server operator's config. The server
+auto-generates an Ed25519 host key on first run if the key file does not
+already exist.
+
+Example:
+  onyx-cli serve --port 2222
+  ssh localhost -p 2222`,
+		RunE: func(cmd *cobra.Command, args []string) error {
+			serverCfg := config.Load()
+			if serverCfg.ServerURL == "" {
+				return fmt.Errorf("server URL is not configured; run 'onyx-cli configure' first")
+			}
+			if rateLimitPerMin <= 0 {
+				return fmt.Errorf("--rate-limit-per-minute must be > 0")
+			}
+			if rateLimitBurst <= 0 {
+				return fmt.Errorf("--rate-limit-burst must be > 0")
+			}
+			if rateLimitCache <= 0 {
+				return fmt.Errorf("--rate-limit-cache must be > 0")
+			}
+
+			addr := net.JoinHostPort(host, fmt.Sprintf("%d", port))
+			connectionLimiter := ratelimiter.NewRateLimiter(
+				rate.Limit(float64(rateLimitPerMin)/60.0),
+				rateLimitBurst,
+				rateLimitCache,
+			)
+
+			handler := func(s ssh.Session) (tea.Model, []tea.ProgramOption) {
+				apiKey := ""
+				if v, ok := sessionAPIKeys.Load(s); ok {
+					apiKey = v.(string)
+				}
+
+				cfg := config.OnyxCliConfig{
+					ServerURL:      serverCfg.ServerURL,
+					APIKey:         apiKey,
+					DefaultAgentID: serverCfg.DefaultAgentID,
+				}
+
+				m := tui.NewModel(cfg)
+				return m, []tea.ProgramOption{
+					tea.WithAltScreen(),
+					tea.WithMouseCellMotion(),
+				}
+			}
+
+			serverOptions := []ssh.Option{
+				wish.WithAddress(addr),
+				wish.WithHostKeyPath(keyPath),
+				wish.WithMiddleware(
+					bubbletea.Middleware(handler),
+					authMiddleware(serverCfg),
+					activeterm.Middleware(),
+					ratelimiter.Middleware(connectionLimiter),
+					logging.Middleware(),
+				),
+			}
+			if idleTimeout > 0 {
+				serverOptions = append(serverOptions, wish.WithIdleTimeout(idleTimeout))
+			}
+			if maxSessionTimeout > 0 {
+				serverOptions = append(serverOptions, wish.WithMaxTimeout(maxSessionTimeout))
+			}
+
+			s, err := wish.NewServer(serverOptions...)
+			if err != nil {
+				return fmt.Errorf("could not create SSH server: %w", err)
+			}
+
+			done := make(chan os.Signal, 1)
+			signal.Notify(done, os.Interrupt, syscall.SIGTERM)
+
+			log.Info("Starting Onyx SSH server", "addr", addr)
+			log.Info("Connect with", "cmd", fmt.Sprintf("ssh %s -p %d", host, port))
+
+			go func() {
+				if err := s.ListenAndServe(); err != nil && !errors.Is(err, ssh.ErrServerClosed) {
+					log.Error("SSH server failed", "error", err)
+					done <- nil
+				}
+			}()
+
+			<-done
+			log.Info("Shutting down SSH server")
+			ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+			defer cancel()
+			return s.Shutdown(ctx)
+		},
+	}
+
+	cmd.Flags().StringVar(&host, "host", "localhost", "Host address to bind to")
+	cmd.Flags().IntVarP(&port, "port", "p", 2222, "Port to listen on")
+	cmd.Flags().StringVar(&keyPath, "host-key", ".ssh/onyx_serve_ed25519",
+		"Path to SSH host key (auto-generated if missing)")
+	cmd.Flags().DurationVar(
+		&idleTimeout,
+		"idle-timeout",
+		defaultServeIdleTimeout,
+		"Disconnect idle clients after this duration (set 0 to disable)",
+	)
+	cmd.Flags().DurationVar(
+		&maxSessionTimeout,
+		"max-session-timeout",
+		defaultServeMaxSessionTimeout,
+		"Maximum lifetime of a client session (set 0 to disable)",
+	)
+	cmd.Flags().IntVar(
+		&rateLimitPerMin,
+		"rate-limit-per-minute",
+		defaultServeRateLimitPerMinute,
+		"Per-IP connection rate limit (new sessions per minute)",
+	)
+	cmd.Flags().IntVar(
+		&rateLimitBurst,
+		"rate-limit-burst",
+		defaultServeRateLimitBurst,
+		"Per-IP burst limit for connection attempts",
+	)
+	cmd.Flags().IntVar(
+		&rateLimitCache,
+		"rate-limit-cache",
+		defaultServeRateLimitCacheSize,
+		"Maximum number of IP limiter entries tracked in memory",
+	)
+
+	return cmd
+}

cli/go.mod

@@ -3,43 +3,60 @@ module github.com/onyx-dot-app/onyx/cli
 go 1.26.0
 
 require (
-	github.com/charmbracelet/bubbles v0.20.0
-	github.com/charmbracelet/bubbletea v1.3.4
-	github.com/charmbracelet/glamour v0.8.0
-	github.com/charmbracelet/lipgloss v1.1.0
-	github.com/spf13/cobra v1.9.1
-	golang.org/x/term v0.30.0
+	github.com/charmbracelet/bubbles v1.0.0
+	github.com/charmbracelet/bubbletea v1.3.10
+	github.com/charmbracelet/glamour v0.10.0
+	github.com/charmbracelet/lipgloss v1.1.1-0.20250404203927-76690c660834
+	github.com/charmbracelet/log v0.4.2
+	github.com/charmbracelet/ssh v0.0.0-20250826160808-ebfa259c7309
+	github.com/charmbracelet/wish v1.4.7
+	github.com/spf13/cobra v1.10.2
+	golang.org/x/term v0.40.0
 	golang.org/x/text v0.34.0
+	golang.org/x/time v0.14.0
 )
 
 require (
-	github.com/alecthomas/chroma/v2 v2.14.0 // indirect
+	github.com/alecthomas/chroma/v2 v2.23.1 // indirect
+	github.com/anmitsu/go-shlex v0.0.0-20200514113438-38f4b401e2be // indirect
 	github.com/atotto/clipboard v0.1.4 // indirect
 	github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect
 	github.com/aymerick/douceur v0.2.0 // indirect
-	github.com/charmbracelet/colorprofile v0.2.3-0.20250311203215-f60798e515dc // indirect
-	github.com/charmbracelet/x/ansi v0.8.0 // indirect
-	github.com/charmbracelet/x/cellbuf v0.0.13-0.20250311204145-2c3ea96c31dd // indirect
-	github.com/charmbracelet/x/term v0.2.1 // indirect
-	github.com/dlclark/regexp2 v1.11.0 // indirect
+	github.com/charmbracelet/colorprofile v0.4.2 // indirect
+	github.com/charmbracelet/keygen v0.5.4 // indirect
+	github.com/charmbracelet/x/ansi v0.11.6 // indirect
+	github.com/charmbracelet/x/cellbuf v0.0.15 // indirect
+	github.com/charmbracelet/x/conpty v0.2.0 // indirect
+	github.com/charmbracelet/x/exp/slice v0.0.0-20260305213658-fe36e8c10185 // indirect
+	github.com/charmbracelet/x/input v0.3.7 // indirect
+	github.com/charmbracelet/x/term v0.2.2 // indirect
+	github.com/charmbracelet/x/termios v0.1.1 // indirect
+	github.com/charmbracelet/x/windows v0.2.2 // indirect
+	github.com/clipperhouse/displaywidth v0.11.0 // indirect
+	github.com/clipperhouse/uax29/v2 v2.7.0 // indirect
+	github.com/creack/pty v1.1.24 // indirect
+	github.com/dlclark/regexp2 v1.11.5 // indirect
 	github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f // indirect
+	github.com/go-logfmt/logfmt v0.6.1 // indirect
 	github.com/gorilla/css v1.0.1 // indirect
+	github.com/hashicorp/golang-lru/v2 v2.0.7 // indirect
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect
-	github.com/lucasb-eyer/go-colorful v1.2.0 // indirect
+	github.com/lucasb-eyer/go-colorful v1.3.0 // indirect
 	github.com/mattn/go-isatty v0.0.20 // indirect
 	github.com/mattn/go-localereader v0.0.1 // indirect
-	github.com/mattn/go-runewidth v0.0.16 // indirect
+	github.com/mattn/go-runewidth v0.0.20 // indirect
 	github.com/microcosm-cc/bluemonday v1.0.27 // indirect
 	github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 // indirect
 	github.com/muesli/cancelreader v0.2.2 // indirect
 	github.com/muesli/reflow v0.3.0 // indirect
 	github.com/muesli/termenv v0.16.0 // indirect
 	github.com/rivo/uniseg v0.4.7 // indirect
-	github.com/spf13/pflag v1.0.6 // indirect
+	github.com/spf13/pflag v1.0.10 // indirect
 	github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e // indirect
-	github.com/yuin/goldmark v1.7.4 // indirect
-	github.com/yuin/goldmark-emoji v1.0.3 // indirect
-	golang.org/x/net v0.38.0 // indirect
-	golang.org/x/sync v0.19.0 // indirect
-	golang.org/x/sys v0.31.0 // indirect
+	github.com/yuin/goldmark v1.7.16 // indirect
+	github.com/yuin/goldmark-emoji v1.0.6 // indirect
+	golang.org/x/crypto v0.48.0 // indirect
+	golang.org/x/exp v0.0.0-20260218203240-3dfff04db8fa // indirect
+	golang.org/x/net v0.51.0 // indirect
+	golang.org/x/sys v0.41.0 // indirect
 )

cli/go.sum

@@ -1,55 +1,89 @@
-github.com/alecthomas/assert/v2 v2.7.0 h1:QtqSACNS3tF7oasA8CU6A6sXZSBDqnm7RfpLl9bZqbE=
-github.com/alecthomas/assert/v2 v2.7.0/go.mod h1:Bze95FyfUr7x34QZrjL+XP+0qgp/zg8yS+TtBj1WA3k=
-github.com/alecthomas/chroma/v2 v2.14.0 h1:R3+wzpnUArGcQz7fCETQBzO5n9IMNi13iIs46aU4V9E=
-github.com/alecthomas/chroma/v2 v2.14.0/go.mod h1:QolEbTfmUHIMVpBqxeDnNBj2uoeI4EbYP4i6n68SG4I=
-github.com/alecthomas/repr v0.4.0 h1:GhI2A8MACjfegCPVq9f1FLvIBS+DrQ2KQBFZP1iFzXc=
-github.com/alecthomas/repr v0.4.0/go.mod h1:Fr0507jx4eOXV7AlPV6AVZLYrLIuIeSOWtW57eE/O/4=
+github.com/alecthomas/assert/v2 v2.11.0 h1:2Q9r3ki8+JYXvGsDyBXwH3LcJ+WK5D0gc5E8vS6K3D0=
+github.com/alecthomas/assert/v2 v2.11.0/go.mod h1:Bze95FyfUr7x34QZrjL+XP+0qgp/zg8yS+TtBj1WA3k=
+github.com/alecthomas/chroma/v2 v2.23.1 h1:nv2AVZdTyClGbVQkIzlDm/rnhk1E9bU9nXwmZ/Vk/iY=
+github.com/alecthomas/chroma/v2 v2.23.1/go.mod h1:NqVhfBR0lte5Ouh3DcthuUCTUpDC9cxBOfyMbMQPs3o=
+github.com/alecthomas/repr v0.5.2 h1:SU73FTI9D1P5UNtvseffFSGmdNci/O6RsqzeXJtP0Qs=
+github.com/alecthomas/repr v0.5.2/go.mod h1:Fr0507jx4eOXV7AlPV6AVZLYrLIuIeSOWtW57eE/O/4=
+github.com/anmitsu/go-shlex v0.0.0-20200514113438-38f4b401e2be h1:9AeTilPcZAjCFIImctFaOjnTIavg87rW78vTPkQqLI8=
+github.com/anmitsu/go-shlex v0.0.0-20200514113438-38f4b401e2be/go.mod h1:ySMOLuWl6zY27l47sB3qLNK6tF2fkHG55UZxx8oIVo4=
 github.com/atotto/clipboard v0.1.4 h1:EH0zSVneZPSuFR11BlR9YppQTVDbh5+16AmcJi4g1z4=
 github.com/atotto/clipboard v0.1.4/go.mod h1:ZY9tmq7sm5xIbd9bOK4onWV4S6X0u6GY7Vn0Yu86PYI=
 github.com/aymanbagabas/go-osc52/v2 v2.0.1 h1:HwpRHbFMcZLEVr42D4p7XBqjyuxQH5SMiErDT4WkJ2k=
 github.com/aymanbagabas/go-osc52/v2 v2.0.1/go.mod h1:uYgXzlJ7ZpABp8OJ+exZzJJhRNQ2ASbcXHWsFqH8hp8=
-github.com/aymanbagabas/go-udiff v0.2.0 h1:TK0fH4MteXUDspT88n8CKzvK0X9O2xu9yQjWpi6yML8=
-github.com/aymanbagabas/go-udiff v0.2.0/go.mod h1:RE4Ex0qsGkTAJoQdQQCA0uG+nAzJO/pI/QwceO5fgrA=
+github.com/aymanbagabas/go-udiff v0.3.1 h1:LV+qyBQ2pqe0u42ZsUEtPiCaUoqgA9gYRDs3vj1nolY=
+github.com/aymanbagabas/go-udiff v0.3.1/go.mod h1:G0fsKmG+P6ylD0r6N/KgQD/nWzgfnl8ZBcNLgcbrw8E=
 github.com/aymerick/douceur v0.2.0 h1:Mv+mAeH1Q+n9Fr+oyamOlAkUNPWPlA8PPGR0QAaYuPk=
 github.com/aymerick/douceur v0.2.0/go.mod h1:wlT5vV2O3h55X9m7iVYN0TBM0NH/MmbLnd30/FjWUq4=
-github.com/charmbracelet/bubbles v0.20.0 h1:jSZu6qD8cRQ6k9OMfR1WlM+ruM8fkPWkHvQWD9LIutE=
-github.com/charmbracelet/bubbles v0.20.0/go.mod h1:39slydyswPy+uVOHZ5x/GjwVAFkCsV8IIVy+4MhzwwU=
-github.com/charmbracelet/bubbletea v1.3.4 h1:kCg7B+jSCFPLYRA52SDZjr51kG/fMUEoPoZrkaDHyoI=
-github.com/charmbracelet/bubbletea v1.3.4/go.mod h1:dtcUCyCGEX3g9tosuYiut3MXgY/Jsv9nKVdibKKRRXo=
-github.com/charmbracelet/colorprofile v0.2.3-0.20250311203215-f60798e515dc h1:4pZI35227imm7yK2bGPcfpFEmuY1gc2YSTShr4iJBfs=
-github.com/charmbracelet/colorprofile v0.2.3-0.20250311203215-f60798e515dc/go.mod h1:X4/0JoqgTIPSFcRA/P6INZzIuyqdFY5rm8tb41s9okk=
-github.com/charmbracelet/glamour v0.8.0 h1:tPrjL3aRcQbn++7t18wOpgLyl8wrOHUEDS7IZ68QtZs=
-github.com/charmbracelet/glamour v0.8.0/go.mod h1:ViRgmKkf3u5S7uakt2czJ272WSg2ZenlYEZXT2x7Bjw=
-github.com/charmbracelet/lipgloss v1.1.0 h1:vYXsiLHVkK7fp74RkV7b2kq9+zDLoEU4MZoFqR/noCY=
-github.com/charmbracelet/lipgloss v1.1.0/go.mod h1:/6Q8FR2o+kj8rz4Dq0zQc3vYf7X+B0binUUBwA0aL30=
-github.com/charmbracelet/x/ansi v0.8.0 h1:9GTq3xq9caJW8ZrBTe0LIe2fvfLR/bYXKTx2llXn7xE=
-github.com/charmbracelet/x/ansi v0.8.0/go.mod h1:wdYl/ONOLHLIVmQaxbIYEC/cRKOQyjTkowiI4blgS9Q=
-github.com/charmbracelet/x/cellbuf v0.0.13-0.20250311204145-2c3ea96c31dd h1:vy0GVL4jeHEwG5YOXDmi86oYw2yuYUGqz6a8sLwg0X8=
-github.com/charmbracelet/x/cellbuf v0.0.13-0.20250311204145-2c3ea96c31dd/go.mod h1:xe0nKWGd3eJgtqZRaN9RjMtK7xUYchjzPr7q6kcvCCs=
-github.com/charmbracelet/x/exp/golden v0.0.0-20240815200342-61de596daa2b h1:MnAMdlwSltxJyULnrYbkZpp4k58Co7Tah3ciKhSNo0Q=
-github.com/charmbracelet/x/exp/golden v0.0.0-20240815200342-61de596daa2b/go.mod h1:wDlXFlCrmJ8J+swcL/MnGUuYnqgQdW9rhSD61oNMb6U=
-github.com/charmbracelet/x/term v0.2.1 h1:AQeHeLZ1OqSXhrAWpYUtZyX1T3zVxfpZuEQMIQaGIAQ=
-github.com/charmbracelet/x/term v0.2.1/go.mod h1:oQ4enTYFV7QN4m0i9mzHrViD7TQKvNEEkHUMCmsxdUg=
+github.com/charmbracelet/bubbles v1.0.0 h1:12J8/ak/uCZEMQ6KU7pcfwceyjLlWsDLAxB5fXonfvc=
+github.com/charmbracelet/bubbles v1.0.0/go.mod h1:9d/Zd5GdnauMI5ivUIVisuEm3ave1XwXtD1ckyV6r3E=
+github.com/charmbracelet/bubbletea v1.3.10 h1:otUDHWMMzQSB0Pkc87rm691KZ3SWa4KUlvF9nRvCICw=
+github.com/charmbracelet/bubbletea v1.3.10/go.mod h1:ORQfo0fk8U+po9VaNvnV95UPWA1BitP1E0N6xJPlHr4=
+github.com/charmbracelet/colorprofile v0.4.2 h1:BdSNuMjRbotnxHSfxy+PCSa4xAmz7szw70ktAtWRYrY=
+github.com/charmbracelet/colorprofile v0.4.2/go.mod h1:0rTi81QpwDElInthtrQ6Ni7cG0sDtwAd4C4le060fT8=
+github.com/charmbracelet/glamour v0.10.0 h1:MtZvfwsYCx8jEPFJm3rIBFIMZUfUJ765oX8V6kXldcY=
+github.com/charmbracelet/glamour v0.10.0/go.mod h1:f+uf+I/ChNmqo087elLnVdCiVgjSKWuXa/l6NU2ndYk=
+github.com/charmbracelet/keygen v0.5.4 h1:XQYgf6UEaTGgQSSmiPpIQ78WfseNQp4Pz8N/c1OsrdA=
+github.com/charmbracelet/keygen v0.5.4/go.mod h1:t4oBRr41bvK7FaJsAaAQhhkUuHslzFXVjOBwA55CZNM=
+github.com/charmbracelet/lipgloss v1.1.1-0.20250404203927-76690c660834 h1:ZR7e0ro+SZZiIZD7msJyA+NjkCNNavuiPBLgerbOziE=
+github.com/charmbracelet/lipgloss v1.1.1-0.20250404203927-76690c660834/go.mod h1:aKC/t2arECF6rNOnaKaVU6y4t4ZeHQzqfxedE/VkVhA=
+github.com/charmbracelet/log v0.4.2 h1:hYt8Qj6a8yLnvR+h7MwsJv/XvmBJXiueUcI3cIxsyig=
+github.com/charmbracelet/log v0.4.2/go.mod h1:qifHGX/tc7eluv2R6pWIpyHDDrrb/AG71Pf2ysQu5nw=
+github.com/charmbracelet/ssh v0.0.0-20250826160808-ebfa259c7309 h1:dCVbCRRtg9+tsfiTXTp0WupDlHruAXyp+YoxGVofHHc=
+github.com/charmbracelet/ssh v0.0.0-20250826160808-ebfa259c7309/go.mod h1:R9cISUs5kAH4Cq/rguNbSwcR+slE5Dfm8FEs//uoIGE=
+github.com/charmbracelet/wish v1.4.7 h1:O+jdLac3s6GaqkOHHSwezejNK04vl6VjO1A+hl8J8Yc=
+github.com/charmbracelet/wish v1.4.7/go.mod h1:OBZ8vC62JC5cvbxJLh+bIWtG7Ctmct+ewziuUWK+G14=
+github.com/charmbracelet/x/ansi v0.11.6 h1:GhV21SiDz/45W9AnV2R61xZMRri5NlLnl6CVF7ihZW8=
+github.com/charmbracelet/x/ansi v0.11.6/go.mod h1:2JNYLgQUsyqaiLovhU2Rv/pb8r6ydXKS3NIttu3VGZQ=
+github.com/charmbracelet/x/cellbuf v0.0.15 h1:ur3pZy0o6z/R7EylET877CBxaiE1Sp1GMxoFPAIztPI=
+github.com/charmbracelet/x/cellbuf v0.0.15/go.mod h1:J1YVbR7MUuEGIFPCaaZ96KDl5NoS0DAWkskup+mOY+Q=
+github.com/charmbracelet/x/conpty v0.2.0 h1:eKtA2hm34qNfgJCDp/M6Dc0gLy7e07YEK4qAdNGOvVY=
+github.com/charmbracelet/x/conpty v0.2.0/go.mod h1:fexgUnVrZgw8scD49f6VSi0Ggj9GWYIrpedRthAwW/8=
+github.com/charmbracelet/x/exp/golden v0.0.0-20241011142426-46044092ad91 h1:payRxjMjKgx2PaCWLZ4p3ro9y97+TVLZNaRZgJwSVDQ=
+github.com/charmbracelet/x/exp/golden v0.0.0-20241011142426-46044092ad91/go.mod h1:wDlXFlCrmJ8J+swcL/MnGUuYnqgQdW9rhSD61oNMb6U=
+github.com/charmbracelet/x/exp/slice v0.0.0-20260305213658-fe36e8c10185 h1:bloHJLweYZeIkBVgi8AF94DrTdx3eoEB57VOpFuFi3U=
+github.com/charmbracelet/x/exp/slice v0.0.0-20260305213658-fe36e8c10185/go.mod h1:vqEfX6xzqW1pKKZUUiFOKg0OQ7bCh54Q2vR/tserrRA=
+github.com/charmbracelet/x/input v0.3.7 h1:UzVbkt1vgM9dBQ+K+uRolBlN6IF2oLchmPKKo/aucXo=
+github.com/charmbracelet/x/input v0.3.7/go.mod h1:ZSS9Cia6Cycf2T6ToKIOxeTBTDwl25AGwArJuGaOBH8=
+github.com/charmbracelet/x/term v0.2.2 h1:xVRT/S2ZcKdhhOuSP4t5cLi5o+JxklsoEObBSgfgZRk=
+github.com/charmbracelet/x/term v0.2.2/go.mod h1:kF8CY5RddLWrsgVwpw4kAa6TESp6EB5y3uxGLeCqzAI=
+github.com/charmbracelet/x/termios v0.1.1 h1:o3Q2bT8eqzGnGPOYheoYS8eEleT5ZVNYNy8JawjaNZY=
+github.com/charmbracelet/x/termios v0.1.1/go.mod h1:rB7fnv1TgOPOyyKRJ9o+AsTU/vK5WHJ2ivHeut/Pcwo=
+github.com/charmbracelet/x/windows v0.2.2 h1:IofanmuvaxnKHuV04sC0eBy/smG6kIKrWG2/jYn2GuM=
+github.com/charmbracelet/x/windows v0.2.2/go.mod h1:/8XtdKZzedat74NQFn0NGlGL4soHB0YQZrETF96h75k=
+github.com/clipperhouse/displaywidth v0.11.0 h1:lBc6kY44VFw+TDx4I8opi/EtL9m20WSEFgwIwO+UVM8=
+github.com/clipperhouse/displaywidth v0.11.0/go.mod h1:bkrFNkf81G8HyVqmKGxsPufD3JhNl3dSqnGhOoSD/o0=
+github.com/clipperhouse/uax29/v2 v2.7.0 h1:+gs4oBZ2gPfVrKPthwbMzWZDaAFPGYK72F0NJv2v7Vk=
+github.com/clipperhouse/uax29/v2 v2.7.0/go.mod h1:EFJ2TJMRUaplDxHKj1qAEhCtQPW2tJSwu5BF98AuoVM=
 github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g=
-github.com/dlclark/regexp2 v1.11.0 h1:G/nrcoOa7ZXlpoa/91N3X7mM3r8eIlMBBJZvsz/mxKI=
-github.com/dlclark/regexp2 v1.11.0/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8=
+github.com/creack/pty v1.1.24 h1:bJrF4RRfyJnbTJqzRLHzcGaZK1NeM5kTC9jGgovnR1s=
+github.com/creack/pty v1.1.24/go.mod h1:08sCNb52WyoAwi2QDyzUCTgcvVFhUzewun7wtTfvcwE=
+github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
+github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/dlclark/regexp2 v1.11.5 h1:Q/sSnsKerHeCkc/jSTNq1oCm7KiVgUMZRDUoRu0JQZQ=
+github.com/dlclark/regexp2 v1.11.5/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8=
 github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f h1:Y/CXytFA4m6baUTXGLOoWe4PQhGxaX0KpnayAqC48p4=
 github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f/go.mod h1:vw97MGsxSvLiUE2X8qFplwetxpGLQrlU1Q9AUEIzCaM=
+github.com/go-logfmt/logfmt v0.6.1 h1:4hvbpePJKnIzH1B+8OR/JPbTx37NktoI9LE2QZBBkvE=
+github.com/go-logfmt/logfmt v0.6.1/go.mod h1:EV2pOAQoZaT1ZXZbqDl5hrymndi4SY9ED9/z6CO0XAk=
+github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
+github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
 github.com/gorilla/css v1.0.1 h1:ntNaBIghp6JmvWnxbZKANoLyuXTPZ4cAMlo6RyhlbO8=
 github.com/gorilla/css v1.0.1/go.mod h1:BvnYkspnSzMmwRK+b8/xgNPLiIuNZr6vbZBTPQ2A3b0=
+github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k=
+github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM=
 github.com/hexops/gotextdiff v1.0.3 h1:gitA9+qJrrTCsiCl7+kh75nPqQt1cx4ZkudSTLoUqJM=
 github.com/hexops/gotextdiff v1.0.3/go.mod h1:pSWU5MAI3yDq+fZBTazCSJysOMbxWL1BSow5/V2vxeg=
 github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
 github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
-github.com/lucasb-eyer/go-colorful v1.2.0 h1:1nnpGOrhyZZuNyfu1QjKiUICQ74+3FNCN69Aj6K7nkY=
-github.com/lucasb-eyer/go-colorful v1.2.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0=
+github.com/lucasb-eyer/go-colorful v1.3.0 h1:2/yBRLdWBZKrf7gB40FoiKfAWYQ0lqNcbuQwVHXptag=
+github.com/lucasb-eyer/go-colorful v1.3.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0=
 github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
 github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
 github.com/mattn/go-localereader v0.0.1 h1:ygSAOl7ZXTx4RdPYinUpg6W99U8jWvWi9Ye2JC/oIi4=
 github.com/mattn/go-localereader v0.0.1/go.mod h1:8fBrzywKY7BI3czFoHkuzRoWE9C+EiG4R1k4Cjx5p88=
 github.com/mattn/go-runewidth v0.0.12/go.mod h1:RAqKPSqVFrSLVXbA8x7dzmKdmGzieGRCM46jaSJTDAk=
-github.com/mattn/go-runewidth v0.0.16 h1:E5ScNMtiwvlvB5paMFdw9p4kSQzbXFikJ5SQO6TULQc=
-github.com/mattn/go-runewidth v0.0.16/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w=
+github.com/mattn/go-runewidth v0.0.20 h1:WcT52H91ZUAwy8+HUkdM3THM6gXqXuLJi9O3rjcQQaQ=
+github.com/mattn/go-runewidth v0.0.20/go.mod h1:XBkDxAl56ILZc9knddidhrOlY5R/pDhgLpndooCuJAs=
 github.com/microcosm-cc/bluemonday v1.0.27 h1:MpEUotklkwCSLeH+Qdx1VJgNqLlpY2KXwXFM08ygZfk=
 github.com/microcosm-cc/bluemonday v1.0.27/go.mod h1:jFi9vgW+H7c3V0lb6nR74Ib/DIB5OBs92Dimizgw2cA=
 github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 h1:ZK8zHtRHOkbHy6Mmr5D264iyp3TiX5OmNcI5cIARiQI=
@@ -60,35 +94,45 @@ github.com/muesli/reflow v0.3.0 h1:IFsN6K9NfGtjeggFP+68I4chLZV2yIKsXJFNZ+eWh6s=
 github.com/muesli/reflow v0.3.0/go.mod h1:pbwTDkVPibjO2kyvBQRBxTWEEGDGq0FlB1BIKtnHY/8=
 github.com/muesli/termenv v0.16.0 h1:S5AlUN9dENB57rsbnkPyfdGuWIlkmzJjbFf0Tf5FWUc=
 github.com/muesli/termenv v0.16.0/go.mod h1:ZRfOIKPFDYQoDFF4Olj7/QJbW60Ol/kL1pU3VfY/Cnk=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/rivo/uniseg v0.1.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc=
 github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc=
 github.com/rivo/uniseg v0.4.7 h1:WUdvkW8uEhrYfLC4ZzdpI2ztxP1I582+49Oc5Mq64VQ=
 github.com/rivo/uniseg v0.4.7/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88=
 github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
-github.com/spf13/cobra v1.9.1 h1:CXSaggrXdbHK9CF+8ywj8Amf7PBRmPCOJugH954Nnlo=
-github.com/spf13/cobra v1.9.1/go.mod h1:nDyEzZ8ogv936Cinf6g1RU9MRY64Ir93oCnqb9wxYW0=
-github.com/spf13/pflag v1.0.6 h1:jFzHGLGAlb3ruxLB8MhbI6A8+AQX/2eW4qeyNZXNp2o=
-github.com/spf13/pflag v1.0.6/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
+github.com/spf13/cobra v1.10.2 h1:DMTTonx5m65Ic0GOoRY2c16WCbHxOOw6xxezuLaBpcU=
+github.com/spf13/cobra v1.10.2/go.mod h1:7C1pvHqHw5A4vrJfjNwvOdzYu0Gml16OCs2GRiTUUS4=
+github.com/spf13/pflag v1.0.9/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
+github.com/spf13/pflag v1.0.10 h1:4EBh2KAYBwaONj6b2Ye1GiHfwjqyROoF4RwYO+vPwFk=
+github.com/spf13/pflag v1.0.10/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
+github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA=
+github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
 github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e h1:JVG44RsyaB9T2KIHavMF/ppJZNG9ZpyihvCd0w101no=
 github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e/go.mod h1:RbqR21r5mrJuqunuUZ/Dhy/avygyECGrLceyNeo4LiM=
-github.com/yuin/goldmark v1.7.1/go.mod h1:uzxRWxtg69N339t3louHJ7+O03ezfj6PlliRlaOzY1E=
-github.com/yuin/goldmark v1.7.4 h1:BDXOHExt+A7gwPCJgPIIq7ENvceR7we7rOS9TNoLZeg=
-github.com/yuin/goldmark v1.7.4/go.mod h1:uzxRWxtg69N339t3louHJ7+O03ezfj6PlliRlaOzY1E=
-github.com/yuin/goldmark-emoji v1.0.3 h1:aLRkLHOuBR2czCY4R8olwMjID+tENfhyFDMCRhbIQY4=
-github.com/yuin/goldmark-emoji v1.0.3/go.mod h1:tTkZEbwu5wkPmgTcitqddVxY9osFZiavD+r4AzQrh1U=
-golang.org/x/exp v0.0.0-20220909182711-5c715a9e8561 h1:MDc5xs78ZrZr3HMQugiXOAkSZtfTpbJLDr/lwfgO53E=
-golang.org/x/exp v0.0.0-20220909182711-5c715a9e8561/go.mod h1:cyybsKvd6eL0RnXn6p/Grxp8F5bW7iYuBgsNCOHpMYE=
-golang.org/x/net v0.38.0 h1:vRMAPTMaeGqVhG5QyLJHqNDwecKTomGeqbnfZyKlBI8=
-golang.org/x/net v0.38.0/go.mod h1:ivrbrMbzFq5J41QOQh0siUuly180yBYtLp+CKbEaFx8=
+github.com/yuin/goldmark v1.7.16 h1:n+CJdUxaFMiDUNnWC3dMWCIQJSkxH4uz3ZwQBkAlVNE=
+github.com/yuin/goldmark v1.7.16/go.mod h1:ip/1k0VRfGynBgxOz0yCqHrbZXhcjxyuS66Brc7iBKg=
+github.com/yuin/goldmark-emoji v1.0.6 h1:QWfF2FYaXwL74tfGOW5izeiZepUDroDJfWubQI9HTHs=
+github.com/yuin/goldmark-emoji v1.0.6/go.mod h1:ukxJDKFpdFb5x0a5HqbdlcKtebh086iJpI31LTKmWuA=
+go.yaml.in/yaml/v3 v3.0.4/go.mod h1:DhzuOOF2ATzADvBadXxruRBLzYTpT36CKvDb3+aBEFg=
+golang.org/x/crypto v0.48.0 h1:/VRzVqiRSggnhY7gNRxPauEQ5Drw9haKdM0jqfcCFts=
+golang.org/x/crypto v0.48.0/go.mod h1:r0kV5h3qnFPlQnBSrULhlsRfryS2pmewsg+XfMgkVos=
+golang.org/x/exp v0.0.0-20260218203240-3dfff04db8fa h1:Zt3DZoOFFYkKhDT3v7Lm9FDMEV06GpzjG2jrqW+QTE0=
+golang.org/x/exp v0.0.0-20260218203240-3dfff04db8fa/go.mod h1:K79w1Vqn7PoiZn+TkNpx3BUWUQksGO3JcVX6qIjytmA=
+golang.org/x/net v0.51.0 h1:94R/GTO7mt3/4wIKpcR5gkGmRLOuE/2hNGeWq/GBIFo=
+golang.org/x/net v0.51.0/go.mod h1:aamm+2QF5ogm02fjy5Bb7CQ0WMt1/WVM7FtyaTLlA9Y=
 golang.org/x/sync v0.19.0 h1:vV+1eWNmZ5geRlYjzm2adRgW2/mcpevXNg50YZtPCE4=
 golang.org/x/sync v0.19.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI=
 golang.org/x/sys v0.0.0-20210809222454-d867a43fc93e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.31.0 h1:ioabZlmFYtWhL+TRYpcnNlLwhyxaM9kWTDEmfnprqik=
-golang.org/x/sys v0.31.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k=
-golang.org/x/term v0.30.0 h1:PQ39fJZ+mfadBm0y5WlL4vlM7Sx1Hgf13sMIY2+QS9Y=
-golang.org/x/term v0.30.0/go.mod h1:NYYFdzHoI5wRh/h5tDMdMqCqPJZEuNqVR5xJLd/n67g=
+golang.org/x/sys v0.41.0 h1:Ivj+2Cp/ylzLiEU89QhWblYnOE9zerudt9Ftecq2C6k=
+golang.org/x/sys v0.41.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=
+golang.org/x/term v0.40.0 h1:36e4zGLqU4yhjlmxEaagx2KuYbJq3EwY8K943ZsHcvg=
+golang.org/x/term v0.40.0/go.mod h1:w2P8uVp06p2iyKKuvXIm7N/y0UCRt3UfJTfZ7oOpglM=
 golang.org/x/text v0.34.0 h1:oL/Qq0Kdaqxa1KbNeMKwQq0reLCCaFtqu2eNuSeNHbk=
 golang.org/x/text v0.34.0/go.mod h1:homfLqTYRFyVYemLBFl5GgL/DWEiH5wcsQ5gSh1yziA=
+golang.org/x/time v0.14.0 h1:MRx4UaLrDotUKUdCIqzPC48t1Y9hANFKIRpNx+Te8PI=
+golang.org/x/time v0.14.0/go.mod h1:eL/Oa2bBBK0TkX57Fyni+NgnyQQN4LitPmob2Hjnqw4=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

docker-bake.hcl

@@ -18,10 +18,6 @@ variable "INTEGRATION_REPOSITORY" {
   default = "onyxdotapp/onyx-integration"
 }
 
-variable "CLI_REPOSITORY" {
-  default = "onyxdotapp/onyx-cli"
-}
-
 variable "TAG" {
   default = "latest"
 }
@@ -68,13 +64,3 @@ target "integration" {
 
   tags      = ["${INTEGRATION_REPOSITORY}:${TAG}"]
 }
-
-target "cli" {
-  context    = "cli"
-  dockerfile = "Dockerfile"
-
-  cache-from = ["type=registry,ref=${CLI_REPOSITORY}:latest"]
-  cache-to   = ["type=inline"]
-
-  tags      = ["${CLI_REPOSITORY}:${TAG}"]
-}

web/src/app/app/message/messageComponents/renderMessageComponent.tsx

@@ -1,14 +1,11 @@
 import React, { JSX, memo } from "react";
 import {
   ChatPacket,
-  CODE_INTERPRETER_TOOL_TYPES,
   ImageGenerationToolPacket,
   Packet,
   PacketType,
   ReasoningPacket,
-  SearchToolStart,
   StopReason,
-  ToolCallArgumentDelta,
 } from "../../services/streamingModels";
 import {
   FullChatState,
@@ -29,6 +26,7 @@ import { DeepResearchPlanRenderer } from "./timeline/renderers/deepresearch/Deep
 import { ResearchAgentRenderer } from "./timeline/renderers/deepresearch/ResearchAgentRenderer";
 import { WebSearchToolRenderer } from "./timeline/renderers/search/WebSearchToolRenderer";
 import { InternalSearchToolRenderer } from "./timeline/renderers/search/InternalSearchToolRenderer";
+import { SearchToolStart } from "../../services/streamingModels";
 
 // Different types of chat packets using discriminated unions
 interface GroupedPackets {
@@ -58,12 +56,7 @@ function isImageToolPacket(packet: Packet) {
 }
 
 function isPythonToolPacket(packet: Packet) {
-  return (
-    packet.obj.type === PacketType.PYTHON_TOOL_START ||
-    (packet.obj.type === PacketType.TOOL_CALL_ARGUMENT_DELTA &&
-      (packet.obj as ToolCallArgumentDelta).tool_type ===
-        CODE_INTERPRETER_TOOL_TYPES.PYTHON)
-  );
+  return packet.obj.type === PacketType.PYTHON_TOOL_START;
 }
 
 function isCustomToolPacket(packet: Packet) {

web/src/app/app/message/messageComponents/timeline/hooks/packetProcessor.ts

@@ -10,8 +10,6 @@ import {
   Stop,
   ImageGenerationToolDelta,
   MessageStart,
-  ToolCallArgumentDelta,
-  CODE_INTERPRETER_TOOL_TYPES,
 } from "@/app/app/services/streamingModels";
 import { CitationMap } from "@/app/app/interfaces";
 import { OnyxDocument } from "@/lib/search/interfaces";
@@ -140,7 +138,6 @@ const CONTENT_PACKET_TYPES_SET = new Set<PacketType>([
   PacketType.SEARCH_TOOL_START,
   PacketType.IMAGE_GENERATION_TOOL_START,
   PacketType.PYTHON_TOOL_START,
-  PacketType.TOOL_CALL_ARGUMENT_DELTA,
   PacketType.CUSTOM_TOOL_START,
   PacketType.FILE_READER_START,
   PacketType.FETCH_TOOL_START,
@@ -152,16 +149,9 @@ const CONTENT_PACKET_TYPES_SET = new Set<PacketType>([
 ]);
 
 function hasContentPackets(packets: Packet[]): boolean {
-  return packets.some((packet) => {
-    const type = packet.obj.type as PacketType;
-    if (type === PacketType.TOOL_CALL_ARGUMENT_DELTA) {
-      return (
-        (packet.obj as ToolCallArgumentDelta).tool_type ===
-        CODE_INTERPRETER_TOOL_TYPES.PYTHON
-      );
-    }
-    return CONTENT_PACKET_TYPES_SET.has(type);
-  });
+  return packets.some((packet) =>
+    CONTENT_PACKET_TYPES_SET.has(packet.obj.type as PacketType)
+  );
 }
 
 /**

web/src/app/app/message/messageComponents/timeline/packetHelpers.ts

@@ -1,13 +1,6 @@
-import {
-  CODE_INTERPRETER_TOOL_TYPES,
-  Packet,
-  PacketType,
-  ToolCallArgumentDelta,
-} from "@/app/app/services/streamingModels";
+import { Packet, PacketType } from "@/app/app/services/streamingModels";
 
-// Packet types with renderers supporting collapsed streaming mode.
-// TOOL_CALL_ARGUMENT_DELTA is intentionally excluded here because it requires
-// a tool_type check — it's handled separately in stepSupportsCollapsedStreaming.
+// Packet types with renderers supporting collapsed streaming mode
 export const COLLAPSED_STREAMING_PACKET_TYPES = new Set<PacketType>([
   PacketType.SEARCH_TOOL_START,
   PacketType.FETCH_TOOL_START,
@@ -28,13 +21,7 @@ export const isSearchToolPackets = (packets: Packet[]): boolean =>
 
 // Check if packets belong to a python tool
 export const isPythonToolPackets = (packets: Packet[]): boolean =>
-  packets.some(
-    (p) =>
-      p.obj.type === PacketType.PYTHON_TOOL_START ||
-      (p.obj.type === PacketType.TOOL_CALL_ARGUMENT_DELTA &&
-        (p.obj as ToolCallArgumentDelta).tool_type ===
-          CODE_INTERPRETER_TOOL_TYPES.PYTHON)
-  );
+  packets.some((p) => p.obj.type === PacketType.PYTHON_TOOL_START);
 
 // Check if packets belong to reasoning
 export const isReasoningPackets = (packets: Packet[]): boolean =>
@@ -42,12 +29,8 @@ export const isReasoningPackets = (packets: Packet[]): boolean =>
 
 // Check if step supports collapsed streaming rendering mode
 export const stepSupportsCollapsedStreaming = (packets: Packet[]): boolean =>
-  packets.some(
-    (p) =>
-      COLLAPSED_STREAMING_PACKET_TYPES.has(p.obj.type as PacketType) ||
-      (p.obj.type === PacketType.TOOL_CALL_ARGUMENT_DELTA &&
-        (p.obj as ToolCallArgumentDelta).tool_type ===
-          CODE_INTERPRETER_TOOL_TYPES.PYTHON)
+  packets.some((p) =>
+    COLLAPSED_STREAMING_PACKET_TYPES.has(p.obj.type as PacketType)
   );
 
 // Check if packets have content worth rendering in collapsed streaming mode.
@@ -84,13 +67,7 @@ export const stepHasCollapsedStreamingContent = (
   // Python tool renders code/output from the start packet onward
   if (
     packetTypes.has(PacketType.PYTHON_TOOL_START) ||
-    packetTypes.has(PacketType.PYTHON_TOOL_DELTA) ||
-    packets.some(
-      (p) =>
-        p.obj.type === PacketType.TOOL_CALL_ARGUMENT_DELTA &&
-        (p.obj as ToolCallArgumentDelta).tool_type ===
-          CODE_INTERPRETER_TOOL_TYPES.PYTHON
-    )
+    packetTypes.has(PacketType.PYTHON_TOOL_DELTA)
   ) {
     return true;
   }

web/src/app/app/message/messageComponents/timeline/renderers/code/PythonToolRenderer.tsx

@@ -4,9 +4,7 @@ import {
   PythonToolPacket,
   PythonToolStart,
   PythonToolDelta,
-  ToolCallArgumentDelta,
   SectionEnd,
-  CODE_INTERPRETER_TOOL_TYPES,
 } from "@/app/app/services/streamingModels";
 import {
   MessageRenderer,
@@ -41,18 +39,6 @@ function HighlightedPythonCode({ code }: { code: string }) {
 
 // Helper function to construct current Python execution state
 function constructCurrentPythonState(packets: PythonToolPacket[]) {
-  // Accumulate streaming code from argument deltas (arrives before PythonToolStart)
-  const streamingCode = packets
-    .filter(
-      (packet) =>
-        packet.obj.type === PacketType.TOOL_CALL_ARGUMENT_DELTA &&
-        (packet.obj as ToolCallArgumentDelta).tool_type ===
-          CODE_INTERPRETER_TOOL_TYPES.PYTHON
-    )
-    .map((packet) =>
-      String((packet.obj as ToolCallArgumentDelta).argument_deltas.code ?? "")
-    )
-    .join("");
   const pythonStart = packets.find(
     (packet) => packet.obj.type === PacketType.PYTHON_TOOL_START
   )?.obj as PythonToolStart | null;
@@ -65,8 +51,7 @@ function constructCurrentPythonState(packets: PythonToolPacket[]) {
       packet.obj.type === PacketType.ERROR
   )?.obj as SectionEnd | null;
 
-  // Use complete code from PythonToolStart if available, else use streamed code.
-  const code = pythonStart?.code || streamingCode;
+  const code = pythonStart?.code || "";
   const stdout = pythonDeltas
     .map((delta) => delta?.stdout || "")
     .filter((s) => s)
@@ -76,7 +61,6 @@ function constructCurrentPythonState(packets: PythonToolPacket[]) {
     .filter((s) => s)
     .join("");
   const fileIds = pythonDeltas.flatMap((delta) => delta?.file_ids || []);
-  const isStreaming = !pythonStart && streamingCode.length > 0;
   const isExecuting = pythonStart && !pythonEnd;
   const isComplete = pythonStart && pythonEnd;
   const hasError = stderr.length > 0;
@@ -86,7 +70,6 @@ function constructCurrentPythonState(packets: PythonToolPacket[]) {
     stdout,
     stderr,
     fileIds,
-    isStreaming,
     isExecuting,
     isComplete,
     hasError,
@@ -99,16 +82,8 @@ export const PythonToolRenderer: MessageRenderer<PythonToolPacket, {}> = ({
   renderType,
   children,
 }) => {
-  const {
-    code,
-    stdout,
-    stderr,
-    fileIds,
-    isStreaming,
-    isExecuting,
-    isComplete,
-    hasError,
-  } = constructCurrentPythonState(packets);
+  const { code, stdout, stderr, fileIds, isExecuting, isComplete, hasError } =
+    constructCurrentPythonState(packets);
 
   useEffect(() => {
     if (isComplete) {
@@ -117,9 +92,6 @@ export const PythonToolRenderer: MessageRenderer<PythonToolPacket, {}> = ({
   }, [isComplete, onComplete]);
 
   const status = useMemo(() => {
-    if (isStreaming) {
-      return "Writing code...";
-    }
     if (isExecuting) {
       return "Executing Python code...";
     }
@@ -130,13 +102,13 @@ export const PythonToolRenderer: MessageRenderer<PythonToolPacket, {}> = ({
       return "Python execution completed";
     }
     return "Python execution";
-  }, [isStreaming, isComplete, isExecuting, hasError]);
+  }, [isComplete, isExecuting, hasError]);
 
   // Shared content for all states - used by both FULL and compact modes
   const content = (
     <div className="flex flex-col mb-1 space-y-2">
-      {/* Loading indicator when streaming or executing */}
-      {(isStreaming || isExecuting) && (
+      {/* Loading indicator when executing */}
+      {isExecuting && (
         <div className="flex items-center gap-2 text-sm text-muted-foreground">
           <div className="flex gap-0.5">
             <div className="w-1 h-1 bg-current rounded-full animate-pulse"></div>
@@ -149,7 +121,7 @@ export const PythonToolRenderer: MessageRenderer<PythonToolPacket, {}> = ({
               style={{ animationDelay: "0.2s" }}
             ></div>
           </div>
-          <span>{isStreaming ? "Writing code..." : "Running code..."}</span>
+          <span>Running code...</span>
         </div>
       )}
 

web/src/app/app/services/packetUtils.ts

@@ -16,7 +16,6 @@ export function isToolPacket(
     PacketType.SEARCH_TOOL_DOCUMENTS_DELTA,
     PacketType.PYTHON_TOOL_START,
     PacketType.PYTHON_TOOL_DELTA,
-    PacketType.TOOL_CALL_ARGUMENT_DELTA,
     PacketType.CUSTOM_TOOL_START,
     PacketType.CUSTOM_TOOL_DELTA,
     PacketType.FILE_READER_START,

web/src/app/app/services/streamingModels.ts

@@ -27,9 +27,6 @@ export enum PacketType {
   FETCH_TOOL_URLS = "open_url_urls",
   FETCH_TOOL_DOCUMENTS = "open_url_documents",
 
-  // Tool call argument delta (streams tool args before tool executes)
-  TOOL_CALL_ARGUMENT_DELTA = "tool_call_argument_delta",
-
   // Custom tool packets
   CUSTOM_TOOL_START = "custom_tool_start",
   CUSTOM_TOOL_DELTA = "custom_tool_delta",
@@ -62,10 +59,6 @@ export enum PacketType {
   INTERMEDIATE_REPORT_CITED_DOCS = "intermediate_report_cited_docs",
 }
 
-export const CODE_INTERPRETER_TOOL_TYPES = {
-  PYTHON: "python",
-} as const;
-
 // Basic Message Packets
 export interface MessageStart extends BaseObj {
   id: string;
@@ -156,13 +149,6 @@ export interface PythonToolDelta extends BaseObj {
   file_ids: string[];
 }
 
-export interface ToolCallArgumentDelta extends BaseObj {
-  type: "tool_call_argument_delta";
-  tool_type: string;
-  tool_id: string;
-  argument_deltas: Record<string, unknown>;
-}
-
 export interface FetchToolStart extends BaseObj {
   type: "open_url_start";
 }
@@ -308,7 +294,6 @@ export type ImageGenerationToolObj =
 export type PythonToolObj =
   | PythonToolStart
   | PythonToolDelta
-  | ToolCallArgumentDelta
   | SectionEnd
   | PacketError;
 export type FetchToolObj =