mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
4 commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
Rod Boev
|
551f473e04
|
fix(proxy): accept Codex websocket before upstream retries (#2203)
## Description Codex Desktop could abandon Headroom's local `/v1/responses` WebSocket handshake before Headroom's upstream retry strategy had a chance to recover. The ChatGPT-auth path waited for an upstream opening handshake with a minimum 30-second timeout before sending the local 101, while the reported Codex Desktop handshake expired after about 34 seconds. This change accepts validated ChatGPT-auth Codex WebSockets before opening the upstream connection, then keeps the existing upstream retries and HTTP fallback behind the established local session. API-key sessions retain connect-before-accept behavior so upstream `x-codex-*` headers can still be attached to their client-facing 101. The change is scoped to the pre-101 timing failure and does not address the separate large-context streaming investigation in #1944. Closes #2184 ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - Accept ChatGPT-auth Codex WebSocket clients before the upstream connect and retry loop. - Preserve API-key connect-before-accept ordering and upstream `x-codex-*` handshake-header forwarding. - Keep upstream retry, relay, usage-state refresh, and WebSocket-to-HTTP fallback behavior after the local 101. - Add a deterministic regression that blocks the first upstream opening handshake and proves the local acceptance deadline is independent of it. ## Testing - [x] Unit tests pass (`uv run pytest tests/test_openai_codex_ws_lifecycle.py -q`) - [x] Linting passes (`uv run ruff check headroom/proxy/handlers/openai.py tests/test_openai_codex_ws_lifecycle.py`) - [ ] Type checking passes (`uv run mypy headroom`) - [x] New tests added for new functionality when applicable - [ ] Manual testing performed ### Test Output ```text uv run pytest tests/test_openai_codex_ws_lifecycle.py -q 28 passed in 2.02s uv run ruff check headroom/proxy/handlers/openai.py tests/test_openai_codex_ws_lifecycle.py All checks passed! uv run ruff format headroom/proxy/handlers/openai.py tests/test_openai_codex_ws_lifecycle.py --check 2 files already formatted ``` ## Real Behavior Proof - Environment: Python 3.12, synced development worktree, local fake Codex client and upstream WebSocket, no live provider - Exact command / steps: Run `uv run pytest tests/test_openai_codex_ws_lifecycle.py::test_chatgpt_ws_accepts_before_stalled_upstream_connect -q`; the fake upstream blocks its first opening handshake while the client enforces a bounded local-accept deadline. - Observed result: `1 passed in 0.46s`; the ChatGPT-auth client receives its local 101 before the blocked upstream connect is released, and the handler continues into its existing upstream recovery path. - Not tested: live Codex Desktop pre-turn compaction against ChatGPT subscription infrastructure ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes - [x] I have updated the CHANGELOG.md if applicable ## Additional Notes `CHANGELOG.md` is unchanged because the release pipeline generates it from conventional commits. No user documentation changes are required; the handler comments and ordered-flow docstring are updated with the auth-mode-specific behavior. The broader #1944 large-context disconnect surface remains out of scope. --------- Co-authored-by: JerrettDavis <mxjerrett@gmail.com> |
||
|
yiihao
|
0c7087539d
|
fix: add DeepSeek V4/V3.2/R1 tokenizer mappings and context limits (#912)
The tokenizer registry routed deepseek-v4-pro, deepseek-v4-flash,
deepseek-chat, deepseek-reasoner, and other modern DeepSeek models
to the 2023-era deepseek-llm-7b-base tokenizer via prefix fallback.
This caused token counts off by 30-50%, broken context-limit detection
(V4-Pro supports 1M but got 32K), and inaccurate savings reports.
## Fix
3 files, +43/-2:
- **huggingface.py**: 16 new MODEL_TO_TOKENIZER entries with verified
HuggingFace IDs (deepseek-ai/DeepSeek-V4-Pro, V4-Flash, V3.2,
V3-0324, R1, R1-0528, Reasoner, Chat, Coder-V2, etc.)
- **openai_compatible.py**: 17 new _DEFAULT_CONTEXT_LIMITS entries
(V4-Pro/Flash -> 1M, R1/Reasoner -> 131K, V3 -> 128K, etc.)
- **openai.py**: 8 new _CONTEXT_LIMITS entries for LiteLLM-fallback.
Existing mappings untouched (backward compatible).
## Real behavior proof
- **Setup**: Windows 11, Python 3.13.14, headroom-ai 0.2.15 wheel +
source checkout at v0.24.0. No Rust extension built (headroom._core
unavailable). Touched files are at parity with v0.24.0.
- **Steps after patch**:
```
python3 -c "
from headroom.tokenizers.huggingface import get_tokenizer_name
for m in
['deepseek-v4-pro','deepseek-chat','deepseek-reasoner','deepseek-v4-flash']:
print(f'{m} -> {get_tokenizer_name(m)}')
from headroom.tokenizers.registry import get_tokenizer
for m in ['deepseek-v4-pro','deepseek-chat','deepseek-reasoner']:
print(f'{m}: {get_tokenizer(m)}')
"
```
- **Observed result**:
```
deepseek-v4-pro -> deepseek-ai/DeepSeek-V4-Pro
deepseek-v4-flash -> deepseek-ai/DeepSeek-V4-Flash
deepseek-chat -> deepseek-ai/DeepSeek-V3
deepseek-reasoner -> deepseek-ai/DeepSeek-R1
```
Previously ALL resolved to deepseek-ai/deepseek-llm-7b-base.
TokenizerRegistry routes correctly. Context limits verified
(1M / 131K / 128K). compress() import smoke-tested OK.
- **Not tested**: full proxy e2e with a live DeepSeek API key
(no available key). HuggingFace AutoTokenizer download confirmed
for V4-Pro/V3/R1 but produced GBK decode errors from hf_hub on
this zh-CN Windows locale during config fetch -- a separate
huggingface_hub issue unrelated to this change.
<!-- headroom-maintainer-template-completion:start -->
## Description
This PR prepares `fix: add DeepSeek V4/V3.2/R1 tokenizer mappings and
context limits` for review by documenting the intended change,
validation evidence, and remaining merge-readiness context.
Linked issues: None declared.
## Type of Change
- [x] Bug fix
- [ ] New feature
- [ ] Documentation
- [ ] Refactor
- [ ] Tests only
## Changes Made
- Commit: fix: add DeepSeek V4/V3.2/R1 tokenizer mappings and context
limits
- Touches `headroom/providers/openai.py`
- Touches `headroom/providers/openai_compatible.py`
- Touches `headroom/tokenizers/huggingface.py`
## Testing
- [x] GitHub checks reviewed
- [x] Metadata/template validation
- [ ] Local functional testing
### Test Output
```text
gh pr view 912 --repo chopratejas/headroom --json statusCheckRollup
- PR Governance / label: SUCCESS
- external / GitGuardian Security Checks: SUCCESS
```
## Real Behavior Proof
- Environment: GitHub PR metadata and checks for `chopratejas/headroom`
PR #912.
- Exact command / steps: Reviewed PR title, commits, changed files,
linked issues, labels, and check rollup; appended this maintainer
template completion block without replacing the author's original
description.
- Observed result: PR body now contains all required governance
sections, checked readiness fields, and a non-placeholder validation
evidence block.
- Not tested: This pass updated PR metadata only; code validation
remains represented by the linked GitHub checks and any author-provided
evidence above.
## Review Readiness
- [x] I have performed a self-review
- [x] This PR is ready for human review
<!-- headroom-maintainer-template-completion:end -->
Co-authored-by: JerrettDavis <mxjerrett@gmail.com>
|
||
|
Abhay Singh
|
e0232df9b4
|
fix(tokenizers): resolve HF tokenizer names by most-specific prefix (#2096)
## Description
`get_tokenizer_name` can pick the wrong tokenizer for a versioned model,
which silently produces wrong token counts.
For a model that isn't a literal key in `MODEL_TO_TOKENIZER`, it falls
back to prefix matching:
```python
for key, value in MODEL_TO_TOKENIZER.items():
if model_lower.startswith(key):
return value
```
That returns the first key the model merely *starts with*, in
dict-insertion order. The table lists short family keys before their
more-specific siblings — `"qwen"` (→ `Qwen/Qwen-7B`) appears before
`"qwen2"`/`"qwen2-7b"`/`"qwen2.5"`. So
`get_tokenizer_name("qwen2-7b-instruct")` matches `"qwen"` first and
returns the **Qwen1** tokenizer, not Qwen2. Qwen1 and Qwen2 have
different vocabularies, so every `count_text`/`count_messages` for that
model is off. `qwen2.5-*` and `deepseek-v2.x` are mis-resolved the same
way.
The sibling tiktoken resolver already documents and guards this exact
pitfall — `get_encoding_for_model` uses an explicit most-specific-first
prefix list with a comment that scanning "for the first key that merely
starts with the prefix is order-dependent and wrong." The HuggingFace
resolver is the one that still scans insertion order.
## Fix
Match the **longest** (most-specific) prefix instead of the first in
insertion order:
```python
for key in sorted(MODEL_TO_TOKENIZER, key=len, reverse=True):
if model_lower.startswith(key):
return MODEL_TO_TOKENIZER[key]
```
Direct-key lookups and the shorter-family fallback (e.g. `deepseek-chat`
→ `deepseek-ai/deepseek-llm-7b-base`) are unchanged.
Closes #
## Type of Change
- [x] Bug fix (non-breaking change that fixes an issue)
- [ ] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [ ] Documentation update
- [ ] Performance improvement
- [ ] Code refactoring (no functional changes)
## Changes Made
- `headroom/tokenizers/huggingface.py`: `get_tokenizer_name` prefix
matching now iterates keys longest-first and returns the most-specific
match.
- `tests/test_huggingface_tokenizer_timeout.py`: add
`test_get_tokenizer_name_prefers_most_specific_prefix`.
- `CHANGELOG.md`: Bug Fixes entry.
## Testing
- [ ] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [ ] Type checking passes (`mypy headroom`)
- [x] New tests added for new functionality
- [ ] Manual testing performed
### Test Output
```text
$ uvx ruff@0.15.17 check headroom/tokenizers/huggingface.py tests/test_huggingface_tokenizer_timeout.py
All checks passed!
$ python -m py_compile headroom/tokenizers/huggingface.py tests/test_huggingface_tokenizer_timeout.py
OK
```
## Real Behavior Proof
- Environment: Windows 11, Python 3.12, `uvx ruff@0.15.17`. Importing
`headroom` pulls in the torch/transformers stack and a full `pytest`
gets OOM-killed on this box, so I verified against the real key table
with a dependency-free script that parses `MODEL_TO_TOKENIZER` out of
the source and runs both the old (insertion-order) and new
(longest-first) scans, then left the full pytest to CI.
- Exact command / steps: resolved `qwen2-7b-instruct`, `qwen2.5-turbo`,
and `deepseek-v2.5` under both strategies, plus `deepseek-chat` as a
regression guard.
- Observed result: old scan returns `Qwen/Qwen-7B` (Qwen1) for both
qwen2 models and `deepseek-ai/deepseek-llm-7b-base` (v1) for
`deepseek-v2.5`; new scan returns `Qwen/Qwen2-7B`, `Qwen/Qwen2.5-7B`,
and `deepseek-ai/DeepSeek-V2` respectively. `deepseek-chat` resolves
identically under both (`deepseek-ai/deepseek-llm-7b-base`), so the
existing timeout test's model is unaffected.
- Not tested: loading the actual HuggingFace tokenizers
(network/`transformers`); full local `pytest` deferred to CI (OOM, per
above).
## Review Readiness
- [x] I have performed a self-review
- [x] This PR is ready for human review
## Checklist
- [x] My code follows the project's style guidelines
- [x] I have performed a self-review of my code
- [x] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [ ] New and existing unit tests pass locally with my changes
- [x] I have updated the CHANGELOG.md if applicable
## Additional Notes
The "unit tests pass locally" and "type checking" boxes are unchecked
because the full suite imports the ML stack, which I can't run in this
environment; the change is a localized swap of the prefix-scan order in
a pure function, verified by the standalone proof (run against the real
key table) and the new regression test for CI. This mirrors the
same-class fix already present in the sibling tiktoken resolver.
Co-authored-by: Tejas Chopra <chopratejas@gmail.com>
|
||
|
Parideboy
|
46d5d685d9
|
fix(proxy): bound HF tokenizer load and offload token counting off event loop (#1738)
## Description Fixes #1701. On Windows, `headroom proxy --anthropic-api-url https://api.deepseek.com/anthropic` froze: the first `/v1/messages` request took ~610s (`optimization_latency_ms=609972`) with only router/lifecycle markers, and afterwards the whole server was a zombie — `/livez`, `/readyz` and `/health` hung until the process was killed. `HEADROOM_DETECT_BACKEND=python` was already set, so this was not the #575/#845 native-detect deadlock. Root cause: DeepSeek model names route to the HuggingFace tokenizer backend (`MODEL_PATTERNS` in `headroom/tokenizers/registry.py`). `HuggingFaceTokenizer` loads lazily, so the registry's construction-time fallback never fires; the first `count_messages` calls `AutoTokenizer.from_pretrained(..., trust_remote_code=True)` — unbounded network downloads/retries — and this ran **synchronously inside the async Anthropic messages handler** (`get_tokenizer(model)` + `tokenizer.count_messages(messages)`), outside the 30s `_run_compression_in_executor` bound. huggingface_hub retry chains on a restricted network easily reach ~10 minutes, blocking the entire asyncio event loop; subsequent on-loop counting kept it pinned. tiktoken got a bounded eager load for the same bug class long ago (#956); the HF backend never did. ## Type of Change - [x] Bug fix (non-breaking change which fixes an issue) - [ ] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Refactoring (no functional changes) ## Changes Made - `headroom/tokenizers/huggingface.py`: `_load_tokenizer` now tries the local HF cache first (`local_files_only=True`, no network), then bounds the network load with `HEADROOM_HF_TOKENIZER_LOAD_TIMEOUT_SECS` (default 10s; `0` disables network loads) on a daemon thread. Timeouts/failures return `None` (cached by `lru_cache`, so the hub is probed at most once per process per tokenizer) and `count_messages` fails open to char-based estimation via the existing `_use_fallback()` path. - `headroom/proxy/handlers/anthropic.py`: new `AnthropicHandlerMixin._count_tokens_offloaded(model, messages)` runs `get_tokenizer` + `count_messages` on the compression executor bounded by `COMPRESSION_TIMEOUT_SECONDS`, failing open to `EstimatingTokenCounter` (downgrade logged once per model). Used in `handle_anthropic_messages` (the issue's hot path, both count sites) and `handle_anthropic_batch_create`; the batch path's inline `anthropic_pipeline.apply()` is now offloaded via `_run_compression_in_executor` (mirrors the #1612 image-compression offload). - `headroom/proxy/handlers/batch.py`: the two remaining inline `openai_pipeline.apply()` calls (`handle_google_batch_create`, `_compress_batch_jsonl`) are offloaded the same way; existing `except` blocks keep the pass-through fail-open semantics. - Tests: `tests/test_huggingface_tokenizer_timeout.py` (cache-first, bounded timeout, failure caching, timeout=0, fail-open estimation), `tests/test_tokenizer_count_offload.py` (wiring guards, runs on `headroom-compress` worker, event loop stays responsive during slow tokenizer work, fail-open), plus `_run_compression_in_executor` stub on the batch test double. ## Testing - [x] All existing tests pass - [x] Added new tests for the changes - [ ] Manual testing performed ``` $ python -m pytest tests/test_huggingface_tokenizer_timeout.py tests/test_tokenizer_count_offload.py tests/test_image_compression_offload.py tests/test_gemini_compression_offload.py tests/test_tokenizers tests/test_proxy_handlers_batch.py -q 50 passed $ ruff check . # No issues found $ ruff format --check . # 1043 files already formatted $ mypy headroom --ignore-missing-imports # 0 errors ``` ## Real Behavior Proof - Environment: Windows 11 Pro (10.0.26200), Python 3.13, local checkout of this branch with the Rust core built. - Exact command / steps: `python -m pytest tests/test_tokenizer_count_offload.py -q` — includes `test_count_tokens_offloaded_keeps_loop_responsive`, which reproduces the issue's mechanism: a tokenizer whose `count_messages` blocks (stand-in for the unbounded `AutoTokenizer.from_pretrained` network load) while an asyncio ticker measures event-loop liveness. Also `python -m pytest tests/test_huggingface_tokenizer_timeout.py -q` with a `from_pretrained` stub that sleeps 60s and `HEADROOM_HF_TOKENIZER_LOAD_TIMEOUT_SECS=0.2`. - Observed result: with the fix, the slow count runs on a `headroom-compress` worker thread and the loop keeps ticking (`ticks >= 5`; inline it yields ~0 — the zombie). The 60s-hung HF load unblocks at the 0.2s timeout, falls back to estimation, and the second call returns instantly (failure cached, no re-probe). All 10 new tests pass. - Not tested: live reproduction against `api.deepseek.com` from a network where HF hub downloads stall (the reporter's exact environment); actual HF vocab download timing on a healthy network. ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review Co-authored-by: Claude Fable 5 <noreply@anthropic.com> |