Skip to content

feat: add custom headers support for LLM model pool#92

Open
Cluas wants to merge 3 commits intodataelement:mainfrom
Cluas:feat/llm-model-custom-headers
Open

feat: add custom headers support for LLM model pool#92
Cluas wants to merge 3 commits intodataelement:mainfrom
Cluas:feat/llm-model-custom-headers

Conversation

@Cluas
Copy link

@Cluas Cluas commented Mar 15, 2026

Summary

Closes #91

Adds support for configuring custom HTTP headers on LLM models in the model pool. This enables gateway authentication tokens, vendor-specific auth supplements, or routing headers required by self-hosted proxies.

Changes

Backend

  • Migration — adds headers_encrypted TEXT column to llm_models (idempotent)
  • security.py — new encrypt_symmetric / decrypt_symmetric helpers (AES-GCM)
  • models/llm.pyheaders_encrypted: Mapped[str | None]
  • schemas.pyheaders: dict | None on Create/Update/Out schemas
  • enterprise.py — encrypt on write, update on edit
  • llm_client.py — all client classes and factory functions accept and merge custom headers
  • Call sites — websocket, agent_tools, task_executor, heartbeat, scheduler, supervision_reminder all decrypt and forward headers

Frontend

  • Replace raw JSON textarea with structured key-value row editor
  • Header values masked with password input

Security

Header values are encrypted at rest using AES-GCM (same key as API keys) and never returned in API responses.

@wisdomqin
Copy link
Contributor

@yaojin3616

Copy link
Collaborator

@yaojin3616 yaojin3616 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Summary: I’m requesting changes. The PR adds custom LLM headers + symmetric encryption, but there are two blocking issues and a couple of sharp edges.

Blocking

  1. Missing dependency for AES
    backend/app/core/security.py imports Crypto.Cipher.AES, but backend/pyproject.toml doesn’t include pycryptodome. This will crash at import time in
    environments without it.
    Fix: add pycryptodome or switch to cryptography (already present).
  2. Header override can silently break auth
    backend/app/services/llm_client.py merges custom_headers over default headers, which allows overriding Authorization, Content-Type, and anthropic-
    version. This is too easy to misconfigure and is a production risk.
    Fix: block overriding these keys, or only allow overrides for provider=custom.

Non-blocking (but please consider)

  • Editing headers UX is incomplete: existing headers aren’t visible and there’s no explicit “clear headers” action. Users can’t confidently manage
    what’s already stored.

Once the dependency and header-override issues are addressed, I’m good to re-review quickly.

@Cluas Cluas force-pushed the feat/llm-model-custom-headers branch from 46fe9f7 to 5d7e33b Compare March 16, 2026 16:11
@Cluas Cluas requested a review from yaojin3616 March 16, 2026 16:45
@Cluas
Copy link
Author

Cluas commented Mar 16, 2026

Summary: I’m requesting changes. The PR adds custom LLM headers + symmetric encryption, but there are two blocking issues and a couple of sharp edges.

Blocking

  1. Missing dependency for AES
    backend/app/core/security.py imports Crypto.Cipher.AES, but backend/pyproject.toml doesn’t include pycryptodome. This will crash at import time in
    environments without it.
    Fix: add pycryptodome or switch to cryptography (already present).
  2. Header override can silently break auth
    backend/app/services/llm_client.py merges custom_headers over default headers, which allows overriding Authorization, Content-Type, and anthropic-
    version. This is too easy to misconfigure and is a production risk.
    Fix: block overriding these keys, or only allow overrides for provider=custom.

Non-blocking (but please consider)

  • Editing headers UX is incomplete: existing headers aren’t visible and there’s no explicit “clear headers” action. Users can’t confidently manage
    what’s already stored.

Once the dependency and header-override issues are addressed, I’m good to re-review quickly.

Thanks for the thorough review! Here's how each issue has been addressed:

Blocking #1 — Missing AES dependency

Switched from pycryptodome (Crypto.Cipher.AES) to the cryptography package (already a project dependency via python-jose[cryptography]), using cryptography.hazmat.primitives.ciphers.aead.AESGCM. Also removed the now-unused pycryptodome entry
from pyproject.toml. No new dependency required.

Blocking #2 — Header override silently breaking auth

Added a module-level constant _PROTECTED_HEADER_KEYS in llm_client.py:

_PROTECTED_HEADER_KEYS: frozenset[str] = frozenset(
{"authorization", "x-api-key", "content-type", "anthropic-version"}
)

Both OpenAICompatibleClient._get_headers() and AnthropicClient._get_headers() now strip any custom header whose key (case-insensitive) matches this set before merging, so auth and protocol headers can never be overridden by user config.

Non-blocking — Editing headers UX

Three improvements made:

  • Existing headers are now shown when opening the edit form — the API decrypts and returns headers in the response, and the frontend maps them into the key-value row editor.
  • Values are masked with type="password" inputs (with a leave blank to keep placeholder in edit mode).
  • "Clear All" button appears whenever there are rows, giving users an explicit way to remove all stored headers.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Add custom headers support for LLM model pool

3 participants