mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
0805e8e410
## Description
`headroom/providers/openai.py::_get_encoding` calls
`tiktoken.get_encoding` directly. tiktoken downloads missing
vocabularies via `requests.get` with **no timeout**, so on a network
that blackholes the vocab CDN (corporate firewall, SSL-intercepting
proxy), whichever thread first counts tokens for an OpenAI model — proxy
startup included — blocks indefinitely.
This is the provider-path hole left by #956: the tokenizer registry
already routes through a bounded loader
(`headroom/tokenizers/tiktoken_counter.py`, worker-thread load +
`HEADROOM_TIKTOKEN_LOAD_TIMEOUT_SECONDS`, default 10s) and falls back to
estimation, but the OpenAI provider path never got the same treatment.
Observed in production (Headroom Desktop fleet, Sentry): a proxy that
never finished booting, with a faulthandler dump wedged in
`tiktoken/registry.py` `get_encoding` on the main thread, reached from
the `headroom` CLI entrypoint via click. The desktop app now also
pre-seeds a persistent `TIKTOKEN_CACHE_DIR`, but the unbounded load
affects every deployment of the proxy, so it should be fixed here too.
Follow-up to #956.
## 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
- `_get_encoding` now routes through the bounded `load_encoding` from
`headroom.tokenizers.tiktoken_counter` instead of calling
`tiktoken.get_encoding` directly, so a stalled vocab download raises
`TiktokenLoadError` after the timeout instead of hanging the calling
thread.
- `OpenAIProvider.get_token_counter` catches `TiktokenLoadError` and
falls back to `EstimatingTokenCounter`, cached per model so later
requests never re-block on the same failed download — mirroring
`TokenizerRegistry._create_tiktoken`.
- `TIKTOKEN_AVAILABLE` uses `importlib.util.find_spec` (the module-level
`import tiktoken` became unused; same pattern as `LITELLM_AVAILABLE`).
- Two regression tests (`TestGuardedEncodingLoad`) covering the
bounded-raise path and the cached estimation fallback.
## Testing
- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [x] Type checking passes (`mypy headroom`)
- [x] New tests added for new functionality
- [ ] Manual testing performed
### Test Output
```text
$ uv run --frozen --extra dev pytest tests/test_providers/ tests/test_tokenizers/
======================= 117 passed, 4 warnings in 27.95s =======================
$ uvx ruff check headroom/providers/openai.py tests/test_providers/test_openai.py
All checks passed!
$ uvx ruff format --check headroom/providers/openai.py tests/test_providers/test_openai.py
2 files already formatted
$ uv run --frozen --extra dev mypy headroom/providers/openai.py
Success: no issues found in 1 source file
```
## Real Behavior Proof
- Environment: macOS (arm64), Python 3.12, uv-managed venv, branch off
`upstream/main` (
|
||
|---|---|---|
| .. | ||
| __init__.py | ||
| test_anthropic.py | ||
| test_cohere.py | ||
| test_deepseek.py | ||
| test_openai.py | ||
| test_universal.py | ||