mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
Merge remote-tracking branch 'upstream/main' into native-e2e-expansion
This commit is contained in:
commit
46f556eead
67 changed files with 5168 additions and 367 deletions
4
.github/ISSUE_TEMPLATE/config.yml
vendored
4
.github/ISSUE_TEMPLATE/config.yml
vendored
|
|
@ -1,8 +1,8 @@
|
|||
blank_issues_enabled: true
|
||||
contact_links:
|
||||
- name: Questions & Discussions
|
||||
url: https://github.com/headroom-sdk/headroom/discussions
|
||||
url: https://github.com/chopratejas/headroom/discussions
|
||||
about: Ask questions and discuss ideas in GitHub Discussions
|
||||
- name: Documentation
|
||||
url: https://headroom.dev/docs
|
||||
url: https://headroom-docs.vercel.app/docs
|
||||
about: Check out the documentation for guides and API reference
|
||||
|
|
|
|||
31
.github/PULL_REQUEST_TEMPLATE.md
vendored
31
.github/PULL_REQUEST_TEMPLATE.md
vendored
|
|
@ -1,8 +1,8 @@
|
|||
## Description
|
||||
|
||||
Brief description of changes and motivation.
|
||||
<!-- Briefly explain the change and why it is needed. -->
|
||||
|
||||
Fixes #(issue number)
|
||||
Closes #
|
||||
|
||||
## Type of Change
|
||||
|
||||
|
|
@ -15,13 +15,11 @@ Fixes #(issue number)
|
|||
|
||||
## Changes Made
|
||||
|
||||
- Change 1
|
||||
- Change 2
|
||||
- Change 3
|
||||
-
|
||||
|
||||
## Testing
|
||||
|
||||
Describe the tests you ran to verify your changes:
|
||||
<!-- Check what you actually ran, then paste the real command output below. -->
|
||||
|
||||
- [ ] Unit tests pass (`pytest`)
|
||||
- [ ] Linting passes (`ruff check .`)
|
||||
|
|
@ -29,12 +27,23 @@ Describe the tests you ran to verify your changes:
|
|||
- [ ] New tests added for new functionality
|
||||
- [ ] Manual testing performed
|
||||
|
||||
## Test Output
|
||||
### Test Output
|
||||
|
||||
```text
|
||||
# Paste relevant command output or artifact links here
|
||||
```
|
||||
# Paste relevant test output here
|
||||
pytest -v tests/test_your_feature.py
|
||||
```
|
||||
|
||||
## Real Behavior Proof
|
||||
|
||||
- Environment:
|
||||
- Exact command / steps:
|
||||
- Observed result:
|
||||
- Not tested:
|
||||
|
||||
## Review Readiness
|
||||
|
||||
- [ ] I have performed a self-review
|
||||
- [ ] This PR is ready for human review
|
||||
|
||||
## Checklist
|
||||
|
||||
|
|
@ -53,4 +62,4 @@ Add screenshots to help explain your changes.
|
|||
|
||||
## Additional Notes
|
||||
|
||||
Any additional information that reviewers should know.
|
||||
<!-- Mention any N/A checklist items, tradeoffs, follow-ups, or maintainer context. -->
|
||||
|
|
|
|||
19
.github/act/pr-governance-invalid.json
vendored
Normal file
19
.github/act/pr-governance-invalid.json
vendored
Normal file
|
|
@ -0,0 +1,19 @@
|
|||
{
|
||||
"action": "opened",
|
||||
"number": 42,
|
||||
"pull_request": {
|
||||
"number": 42,
|
||||
"draft": false,
|
||||
"title": "feat: add PR governance",
|
||||
"body": "## Description\n\nFixes #123\n",
|
||||
"user": {
|
||||
"login": "octocat"
|
||||
},
|
||||
"base": {
|
||||
"sha": "dff6a199"
|
||||
}
|
||||
},
|
||||
"repository": {
|
||||
"full_name": "JerrettDavis/headroom"
|
||||
}
|
||||
}
|
||||
19
.github/act/pr-governance-valid.json
vendored
Normal file
19
.github/act/pr-governance-valid.json
vendored
Normal file
|
|
@ -0,0 +1,19 @@
|
|||
{
|
||||
"action": "ready_for_review",
|
||||
"number": 42,
|
||||
"pull_request": {
|
||||
"number": 42,
|
||||
"draft": false,
|
||||
"title": "feat: add PR governance",
|
||||
"body": "## Description\n\nAdd a required PR governance check and commit-msg enforcement.\n\nCloses #123\n\n## Type of Change\n\n- [x] New feature (non-breaking change that adds functionality)\n\n## Changes Made\n\n- Added workflow validation for PR template completeness.\n- Added a commit-msg hook that runs commitlint locally.\n\n## Testing\n\n- [x] Unit tests pass (`pytest`)\n- [x] Manual testing performed\n\n### Test Output\n\n```text\npytest scripts/tests/test_pr_governance.py -q\n```\n\n## Real Behavior Proof\n\n- Environment: Ubuntu runner, Python 3.12\n- Exact command / steps: Opened a PR with an incomplete template, then fixed the body.\n- Observed result: The governance check failed until the template and readiness boxes were complete.\n- Not tested: Repository-level automatic Copilot rulesets.\n\n## Review Readiness\n\n- [x] I have performed a self-review\n- [x] This PR is ready for human review\n",
|
||||
"user": {
|
||||
"login": "octocat"
|
||||
},
|
||||
"base": {
|
||||
"sha": "dff6a199"
|
||||
}
|
||||
},
|
||||
"repository": {
|
||||
"full_name": "JerrettDavis/headroom"
|
||||
}
|
||||
}
|
||||
7
.github/copilot-instructions.md
vendored
Normal file
7
.github/copilot-instructions.md
vendored
Normal file
|
|
@ -0,0 +1,7 @@
|
|||
When performing a pull request review in this repository:
|
||||
|
||||
1. Treat `.github/PULL_REQUEST_TEMPLATE.md` and `CONTRIBUTING.md` as required policy, not optional guidance.
|
||||
2. Flag pull requests that do not include concrete "Real Behavior Proof" with environment, exact commands or steps, observed result, and what was not tested.
|
||||
3. Be strict about contributor verification: missing tests, missing runtime evidence, or placeholder PR text should be called out.
|
||||
4. For user-facing, release, dependency, workflow, or security-sensitive changes, prefer blocking feedback over optional suggestions.
|
||||
5. Focus on correctness, safety, and whether the PR is actually ready for human maintainer review.
|
||||
11
.github/workflows/ci.yml
vendored
11
.github/workflows/ci.yml
vendored
|
|
@ -195,8 +195,19 @@ jobs:
|
|||
run: |
|
||||
pytest tests scripts/tests \
|
||||
--splits 4 --group ${{ matrix.shard }} \
|
||||
--cov=headroom --cov-branch \
|
||||
--cov-report=xml:coverage-${{ matrix.shard }}.xml \
|
||||
--cov-report= \
|
||||
--tb=short -q
|
||||
|
||||
- name: Upload coverage shard ${{ matrix.shard }} to Codecov
|
||||
uses: codecov/codecov-action@v5
|
||||
with:
|
||||
files: coverage-${{ matrix.shard }}.xml
|
||||
flags: python
|
||||
name: python-shard-${{ matrix.shard }}
|
||||
fail_ci_if_error: true
|
||||
|
||||
test-extras:
|
||||
needs: [changes, build-wheel]
|
||||
if: needs.changes.outputs.code == 'true'
|
||||
|
|
|
|||
314
.github/workflows/pr-health.yml
vendored
314
.github/workflows/pr-health.yml
vendored
|
|
@ -1,91 +1,223 @@
|
|||
name: PR Health
|
||||
|
||||
on:
|
||||
pull_request_target:
|
||||
types: [opened, reopened, synchronize, ready_for_review]
|
||||
schedule:
|
||||
# Keep labels fresh even when base branches move or checks finish later.
|
||||
- cron: '23 14 * * 1-5'
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
issues: write
|
||||
pull-requests: write
|
||||
checks: read
|
||||
statuses: read
|
||||
|
||||
concurrency:
|
||||
group: pr-health-${{ github.event.pull_request.number || github.ref }}
|
||||
cancel-in-progress: true
|
||||
|
||||
jobs:
|
||||
label:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 10
|
||||
env:
|
||||
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
REPO: ${{ github.repository }}
|
||||
steps:
|
||||
- name: Ensure maintenance labels exist
|
||||
run: |
|
||||
set -euo pipefail
|
||||
|
||||
gh label create "status: needs rebase" \
|
||||
--repo "$REPO" \
|
||||
--color "fbca04" \
|
||||
--description "Pull request branch is behind the base branch" \
|
||||
--force
|
||||
gh label create "status: has conflicts" \
|
||||
--repo "$REPO" \
|
||||
--color "d73a4a" \
|
||||
--description "Pull request has merge conflicts with the base branch" \
|
||||
--force
|
||||
gh label create "status: ci failing" \
|
||||
--repo "$REPO" \
|
||||
--color "d73a4a" \
|
||||
--description "Required or reported CI checks are failing" \
|
||||
--force
|
||||
|
||||
- name: Label open pull requests
|
||||
run: |
|
||||
set -euo pipefail
|
||||
|
||||
if jq -e '.pull_request.number' "$GITHUB_EVENT_PATH" >/dev/null; then
|
||||
pr_numbers="$(jq -r '.pull_request.number' "$GITHUB_EVENT_PATH")"
|
||||
else
|
||||
pr_numbers="$(gh pr list --repo "$REPO" --state open --limit 100 --json number --jq '.[].number')"
|
||||
fi
|
||||
|
||||
for pr in $pr_numbers; do
|
||||
data="$(gh pr view "$pr" --repo "$REPO" \
|
||||
--json mergeStateStatus,statusCheckRollup)"
|
||||
|
||||
merge_state="$(jq -r '.mergeStateStatus // "UNKNOWN"' <<<"$data")"
|
||||
check_state="$(jq -r '
|
||||
[
|
||||
.statusCheckRollup[]
|
||||
| select((.conclusion // .state // "") as $s
|
||||
| ["FAILURE", "TIMED_OUT", "ACTION_REQUIRED", "CANCELLED", "ERROR"] | index($s))
|
||||
]
|
||||
| if length > 0 then "failing" else "passing" end
|
||||
' <<<"$data")"
|
||||
|
||||
if [[ "$merge_state" == "BEHIND" ]]; then
|
||||
gh pr edit "$pr" --repo "$REPO" --add-label "status: needs rebase"
|
||||
else
|
||||
gh pr edit "$pr" --repo "$REPO" --remove-label "status: needs rebase" || true
|
||||
fi
|
||||
|
||||
if [[ "$merge_state" == "DIRTY" ]]; then
|
||||
gh pr edit "$pr" --repo "$REPO" --add-label "status: has conflicts"
|
||||
else
|
||||
gh pr edit "$pr" --repo "$REPO" --remove-label "status: has conflicts" || true
|
||||
fi
|
||||
|
||||
if [[ "$check_state" == "failing" ]]; then
|
||||
gh pr edit "$pr" --repo "$REPO" --add-label "status: ci failing"
|
||||
else
|
||||
gh pr edit "$pr" --repo "$REPO" --remove-label "status: ci failing" || true
|
||||
fi
|
||||
done
|
||||
name: PR Governance
|
||||
|
||||
on:
|
||||
pull_request_target:
|
||||
types: [opened, edited, reopened, synchronize, ready_for_review, converted_to_draft]
|
||||
schedule:
|
||||
# Keep labels fresh even when base branches move or checks finish later.
|
||||
- cron: '23 14 * * 1-5'
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
issues: write
|
||||
pull-requests: write
|
||||
checks: read
|
||||
statuses: read
|
||||
|
||||
concurrency:
|
||||
group: pr-health-${{ github.event.pull_request.number || github.ref }}
|
||||
cancel-in-progress: true
|
||||
|
||||
jobs:
|
||||
template:
|
||||
if: github.event_name == 'pull_request_target'
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 10
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.base.sha }}
|
||||
|
||||
- name: Validate PR template
|
||||
id: validate
|
||||
run: python3 scripts/pr-governance.py --event "$GITHUB_EVENT_PATH" --report .pr-governance-report.json
|
||||
|
||||
- name: Append governance summary
|
||||
run: |
|
||||
python3 - <<'PY'
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
report = json.loads(Path(".pr-governance-report.json").read_text(encoding="utf-8"))
|
||||
summary = report["summary_markdown"].strip()
|
||||
with Path(os.environ["GITHUB_STEP_SUMMARY"]).open("a", encoding="utf-8") as handle:
|
||||
handle.write(f"{summary}\n")
|
||||
PY
|
||||
|
||||
- name: Ensure governance labels exist
|
||||
env:
|
||||
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
REPO: ${{ github.repository }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
|
||||
gh label create "status: needs author action" \
|
||||
--repo "$REPO" \
|
||||
--color "d93f0b" \
|
||||
--description "Pull request body or readiness checklist still needs author updates" \
|
||||
--force
|
||||
gh label create "status: ready for review" \
|
||||
--repo "$REPO" \
|
||||
--color "0e8a16" \
|
||||
--description "Pull request body is complete and the author marked it ready for human review" \
|
||||
--force
|
||||
|
||||
- name: Sync governance comment and labels
|
||||
if: steps.validate.outputs.is_bot_pr != 'true'
|
||||
uses: actions/github-script@v7
|
||||
env:
|
||||
REPORT_PATH: .pr-governance-report.json
|
||||
with:
|
||||
script: |
|
||||
const fs = require('fs');
|
||||
const report = JSON.parse(fs.readFileSync(process.env.REPORT_PATH, 'utf8'));
|
||||
const owner = context.repo.owner;
|
||||
const repo = context.repo.repo;
|
||||
const issue_number = context.payload.pull_request.number;
|
||||
const marker = report.comment_marker;
|
||||
const body = `${marker}\n${report.comment_markdown}`.trim();
|
||||
|
||||
const comments = await github.paginate(github.rest.issues.listComments, {
|
||||
owner,
|
||||
repo,
|
||||
issue_number,
|
||||
per_page: 100,
|
||||
});
|
||||
|
||||
const existing = comments.find(
|
||||
(comment) =>
|
||||
comment.user?.type === 'Bot' && typeof comment.body === 'string' && comment.body.includes(marker),
|
||||
);
|
||||
|
||||
if (existing) {
|
||||
await github.rest.issues.updateComment({
|
||||
owner,
|
||||
repo,
|
||||
comment_id: existing.id,
|
||||
body,
|
||||
});
|
||||
} else {
|
||||
await github.rest.issues.createComment({
|
||||
owner,
|
||||
repo,
|
||||
issue_number,
|
||||
body,
|
||||
});
|
||||
}
|
||||
|
||||
if (report.labels_to_add.length > 0) {
|
||||
await github.rest.issues.addLabels({
|
||||
owner,
|
||||
repo,
|
||||
issue_number,
|
||||
labels: report.labels_to_add,
|
||||
});
|
||||
}
|
||||
|
||||
for (const label of report.labels_to_remove) {
|
||||
try {
|
||||
await github.rest.issues.removeLabel({
|
||||
owner,
|
||||
repo,
|
||||
issue_number,
|
||||
name: label,
|
||||
});
|
||||
} catch (error) {
|
||||
if (error.status !== 404) {
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
- name: Fail when the PR body is incomplete
|
||||
if: steps.validate.outputs.valid != 'true'
|
||||
run: |
|
||||
echo "PR template validation failed. Update the PR body or move the PR back to draft."
|
||||
exit 1
|
||||
|
||||
label:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 10
|
||||
env:
|
||||
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
REPO: ${{ github.repository }}
|
||||
steps:
|
||||
- name: Ensure maintenance labels exist
|
||||
run: |
|
||||
set -euo pipefail
|
||||
|
||||
gh label create "status: needs rebase" \
|
||||
--repo "$REPO" \
|
||||
--color "fbca04" \
|
||||
--description "Pull request branch is behind the base branch" \
|
||||
--force
|
||||
gh label create "status: has conflicts" \
|
||||
--repo "$REPO" \
|
||||
--color "d73a4a" \
|
||||
--description "Pull request has merge conflicts with the base branch" \
|
||||
--force
|
||||
gh label create "status: ci failing" \
|
||||
--repo "$REPO" \
|
||||
--color "d73a4a" \
|
||||
--description "Required or reported CI checks are failing" \
|
||||
--force
|
||||
gh label create "status: needs author action" \
|
||||
--repo "$REPO" \
|
||||
--color "d93f0b" \
|
||||
--description "Pull request body or readiness checklist still needs author updates" \
|
||||
--force
|
||||
gh label create "status: ready for review" \
|
||||
--repo "$REPO" \
|
||||
--color "0e8a16" \
|
||||
--description "Pull request body is complete and the author marked it ready for human review" \
|
||||
--force
|
||||
|
||||
- name: Label open pull requests
|
||||
run: |
|
||||
set -euo pipefail
|
||||
|
||||
if jq -e '.pull_request.number' "$GITHUB_EVENT_PATH" >/dev/null; then
|
||||
pr_numbers="$(jq -r '.pull_request.number' "$GITHUB_EVENT_PATH")"
|
||||
else
|
||||
pr_numbers="$(gh pr list --repo "$REPO" --state open --limit 100 --json number --jq '.[].number')"
|
||||
fi
|
||||
|
||||
for pr in $pr_numbers; do
|
||||
data="$(gh pr view "$pr" --repo "$REPO" \
|
||||
--json isDraft,labels,mergeStateStatus,statusCheckRollup)"
|
||||
|
||||
merge_state="$(jq -r '.mergeStateStatus // "UNKNOWN"' <<<"$data")"
|
||||
check_state="$(jq -r '
|
||||
[
|
||||
(.statusCheckRollup // [])[]
|
||||
| select((.conclusion // .state // "") as $s
|
||||
| ["FAILURE", "TIMED_OUT", "ACTION_REQUIRED", "CANCELLED", "ERROR"] | index($s))
|
||||
]
|
||||
| if length > 0 then "failing" else "passing" end
|
||||
' <<<"$data")"
|
||||
is_draft="$(jq -r '.isDraft' <<<"$data")"
|
||||
|
||||
if [[ "$merge_state" == "BEHIND" ]]; then
|
||||
gh pr edit "$pr" --repo "$REPO" --add-label "status: needs rebase"
|
||||
else
|
||||
gh pr edit "$pr" --repo "$REPO" --remove-label "status: needs rebase" || true
|
||||
fi
|
||||
|
||||
if [[ "$merge_state" == "DIRTY" ]]; then
|
||||
gh pr edit "$pr" --repo "$REPO" --add-label "status: has conflicts"
|
||||
else
|
||||
gh pr edit "$pr" --repo "$REPO" --remove-label "status: has conflicts" || true
|
||||
fi
|
||||
|
||||
if [[ "$check_state" == "failing" ]]; then
|
||||
gh pr edit "$pr" --repo "$REPO" --add-label "status: ci failing"
|
||||
else
|
||||
gh pr edit "$pr" --repo "$REPO" --remove-label "status: ci failing" || true
|
||||
fi
|
||||
|
||||
if [[ "$merge_state" == "BEHIND" || "$merge_state" == "DIRTY" || "$check_state" == "failing" || "$is_draft" == "true" ]]; then
|
||||
gh pr edit "$pr" --repo "$REPO" --remove-label "status: ready for review" || true
|
||||
fi
|
||||
done
|
||||
|
|
|
|||
1
.gitignore
vendored
1
.gitignore
vendored
|
|
@ -17,6 +17,7 @@ scripts/*
|
|||
!scripts/sync-plugin-versions.py
|
||||
!scripts/changelog-gen.py
|
||||
!scripts/verify-versions.py
|
||||
!scripts/pr-governance.py
|
||||
!scripts/tests/
|
||||
!scripts/README.md
|
||||
!scripts/repro_codex_replay.py
|
||||
|
|
|
|||
|
|
@ -7,6 +7,11 @@ repos:
|
|||
language: system
|
||||
pass_filenames: false
|
||||
always_run: true
|
||||
- id: commitlint
|
||||
name: Commitlint
|
||||
entry: bash -lc 'npx --yes --package=@commitlint/cli --package=@commitlint/config-conventional -- commitlint --edit "$1" --config .commitlintrc.json' --
|
||||
language: system
|
||||
stages: [commit-msg]
|
||||
- repo: https://github.com/astral-sh/ruff-pre-commit
|
||||
rev: v0.9.4
|
||||
hooks:
|
||||
|
|
|
|||
|
|
@ -19,6 +19,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
|||
### Bug Fixes
|
||||
|
||||
* **ccr:** make retrieval store TTL configurable with `HEADROOM_CCR_TTL_SECONDS`, expose the effective TTL in `/v1/retrieve/stats`, and distinguish expired retrievals from missing hashes.
|
||||
* **proxy:** add native Bedrock `/model/{id}/converse-stream` route and forward it through the existing streaming EventStream/SSE pipeline.
|
||||
|
||||
|
||||
## [0.25.0](https://github.com/chopratejas/headroom/compare/v0.24.0...v0.25.0) (2026-06-12)
|
||||
|
|
|
|||
|
|
@ -73,15 +73,17 @@ A human maintainer reviews every dep change. PRs that add or bump a package must
|
|||
## PR workflow
|
||||
|
||||
1. Fork, branch from `main`.
|
||||
2. `pip install -e ".[dev]"` then `make install-git-hooks` — installs repo pre-commit checks on every commit and ci-precheck on every push.
|
||||
2. Install **Node 18+** and run `uv sync --extra dev` then `make install-git-hooks` — installs repo pre-commit checks on every commit, commitlint on every commit message, and ci-precheck on every push.
|
||||
3. One logical change per PR.
|
||||
4. Add tests.
|
||||
5. `pytest` · `ruff check .` · `ruff format .`
|
||||
5. `uv run pytest` · `uv run ruff check .` · `uv run ruff format .`
|
||||
6. Update `CHANGELOG.md` for user-facing changes.
|
||||
7. Open the PR with a clear description + `Real behavior proof` + any spec/justification required.
|
||||
7. Open the PR with a clear description + `Real behavior proof` + any spec/justification required, and keep the PR in draft until the `Review Readiness` boxes are complete.
|
||||
|
||||
**Title format** (conventional commits): `feat:`, `fix:`, `docs:`, `test:`, `refactor:`.
|
||||
|
||||
**Commit message format** is enforced locally by the repo's `commit-msg` hook and again in CI.
|
||||
|
||||
**Review:** CI green, one maintainer review, coverage held/improved.
|
||||
|
||||
## Development setup
|
||||
|
|
@ -90,10 +92,16 @@ A human maintainer reviews every dep change. PRs that add or bump a package must
|
|||
git clone https://github.com/chopratejas/headroom.git
|
||||
cd headroom
|
||||
python -m venv .venv && source .venv/bin/activate
|
||||
pip install -e ".[dev,relevance,proxy]"
|
||||
pytest
|
||||
node --version # Node 18+ required for commitlint hooks
|
||||
python -m pip install --upgrade pip
|
||||
python -m pip install -e ".[dev,relevance,proxy]"
|
||||
python -m pytest
|
||||
```
|
||||
|
||||
Headroom uses a `pyproject.toml`/`maturin` build backend. Older `pip`
|
||||
versions may fail editable installs by looking for `setup.py`; upgrade `pip`
|
||||
first or use `uv sync --extra dev`.
|
||||
|
||||
### Dev Containers
|
||||
|
||||
Two configs ship for VS Code / Codespaces:
|
||||
|
|
@ -103,6 +111,12 @@ Two configs ship for VS Code / Codespaces:
|
|||
|
||||
Inside, use: `uv run ruff check .`, `uv run pytest`, etc.
|
||||
|
||||
## Optional automated review
|
||||
|
||||
This repository includes `.github/copilot-instructions.md` so maintainers can opt into GitHub Copilot code review without adding workflow billing noise to every PR.
|
||||
|
||||
Enable or disable automatic Copilot review in **Settings → Rules → Rulesets → Automatically request Copilot code review**. Keep it off unless maintainers explicitly want the extra review traffic.
|
||||
|
||||
## Coding standards
|
||||
|
||||
- [Ruff](https://github.com/astral-sh/ruff) for lint + format, line length 100, PEP 8.
|
||||
|
|
|
|||
11
Makefile
11
Makefile
|
|
@ -27,7 +27,7 @@ help:
|
|||
@echo " make ci-precheck-rust - cargo fmt --check + clippy + test"
|
||||
@echo " make ci-precheck-python - smart_crusher-affected python tests"
|
||||
@echo " make ci-precheck-commitlint - lint commits since origin/main"
|
||||
@echo " make install-git-hooks - install a pre-push hook that runs ci-precheck"
|
||||
@echo " make install-git-hooks - install pre-commit, commit-msg, and pre-push hooks"
|
||||
|
||||
test:
|
||||
$(CARGO) test --workspace
|
||||
|
|
@ -123,16 +123,15 @@ ci-precheck-python:
|
|||
tests/test_toin_integration.py
|
||||
|
||||
# Lint commits since `origin/main`. Requires npx (Node 18+) on PATH.
|
||||
# Skips silently if npx is unavailable; install nodejs to enable.
|
||||
ci-precheck-commitlint:
|
||||
@echo "── ci-precheck-commitlint ─────────────────────────────────────"
|
||||
@if ! command -v npx >/dev/null 2>&1; then \
|
||||
echo "skip: npx not on PATH (install node 18+ to enable commitlint pre-check)"; \
|
||||
exit 0; \
|
||||
echo "error: npx not on PATH (install Node 18+ to enable commitlint checks)"; \
|
||||
exit 1; \
|
||||
fi
|
||||
@if ! git rev-parse --verify origin/main >/dev/null 2>&1; then \
|
||||
echo "skip: origin/main not fetched (run 'git fetch origin main')"; \
|
||||
exit 0; \
|
||||
echo "error: origin/main not fetched (run 'git fetch origin main')"; \
|
||||
exit 1; \
|
||||
fi
|
||||
npx --yes --package=@commitlint/cli --package=@commitlint/config-conventional -- \
|
||||
commitlint --from origin/main --to HEAD --config .commitlintrc.json
|
||||
|
|
|
|||
22
README.md
22
README.md
|
|
@ -147,10 +147,28 @@ Any OpenAI-compatible client works via `headroom proxy`. MCP-native: `headroom m
|
|||
Headroom can route GitHub Copilot CLI subscription traffic through the local proxy:
|
||||
|
||||
```bash
|
||||
headroom copilot-auth login
|
||||
headroom wrap copilot --subscription -- --model gpt-4o
|
||||
```
|
||||
|
||||
This lets Headroom intercept OpenAI-compatible Copilot CLI requests and apply the same proxy compression pipeline before forwarding to GitHub Copilot's hosted API. The wrapper resolves the account-specific Copilot API endpoint and prints it as `COPILOT_PROVIDER_API_URL=...` during launch.
|
||||
This lets Headroom intercept OpenAI-compatible Copilot CLI requests and apply the same proxy compression pipeline before forwarding to GitHub Copilot's hosted API. The wrapper exchanges Headroom's reusable GitHub OAuth token for Copilot's short-lived API token and prints the upstream endpoint as `COPILOT_PROVIDER_API_URL=...` during launch.
|
||||
|
||||
`headroom copilot-auth login` stores a Headroom-specific Copilot OAuth token.
|
||||
This avoids relying on generic GitHub or Copilot CLI tokens that can read
|
||||
Copilot account metadata but may still be rejected by Copilot's token-exchange
|
||||
endpoint.
|
||||
|
||||
For GitHub Enterprise Server or custom-domain Copilot deployments, set the
|
||||
deployment domain before launching:
|
||||
|
||||
```bash
|
||||
export GITHUB_COPILOT_ENTERPRISE_DOMAIN=ghe.example.com
|
||||
```
|
||||
|
||||
For GitHub.com Enterprise Cloud URLs such as
|
||||
`github.com/enterprises/your-enterprise`, do not set an enterprise-domain
|
||||
override. Headroom uses GitHub's normal token-exchange endpoint and the Copilot
|
||||
API endpoint advertised for the signed-in account.
|
||||
|
||||
Platform support note: macOS auth reuse via Copilot CLI Keychain storage has been smoke-tested. Windows Credential Manager, Linux Secret Service / `secret-tool`, and Docker/CI token-injection paths are implemented or planned as auth-discovery paths, but still need real OS validation before they should be considered fully vetted. For Docker and CI, prefer passing an explicit `GITHUB_COPILOT_TOKEN` or `GITHUB_COPILOT_GITHUB_TOKEN` rather than relying on host keychain access.
|
||||
|
||||
|
|
@ -302,7 +320,7 @@ Headroom runs **locally**, covers **every** content type, works with every major
|
|||
|
||||
```bash
|
||||
git clone https://github.com/chopratejas/headroom.git && cd headroom
|
||||
pip install -e ".[dev]" && pytest
|
||||
uv sync --extra dev && uv run pytest
|
||||
```
|
||||
|
||||
Devcontainers in `.devcontainer/` (default + `memory-stack` with Qdrant & Neo4j). See [CONTRIBUTING.md](CONTRIBUTING.md).
|
||||
|
|
|
|||
53
codecov.yml
53
codecov.yml
|
|
@ -1,19 +1,34 @@
|
|||
codecov:
|
||||
require_ci_to_pass: true
|
||||
|
||||
coverage:
|
||||
status:
|
||||
project:
|
||||
default:
|
||||
target: auto
|
||||
patch:
|
||||
default:
|
||||
target: auto
|
||||
|
||||
ignore:
|
||||
- "tests/**"
|
||||
- "scripts/tests/**"
|
||||
- ".github/**"
|
||||
- ".claude-plugin/**"
|
||||
- "plugins/headroom-agent-hooks/.claude-plugin/**"
|
||||
- "plugins/headroom-agent-hooks/.github/**"
|
||||
codecov:
|
||||
require_ci_to_pass: true
|
||||
|
||||
coverage:
|
||||
status:
|
||||
# Gate on the comprehensive unit suite (`python` flag from ci.yml's 4 test
|
||||
# shards), NOT the narrow native-e2e smoke flags (install-native /
|
||||
# wrap-native). Those e2e jobs upload first and barely exercise new code,
|
||||
# so an unscoped status computes patch at ~6% off the e2e flags alone and
|
||||
# flaps to FAILURE before the unit shards report. Scoping to `python` makes
|
||||
# the status reflect real coverage of the diff.
|
||||
project:
|
||||
default:
|
||||
target: auto
|
||||
flags:
|
||||
- python
|
||||
patch:
|
||||
default:
|
||||
target: auto
|
||||
flags:
|
||||
- python
|
||||
|
||||
flags:
|
||||
python:
|
||||
carryforward: false
|
||||
|
||||
ignore:
|
||||
- "tests/**"
|
||||
- "scripts/tests/**"
|
||||
- ".github/**"
|
||||
- ".claude-plugin/**"
|
||||
- "headroom/dashboard/templates/**"
|
||||
- "plugins/headroom-agent-hooks/.claude-plugin/**"
|
||||
- "plugins/headroom-agent-hooks/.github/**"
|
||||
|
|
|
|||
|
|
@ -239,21 +239,24 @@ impl CompressionPolicy {
|
|||
/// removes `delta_t` tokens from a message whose cached suffix is
|
||||
/// `suffix_tokens` long (#856).
|
||||
///
|
||||
/// Mutating message K invalidates every cached token after it: the
|
||||
/// suffix is re-written once at the write multiplier instead of
|
||||
/// being read at the read multiplier, costing
|
||||
/// `P_alive · (w − r) · S`. In exchange, `delta_t` tokens are gone
|
||||
/// from the current write and every one of the `expected_reads`
|
||||
/// remaining reads of the chain, saving `ΔT · (w + r·(R − 1))`.
|
||||
/// Mutating message K invalidates every cached token after it. When
|
||||
/// the cache is warm the mutated ΔT tokens are themselves already
|
||||
/// cache-written, so keeping them costs only reads (`ΔT · r · R`)
|
||||
/// while mutating re-writes the suffix: alive-case saving is
|
||||
/// `ΔT·r·R − (w−r)·S`. When the cache is dead there is no suffix
|
||||
/// penalty and the full `ΔT·(w + r·(R−1))` is saved. Taking the
|
||||
/// expectation over `P_alive`:
|
||||
///
|
||||
/// gain = ΔT · (w + r·(R − 1)) − P_alive · (w − r) · S
|
||||
/// gain = ΔT · (w + r·(R − 1)) − P_alive · (w − r) · (S + ΔT)
|
||||
///
|
||||
/// Sanity anchors (Anthropic w=1.25, r=0.1), matching the unit
|
||||
/// tests below: a 2K shave under a 50K warm suffix needs ~276
|
||||
/// tests below: a 2K shave under a 50K warm suffix needs 287.5
|
||||
/// remaining reads to pay off (rarely profitable); a 50K shave
|
||||
/// under a 10K suffix is profitable from the first write (its
|
||||
/// break-even read count is negative); a live-zone edit (S = 0)
|
||||
/// is always profitable.
|
||||
/// under a 10K suffix breaks even at 2.3 reads (profitable in any
|
||||
/// session with a few turns left); an edit with S = 0 is profitable
|
||||
/// whenever at least one read remains. Callers gating not-yet-cached
|
||||
/// content (live-zone edits) should bypass this formula — it prices
|
||||
/// mutations of content the cache has already written.
|
||||
///
|
||||
/// Takes `&self` so a follow-up can apply per-mode margins; today
|
||||
/// the arithmetic is mode-independent. Inputs are clamped:
|
||||
|
|
@ -276,7 +279,15 @@ impl CompressionPolicy {
|
|||
} else {
|
||||
p_alive.clamp(0.0, 1.0)
|
||||
};
|
||||
(delta_t as f32) * (w + r * (reads - 1.0)) - alive * (w - r) * (suffix_tokens as f32)
|
||||
// Corrected warm-case penalty (#856 follow-up): when the cache is
|
||||
// alive, the ΔT tokens are already cache-written, so keeping them
|
||||
// costs only reads — a mutation can avoid at most ΔT·r·R, not a
|
||||
// fresh write. Blending alive (ΔT·r·R − (w−r)·S) and dead
|
||||
// (ΔT·(w + r·(R−1))) cases over P_alive gives a penalty over
|
||||
// S + ΔT, not S alone. The looser ·S form overstated gain by
|
||||
// P_alive·(w−r)·ΔT — always pro-mutation, largest for big shaves.
|
||||
(delta_t as f32) * (w + r * (reads - 1.0))
|
||||
- alive * (w - r) * ((suffix_tokens as f32) + (delta_t as f32))
|
||||
}
|
||||
|
||||
/// Decision form of [`Self::net_mutation_gain`]: mutate iff the
|
||||
|
|
@ -292,9 +303,13 @@ impl CompressionPolicy {
|
|||
}
|
||||
|
||||
/// Remaining-read count at which a warm-cache (P_alive = 1)
|
||||
/// mutation breaks even:
|
||||
/// mutation breaks even. With the corrected penalty this is exactly
|
||||
///
|
||||
/// R = ((w − r) / r) · (S/ΔT − 1) ≈ 11.5 · S/ΔT for S ≫ ΔT
|
||||
/// R = ((w − r) / r) · S/ΔT = 11.5 · S/ΔT (Anthropic 5-min)
|
||||
///
|
||||
/// reproducing the #856 anchors precisely: 2K shave / 50K suffix →
|
||||
/// 287.5 (~290 reads, rarely profitable); 50K shave / 10K suffix →
|
||||
/// 2.3 (profitable in any session with a few turns left).
|
||||
///
|
||||
/// Useful for decision telemetry ("this edit pays off if the
|
||||
/// session lasts N more turns"). Returns 0 when `delta_t` is 0
|
||||
|
|
@ -305,7 +320,7 @@ impl CompressionPolicy {
|
|||
}
|
||||
let w = CACHE_WRITE_MULTIPLIER;
|
||||
let r = CACHE_READ_MULTIPLIER;
|
||||
((w - r) / r) * ((suffix_tokens as f32) / (delta_t as f32) - 1.0)
|
||||
((w - r) / r) * ((suffix_tokens as f32) / (delta_t as f32))
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -414,31 +429,36 @@ mod tests {
|
|||
#[test]
|
||||
fn net_gain_small_shave_deep_suffix_is_loss() {
|
||||
// Shave 2K under a 50K warm suffix at R=10 remaining reads:
|
||||
// 2000·(1.25 + 0.1·9) − 1.0·1.15·50000 = 4300 − 57500 = −53200.
|
||||
// 2000·(1.25 + 0.1·9) − 1.0·1.15·52000 = 4300 − 59800 = −55500.
|
||||
let p = CompressionPolicy::for_mode(AuthMode::Payg);
|
||||
let gain = p.net_mutation_gain(2_000, 50_000, 10.0, 1.0);
|
||||
assert!((gain - (-53_200.0)).abs() < 1.0, "gain = {gain}");
|
||||
assert!((gain - (-55_500.0)).abs() < 1.0, "gain = {gain}");
|
||||
assert!(!p.should_mutate_deep(2_000, 50_000, 10.0, 1.0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn net_gain_big_shave_shallow_suffix_is_win() {
|
||||
// Shave 50K under a 10K warm suffix at R=3:
|
||||
// 50000·(1.25 + 0.1·2) − 1.0·1.15·10000 = 72500 − 11500 = 61000.
|
||||
// 50000·(1.25 + 0.1·2) − 1.0·1.15·60000 = 72500 − 69000 = 3500.
|
||||
// Tight but positive — consistent with the 2.3-read break-even.
|
||||
let p = CompressionPolicy::for_mode(AuthMode::Payg);
|
||||
let gain = p.net_mutation_gain(50_000, 10_000, 3.0, 1.0);
|
||||
assert!((gain - 61_000.0).abs() < 1.0, "gain = {gain}");
|
||||
assert!((gain - 3_500.0).abs() < 1.0, "gain = {gain}");
|
||||
assert!(p.should_mutate_deep(50_000, 10_000, 3.0, 1.0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn net_gain_live_zone_edit_always_profitable() {
|
||||
// S = 0 derives the existing Subscription live-zone policy as a
|
||||
// special case: nothing cached is invalidated, so any positive
|
||||
// shave wins even at R=0 (gain = ΔT·(w − r) > 0).
|
||||
fn net_gain_no_suffix_edit_profitable_with_reads_remaining() {
|
||||
// S = 0: nothing cached after the edit is invalidated. Warm-case
|
||||
// saving is the avoided rereads, ΔT·r·R — positive whenever at
|
||||
// least one read remains. At R=0 with a warm cache the gain is
|
||||
// exactly 0 (already written, never read again): the boundary
|
||||
// where mutating is pointless rather than harmful.
|
||||
let p = CompressionPolicy::for_mode(AuthMode::Subscription);
|
||||
assert!(p.should_mutate_deep(1, 0, 0.0, 1.0));
|
||||
assert!(p.should_mutate_deep(2_000, 0, 0.0, 1.0));
|
||||
assert!(p.should_mutate_deep(1, 0, 1.0, 1.0));
|
||||
assert!(p.should_mutate_deep(2_000, 0, 1.0, 1.0));
|
||||
let boundary = p.net_mutation_gain(2_000, 0, 0.0, 1.0);
|
||||
assert!(boundary.abs() < f32::EPSILON, "boundary = {boundary}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
|
@ -472,13 +492,14 @@ mod tests {
|
|||
|
||||
#[test]
|
||||
fn break_even_reads_matches_research_anchor() {
|
||||
// R = 11.5·(S/ΔT − 1): 2K shave / 50K suffix → 11.5·24 = 276
|
||||
// (rarely profitable); 50K shave / 10K suffix →
|
||||
// 11.5·(0.2 − 1) < 0 → profitable from the first read.
|
||||
// R = 11.5·S/ΔT, the #856 anchors exactly: 2K shave / 50K
|
||||
// suffix → 11.5·25 = 287.5 (rarely profitable); 50K shave /
|
||||
// 10K suffix → 11.5·0.2 = 2.3 (profitable within a few turns).
|
||||
let p = CompressionPolicy::for_mode(AuthMode::Payg);
|
||||
let r = p.break_even_reads(2_000, 50_000);
|
||||
assert!((r - 276.0).abs() < 0.5, "break-even = {r}");
|
||||
assert!(p.break_even_reads(50_000, 10_000) < 0.0);
|
||||
assert!((r - 287.5).abs() < 0.5, "break-even = {r}");
|
||||
let shallow = p.break_even_reads(50_000, 10_000);
|
||||
assert!((shallow - 2.3).abs() < 0.05, "break-even = {shallow}");
|
||||
assert_eq!(p.break_even_reads(0, 10_000), 0.0);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -154,12 +154,18 @@ pub fn translate_message(
|
|||
event_type: event_type.to_string(),
|
||||
})
|
||||
}
|
||||
(OutputMode::Sse, "chunk") => {
|
||||
(OutputMode::Sse, "chunk")
|
||||
| (OutputMode::Sse, "messageStart")
|
||||
| (OutputMode::Sse, "contentBlockStart")
|
||||
| (OutputMode::Sse, "contentBlockDelta")
|
||||
| (OutputMode::Sse, "contentBlockStop")
|
||||
| (OutputMode::Sse, "messageStop")
|
||||
| (OutputMode::Sse, "metadata") => {
|
||||
tracing::info!(
|
||||
event = "bedrock_eventstream_translated_to_sse",
|
||||
event_type = event_type,
|
||||
payload_bytes = message.payload.len(),
|
||||
"translated bedrock eventstream chunk to sse frame"
|
||||
"translated bedrock eventstream message to sse frame"
|
||||
);
|
||||
Ok(TranslateOutcome::Emit(payload_to_sse_frame(
|
||||
&message.payload,
|
||||
|
|
@ -380,6 +386,28 @@ mod tests {
|
|||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn translate_converse_event_to_sse_frame() {
|
||||
let bytes = crate::bedrock::eventstream::MessageBuilder::new()
|
||||
.header_string(":event-type", "contentBlockDelta")
|
||||
.header_string(":message-type", "event")
|
||||
.payload(Bytes::from_static(
|
||||
br#"{"contentBlockIndex":0,"delta":{"text":"hi"}}"#,
|
||||
))
|
||||
.build();
|
||||
let msg = parse(&bytes).unwrap();
|
||||
let outcome = translate_message(&msg, OutputMode::Sse).unwrap();
|
||||
match outcome {
|
||||
TranslateOutcome::Emit(b) => {
|
||||
let s = std::str::from_utf8(&b).unwrap();
|
||||
assert!(s.starts_with("data: "));
|
||||
assert!(s.ends_with("\n\n"));
|
||||
assert!(s.contains("contentBlockIndex"));
|
||||
}
|
||||
other => panic!("expected Emit; got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn missing_event_type_is_loud() {
|
||||
// A message lacking :event-type must not silently translate.
|
||||
|
|
|
|||
|
|
@ -80,8 +80,9 @@ const ANTHROPIC_VENDOR_PREFIX: &str = "anthropic.";
|
|||
/// AWS Bedrock Runtime DNS template.
|
||||
const BEDROCK_RUNTIME_HOST_TEMPLATE: &str = "bedrock-runtime.{region}.amazonaws.com";
|
||||
|
||||
/// Path action for the streaming route.
|
||||
/// Path action for the streaming routes.
|
||||
const STREAMING_ACTION: &str = "invoke-with-response-stream";
|
||||
const CONVERSE_STREAM_ACTION: &str = "converse-stream";
|
||||
|
||||
/// RAII guard that observes the `bedrock_invoke_latency_seconds`
|
||||
/// histogram on drop. Mirrors the [`crate::bedrock::invoke`] guard
|
||||
|
|
@ -166,8 +167,26 @@ pub async fn handle_invoke_streaming(
|
|||
body.clone()
|
||||
};
|
||||
|
||||
// 2. Build upstream URL.
|
||||
let upstream_url = match build_bedrock_streaming_upstream(&state, &model_id, &uri) {
|
||||
// 2. Resolve the Bedrock streaming action from the inbound path and
|
||||
// build the upstream URL.
|
||||
let action = match extract_streaming_action(uri.path()) {
|
||||
Some(a) => a,
|
||||
None => {
|
||||
tracing::error!(
|
||||
event = "bedrock_streaming_action_invalid",
|
||||
request_id = %request_id,
|
||||
path = %uri.path(),
|
||||
"bedrock invoke-streaming: unrecognized streaming action path"
|
||||
);
|
||||
return error_response(
|
||||
StatusCode::BAD_REQUEST,
|
||||
"bedrock_streaming_action_invalid",
|
||||
"Unsupported Bedrock streaming action path",
|
||||
);
|
||||
}
|
||||
};
|
||||
|
||||
let upstream_url = match build_bedrock_streaming_upstream(&state, &model_id, &uri, action) {
|
||||
Ok(u) => u,
|
||||
Err(msg) => {
|
||||
tracing::error!(
|
||||
|
|
@ -888,6 +907,7 @@ fn build_bedrock_streaming_upstream(
|
|||
state: &AppState,
|
||||
model_id: &str,
|
||||
uri: &Uri,
|
||||
action: &str,
|
||||
) -> Result<Url, String> {
|
||||
let base = match state.config.bedrock_endpoint.as_ref() {
|
||||
Some(u) => u.clone(),
|
||||
|
|
@ -901,7 +921,7 @@ fn build_bedrock_streaming_upstream(
|
|||
let path = format!(
|
||||
"/model/{model_id}/{action}",
|
||||
model_id = model_id,
|
||||
action = STREAMING_ACTION,
|
||||
action = action,
|
||||
);
|
||||
let mut joined = base;
|
||||
joined.set_path(&path);
|
||||
|
|
@ -911,6 +931,16 @@ fn build_bedrock_streaming_upstream(
|
|||
Ok(joined)
|
||||
}
|
||||
|
||||
fn extract_streaming_action(path: &str) -> Option<&'static str> {
|
||||
if path.ends_with(&format!("/{STREAMING_ACTION}")) {
|
||||
Some(STREAMING_ACTION)
|
||||
} else if path.ends_with(&format!("/{CONVERSE_STREAM_ACTION}")) {
|
||||
Some(CONVERSE_STREAM_ACTION)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
fn collect_signed_headers(headers: &HeaderMap, upstream_url: &Url) -> Vec<(String, String)> {
|
||||
let mut out: Vec<(String, String)> = Vec::with_capacity(headers.len() + 1);
|
||||
for (name, value) in headers.iter() {
|
||||
|
|
@ -996,6 +1026,7 @@ mod tests {
|
|||
&state,
|
||||
"anthropic.claude-3-haiku-20240307-v1:0",
|
||||
&uri,
|
||||
STREAMING_ACTION,
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
|
|
@ -1012,4 +1043,54 @@ mod tests {
|
|||
assert!(s.ends_with("\n\n"));
|
||||
assert!(s.contains("bedrock_eventstream_crc_mismatch"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_streaming_upstream_supports_converse_stream_action() {
|
||||
use crate::config::Config;
|
||||
let mut config = Config::for_test(Url::parse("http://up:8080").unwrap());
|
||||
config.bedrock_region = "eu-west-1".to_string();
|
||||
let state = AppState {
|
||||
config: std::sync::Arc::new(config),
|
||||
client: reqwest::Client::new(),
|
||||
bedrock_credentials: None,
|
||||
drift_state: crate::cache_stabilization::drift_detector::DriftState::new(8),
|
||||
vertex_token_source: std::sync::Arc::new(crate::vertex::StaticTokenSource::new(
|
||||
"test".to_string(),
|
||||
)),
|
||||
};
|
||||
let uri: Uri = "/model/anthropic.claude-3-haiku-20240307-v1:0/converse-stream"
|
||||
.parse()
|
||||
.unwrap();
|
||||
let url = build_bedrock_streaming_upstream(
|
||||
&state,
|
||||
"anthropic.claude-3-haiku-20240307-v1:0",
|
||||
&uri,
|
||||
CONVERSE_STREAM_ACTION,
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
url.as_str(),
|
||||
"https://bedrock-runtime.eu-west-1.amazonaws.com/model/anthropic.claude-3-haiku-20240307-v1:0/converse-stream"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_streaming_action_supports_both_bedrock_paths() {
|
||||
assert_eq!(
|
||||
extract_streaming_action(
|
||||
"/model/anthropic.claude-3-haiku-20240307-v1:0/invoke-with-response-stream"
|
||||
),
|
||||
Some(STREAMING_ACTION)
|
||||
);
|
||||
assert_eq!(
|
||||
extract_streaming_action(
|
||||
"/model/anthropic.claude-3-haiku-20240307-v1:0/converse-stream"
|
||||
),
|
||||
Some(CONVERSE_STREAM_ACTION)
|
||||
);
|
||||
assert_eq!(
|
||||
extract_streaming_action("/model/anthropic.claude-3-haiku-20240307-v1:0/invoke"),
|
||||
None
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -209,15 +209,22 @@ pub fn build_app(state: AppState) -> Router {
|
|||
"/model/:model_id/converse",
|
||||
post(crate::bedrock::invoke::handle_invoke),
|
||||
)
|
||||
// PR-D2: streaming counterpart. Bedrock's protocol is
|
||||
// PR-D2/PR-D5: streaming counterparts. Bedrock's protocol is
|
||||
// binary EventStream; the handler parses incrementally,
|
||||
// optionally translates each chunk to an SSE frame, and
|
||||
// tees translated frames into AnthropicStreamState for
|
||||
// telemetry. See `bedrock::invoke_streaming`.
|
||||
// telemetry. `invoke-with-response-stream` and
|
||||
// `converse-stream` share the same wire framing and
|
||||
// processing pipeline, so both route to the same handler.
|
||||
// See `bedrock::invoke_streaming`.
|
||||
.route(
|
||||
"/model/:model_id/invoke-with-response-stream",
|
||||
post(crate::bedrock::invoke_streaming::handle_invoke_streaming),
|
||||
)
|
||||
.route(
|
||||
"/model/:model_id/converse-stream",
|
||||
post(crate::bedrock::invoke_streaming::handle_invoke_streaming),
|
||||
)
|
||||
.route_layer(axum::middleware::from_fn(
|
||||
crate::bedrock::classify_and_attach_auth_mode,
|
||||
))
|
||||
|
|
|
|||
|
|
@ -218,6 +218,18 @@ async fn mount_eventstream_upstream(upstream: &MockServer, body: Bytes) {
|
|||
.await;
|
||||
}
|
||||
|
||||
async fn mount_eventstream_upstream_for_action(upstream: &MockServer, body: Bytes, action: &str) {
|
||||
Mock::given(method("POST"))
|
||||
.and(path(format!("/model/{TEST_MODEL}/{action}")))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.insert_header("content-type", "application/vnd.amazon.eventstream")
|
||||
.set_body_bytes(body.to_vec()),
|
||||
)
|
||||
.mount(upstream)
|
||||
.await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn eventstream_translated_to_sse() {
|
||||
let _ = tracing_subscriber::fmt()
|
||||
|
|
@ -465,6 +477,53 @@ async fn client_can_choose_eventstream_or_sse() {
|
|||
proxy.shutdown().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn converse_stream_route_translates_to_sse() {
|
||||
let upstream = MockServer::start().await;
|
||||
let bedrock_bytes = synthesize_bedrock_stream();
|
||||
mount_eventstream_upstream_for_action(&upstream, bedrock_bytes, "converse-stream").await;
|
||||
let proxy = bedrock_proxy(&upstream, |c| {
|
||||
c.compression_mode = headroom_proxy::config::CompressionMode::Off;
|
||||
})
|
||||
.await;
|
||||
|
||||
let body = serde_json::to_vec(&json!({
|
||||
"anthropic_version": "bedrock-2023-05-31",
|
||||
"max_tokens": 16,
|
||||
"messages": [{"role":"user","content":"hi"}]
|
||||
}))
|
||||
.unwrap();
|
||||
|
||||
let resp = reqwest::Client::new()
|
||||
.post(format!(
|
||||
"{}/model/{TEST_MODEL}/converse-stream",
|
||||
proxy.url()
|
||||
))
|
||||
.header("content-type", "application/json")
|
||||
.header("accept", "text/event-stream")
|
||||
.body(body)
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(resp.status(), 200);
|
||||
let ct = resp
|
||||
.headers()
|
||||
.get("content-type")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.unwrap_or("");
|
||||
assert!(
|
||||
ct.starts_with("text/event-stream"),
|
||||
"converse-stream should emit SSE when client asks for SSE; got {ct}"
|
||||
);
|
||||
|
||||
let text = resp.text().await.unwrap();
|
||||
assert!(text.contains("event: content_block_delta"));
|
||||
assert!(text.contains("\"text\":\"OK\""));
|
||||
|
||||
proxy.shutdown().await;
|
||||
}
|
||||
|
||||
// ─── Test 5: Property test — never panic on adversarial bytes ─────
|
||||
|
||||
proptest! {
|
||||
|
|
|
|||
|
|
@ -421,4 +421,4 @@ print(f"Tokens: {result.tokens_before} -> {result.tokens_after}")
|
|||
1. Enable debug logging and check the output
|
||||
2. Use `simulate()` to see what transforms would apply
|
||||
3. Run `validate_setup()` for configuration issues
|
||||
4. File an issue at [github.com/headroom-sdk/headroom](https://github.com/headroom-sdk/headroom/issues) with your Headroom version, Python version, provider, debug log output, and minimal reproduction code
|
||||
4. File an issue at [github.com/chopratejas/headroom](https://github.com/chopratejas/headroom/issues) with your Headroom version, Python version, provider, debug log output, and minimal reproduction code
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ survives that kind of sys.modules mutation.
|
|||
|
||||
from . import ( # noqa: F401
|
||||
capture,
|
||||
copilot_auth,
|
||||
evals,
|
||||
init,
|
||||
install,
|
||||
|
|
|
|||
81
headroom/cli/copilot_auth.py
Normal file
81
headroom/cli/copilot_auth.py
Normal file
|
|
@ -0,0 +1,81 @@
|
|||
"""GitHub Copilot authentication commands."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import click
|
||||
|
||||
from headroom.cli.main import main
|
||||
from headroom.copilot_auth import (
|
||||
DEFAULT_GITHUB_HOST,
|
||||
headroom_copilot_auth_path,
|
||||
poll_copilot_device_authorization,
|
||||
read_headroom_copilot_oauth_token,
|
||||
save_headroom_copilot_oauth_token,
|
||||
start_copilot_device_authorization,
|
||||
token_fingerprint,
|
||||
)
|
||||
|
||||
|
||||
@main.group("copilot-auth")
|
||||
def copilot_auth() -> None:
|
||||
"""Manage Headroom's GitHub Copilot OAuth token."""
|
||||
|
||||
|
||||
@copilot_auth.command("login")
|
||||
@click.option(
|
||||
"--domain",
|
||||
default=DEFAULT_GITHUB_HOST,
|
||||
show_default=True,
|
||||
help=(
|
||||
"GitHub login domain. Use github.com for GitHub.com Enterprise Cloud; "
|
||||
"only pass a custom hostname for GitHub Enterprise Server."
|
||||
),
|
||||
)
|
||||
def login(domain: str) -> None:
|
||||
"""Sign in with GitHub's Copilot OAuth device-code flow."""
|
||||
|
||||
try:
|
||||
device = start_copilot_device_authorization(domain=domain)
|
||||
except Exception as exc:
|
||||
raise click.ClickException(f"Unable to start GitHub device login: {exc}") from exc
|
||||
|
||||
verification_uri = str(device.get("verification_uri") or "").strip()
|
||||
user_code = str(device.get("user_code") or "").strip()
|
||||
device_code = str(device.get("device_code") or "").strip()
|
||||
interval = int(device.get("interval") or 5)
|
||||
expires_in = int(device.get("expires_in") or 900)
|
||||
if not verification_uri or not user_code or not device_code:
|
||||
raise click.ClickException("GitHub device login returned an incomplete response.")
|
||||
|
||||
click.echo("GitHub Copilot OAuth login")
|
||||
click.echo(f" Open: {verification_uri}")
|
||||
click.echo(f" Code: {user_code}")
|
||||
click.echo(" Waiting for authorization...")
|
||||
|
||||
try:
|
||||
token = poll_copilot_device_authorization(
|
||||
device_code,
|
||||
domain=domain,
|
||||
interval=interval,
|
||||
expires_in=expires_in,
|
||||
)
|
||||
except Exception as exc:
|
||||
raise click.ClickException(f"GitHub device login failed: {exc}") from exc
|
||||
|
||||
path = save_headroom_copilot_oauth_token(token, domain=domain)
|
||||
click.echo(f" Saved: {path}")
|
||||
click.echo(f" Token fingerprint: {token_fingerprint(token)}")
|
||||
|
||||
|
||||
@copilot_auth.command("status")
|
||||
def status() -> None:
|
||||
"""Show whether Headroom has a saved Copilot OAuth token."""
|
||||
|
||||
token = read_headroom_copilot_oauth_token()
|
||||
path = headroom_copilot_auth_path()
|
||||
click.echo(f"Auth file: {path}")
|
||||
if not token:
|
||||
click.echo("Status: not logged in")
|
||||
return
|
||||
click.echo("Status: logged in")
|
||||
click.echo(f"Token fingerprint: {token_fingerprint(token)}")
|
||||
|
|
@ -38,6 +38,7 @@ from headroom.install.runtime import (
|
|||
)
|
||||
from headroom.install.state import load_manifest, save_manifest
|
||||
from headroom.install.supervisors import start_supervisor
|
||||
from headroom.providers.codex.install import codex_uses_chatgpt_auth
|
||||
|
||||
from .main import main
|
||||
|
||||
|
|
@ -291,6 +292,14 @@ def _ensure_codex_provider(path: Path, port: int) -> None:
|
|||
import re
|
||||
|
||||
logger.debug("ensure codex provider block: %s (port=%s)", path, port)
|
||||
# Emit requires_openai_auth only for ChatGPT-OAuth users (restores the
|
||||
# account menu); omitting it for API-key users avoids forcing an OAuth
|
||||
# login (#406).
|
||||
requires_openai_auth = (
|
||||
"requires_openai_auth = true\n"
|
||||
if codex_uses_chatgpt_auth(path.parent / "auth.json")
|
||||
else ""
|
||||
)
|
||||
block = (
|
||||
f"{_CODEX_PROVIDER_MARKER_START}\n"
|
||||
'model_provider = "headroom"\n'
|
||||
|
|
@ -299,6 +308,7 @@ def _ensure_codex_provider(path: Path, port: int) -> None:
|
|||
'name = "Headroom init proxy"\n'
|
||||
f'base_url = "http://127.0.0.1:{port}/v1"\n'
|
||||
"supports_websockets = true\n"
|
||||
f"{requires_openai_auth}"
|
||||
f"{_CODEX_PROVIDER_MARKER_END}"
|
||||
)
|
||||
content = path.read_text(encoding="utf-8") if path.exists() else ""
|
||||
|
|
|
|||
|
|
@ -38,6 +38,7 @@ def _register_commands() -> None:
|
|||
from . import (
|
||||
agent_savings, # noqa: F401
|
||||
capture, # noqa: F401
|
||||
copilot_auth, # noqa: F401
|
||||
evals, # noqa: F401
|
||||
init, # noqa: F401
|
||||
install, # noqa: F401
|
||||
|
|
|
|||
|
|
@ -783,12 +783,6 @@ def proxy(
|
|||
optimize=not no_optimize,
|
||||
cache_enabled=not no_cache,
|
||||
rate_limit_enabled=not no_rate_limit,
|
||||
# CCR opt-outs for compression-only deployments (streaming / non-MCP
|
||||
# clients that can't resolve the injected retrieve tool). Defaults keep
|
||||
# CCR fully on; each flag flips one dataclass default to False.
|
||||
ccr_inject_tool=not no_ccr_inject_tool,
|
||||
ccr_inject_marker=not no_ccr_marker,
|
||||
ccr_proactive_expansion=not no_ccr_proactive_expansion,
|
||||
compress_user_messages=_get_env_bool("HEADROOM_COMPRESS_USER_MESSAGES", False),
|
||||
min_tokens_to_crush=_get_env_int_optional("HEADROOM_MIN_TOKENS") or 500,
|
||||
max_items_after_crush=_get_env_int_optional("HEADROOM_MAX_ITEMS") or 50,
|
||||
|
|
@ -799,6 +793,12 @@ def proxy(
|
|||
protect_recent=_get_env_int_optional("HEADROOM_PROTECT_RECENT"),
|
||||
protect_analysis_context=_get_env_bool_optional("HEADROOM_PROTECT_ANALYSIS_CONTEXT"),
|
||||
accuracy_guard=os.environ.get("HEADROOM_ACCURACY_GUARD") or None,
|
||||
# CCR opt-outs for compression-only deployments (streaming / non-MCP
|
||||
# clients that can't resolve the injected retrieve tool). Defaults keep
|
||||
# CCR fully on; each flag flips one dataclass default to False.
|
||||
ccr_inject_tool=not no_ccr_inject_tool,
|
||||
ccr_inject_marker=not no_ccr_marker,
|
||||
ccr_proactive_expansion=not no_ccr_proactive_expansion,
|
||||
# Flatten repeat-flag tuple AND any comma-separated values inside it.
|
||||
# `--proxy-extension a,b --proxy-extension c` and `HEADROOM_PROXY_EXTENSIONS=a,b,c`
|
||||
# both yield ["a", "b", "c"]. None when nothing was supplied.
|
||||
|
|
|
|||
|
|
@ -38,15 +38,19 @@ if sys.platform == "win32" and hasattr(sys.stdout, "buffer"):
|
|||
import click
|
||||
|
||||
from headroom._version import __version__ as _HEADROOM_VERSION
|
||||
from headroom.agent_savings import (
|
||||
apply_agent_savings_env_defaults,
|
||||
)
|
||||
from headroom.copilot_auth import (
|
||||
has_oauth_auth,
|
||||
resolve_client_bearer_token,
|
||||
resolve_copilot_api_url,
|
||||
resolve_subscription_bearer_token,
|
||||
resolve_subscription_bearer_token_details,
|
||||
)
|
||||
from headroom.providers.aider import build_launch_env as _build_aider_launch_env
|
||||
from headroom.providers.claude import proxy_base_url as _claude_proxy_base_url
|
||||
from headroom.providers.codex import build_launch_env as _build_codex_launch_env
|
||||
from headroom.providers.codex.install import codex_uses_chatgpt_auth
|
||||
from headroom.providers.copilot import (
|
||||
build_launch_env as _build_copilot_launch_env,
|
||||
)
|
||||
|
|
@ -95,6 +99,7 @@ _CONTEXT_TOOL_ENV = "HEADROOM_CONTEXT_TOOL"
|
|||
_CONTEXT_TOOL_RTK = "rtk"
|
||||
_CONTEXT_TOOL_LEAN_CTX = "lean-ctx"
|
||||
_VALID_CONTEXT_TOOLS = {_CONTEXT_TOOL_RTK, _CONTEXT_TOOL_LEAN_CTX}
|
||||
_AGENT_SAVINGS_TARGET_AGENTS = {"claude", "codex", "cursor"}
|
||||
_WRAP_PROXY_TIMEOUT_ENV = "HEADROOM_WRAP_PROXY_TIMEOUT"
|
||||
_WRAP_PROXY_TIMEOUT_DEFAULT_SECONDS = 45
|
||||
_WRAP_PROXY_TIMEOUT_ML_DEFAULT_SECONDS = 90
|
||||
|
|
@ -360,6 +365,8 @@ def _start_proxy(
|
|||
# Ensure proxy subprocess uses UTF-8 (Windows defaults to cp1252)
|
||||
proxy_env = os.environ.copy()
|
||||
proxy_env["PYTHONIOENCODING"] = "utf-8"
|
||||
if agent_type in {"claude", "codex", "cursor"}:
|
||||
apply_agent_savings_env_defaults(proxy_env)
|
||||
|
||||
# Tell the proxy which agent is being wrapped (for traffic learning output)
|
||||
if agent_type != "unknown":
|
||||
|
|
@ -367,8 +374,6 @@ def _start_proxy(
|
|||
proxy_env.setdefault("HEADROOM_STACK", f"wrap_{agent_type}")
|
||||
savings_profile = _wrap_agent_savings_profile(agent_type)
|
||||
if savings_profile is not None:
|
||||
from headroom.agent_savings import apply_agent_savings_env_defaults
|
||||
|
||||
apply_agent_savings_env_defaults(proxy_env, savings_profile)
|
||||
if openai_api_url:
|
||||
proxy_env["OPENAI_TARGET_API_URL"] = openai_api_url
|
||||
|
|
@ -380,6 +385,8 @@ def _start_proxy(
|
|||
# GITHUB_COPILOT_API_TOKEN directly, making upstream auth deterministic.
|
||||
if copilot_api_token:
|
||||
proxy_env["GITHUB_COPILOT_API_TOKEN"] = copilot_api_token
|
||||
if openai_api_url:
|
||||
proxy_env["GITHUB_COPILOT_API_URL"] = openai_api_url
|
||||
|
||||
proc = subprocess.Popen(
|
||||
cmd,
|
||||
|
|
@ -1036,12 +1043,19 @@ def _inject_codex_provider_config(port: int) -> None:
|
|||
f'openai_base_url = "http://127.0.0.1:{port}/v1"\n'
|
||||
f"{_CODEX_END_MARKER}\n"
|
||||
)
|
||||
# Emit requires_openai_auth only for ChatGPT-OAuth users (restores the
|
||||
# account menu); omitting it for API-key users avoids forcing an OAuth
|
||||
# login (#406).
|
||||
requires_openai_auth = (
|
||||
"requires_openai_auth = true\n" if codex_uses_chatgpt_auth(config_dir / "auth.json") else ""
|
||||
)
|
||||
provider_section = (
|
||||
f"{_CODEX_TOP_LEVEL_MARKER}\n"
|
||||
"[model_providers.headroom]\n"
|
||||
'name = "OpenAI via Headroom proxy"\n'
|
||||
f'base_url = "http://127.0.0.1:{port}/v1"\n'
|
||||
f"supports_websockets = true\n"
|
||||
f"{requires_openai_auth}"
|
||||
# Per-project savings: Codex sends the header only when the mapped
|
||||
# env var (HEADROOM_PROJECT, set by `headroom wrap codex`) exists at
|
||||
# Codex runtime. Inline table keeps the key inside this section so
|
||||
|
|
@ -1553,6 +1567,77 @@ def _proxy_health_config(payload: dict[str, Any] | None) -> dict[str, Any] | Non
|
|||
return config if isinstance(config, dict) else None
|
||||
|
||||
|
||||
def _env_bool_value(value: str) -> bool:
|
||||
return value.strip().lower() in {"1", "true", "yes", "on"}
|
||||
|
||||
|
||||
def _agent_savings_config_mismatches(
|
||||
running_config: dict[str, Any],
|
||||
agent_type: str,
|
||||
) -> list[str]:
|
||||
"""Return restart reasons when a running proxy lacks target agent savings."""
|
||||
|
||||
if agent_type not in _AGENT_SAVINGS_TARGET_AGENTS:
|
||||
return []
|
||||
|
||||
desired_env = os.environ.copy()
|
||||
apply_agent_savings_env_defaults(desired_env)
|
||||
checks: tuple[tuple[str, str, str, str], ...] = (
|
||||
("HEADROOM_SAVINGS_PROFILE", "savings_profile", "savings-profile", "str"),
|
||||
("HEADROOM_TARGET_RATIO", "target_ratio", "target-ratio", "float"),
|
||||
(
|
||||
"HEADROOM_COMPRESS_USER_MESSAGES",
|
||||
"compress_user_messages",
|
||||
"compress-user-messages",
|
||||
"bool",
|
||||
),
|
||||
(
|
||||
"HEADROOM_COMPRESS_SYSTEM_MESSAGES",
|
||||
"compress_system_messages",
|
||||
"compress-system-messages",
|
||||
"bool",
|
||||
),
|
||||
("HEADROOM_PROTECT_RECENT", "protect_recent", "protect-recent", "int"),
|
||||
(
|
||||
"HEADROOM_PROTECT_ANALYSIS_CONTEXT",
|
||||
"protect_analysis_context",
|
||||
"protect-analysis-context",
|
||||
"bool",
|
||||
),
|
||||
("HEADROOM_MIN_TOKENS", "min_tokens_to_crush", "min-tokens", "int"),
|
||||
("HEADROOM_MAX_ITEMS", "max_items_after_crush", "max-items", "int"),
|
||||
(
|
||||
"HEADROOM_SMART_CRUSHER_COMPACTION",
|
||||
"smart_crusher_with_compaction",
|
||||
"smart-crusher-compaction",
|
||||
"bool",
|
||||
),
|
||||
("HEADROOM_ACCURACY_GUARD", "accuracy_guard", "accuracy-guard", "str"),
|
||||
)
|
||||
|
||||
mismatches: list[str] = []
|
||||
for env_key, config_key, label, value_type in checks:
|
||||
expected = desired_env.get(env_key)
|
||||
if expected is None:
|
||||
continue
|
||||
actual = running_config.get(config_key)
|
||||
try:
|
||||
if value_type == "float":
|
||||
matches = actual is not None and abs(float(actual) - float(expected)) < 1e-9
|
||||
elif value_type == "int":
|
||||
matches = actual is not None and int(actual) == int(expected)
|
||||
elif value_type == "bool":
|
||||
matches = actual is not None and bool(actual) is _env_bool_value(expected)
|
||||
else:
|
||||
matches = str(actual or "").strip().lower() == expected.strip().lower()
|
||||
except (TypeError, ValueError):
|
||||
matches = False
|
||||
if not matches:
|
||||
mismatches.append(label)
|
||||
|
||||
return mismatches
|
||||
|
||||
|
||||
def _proxy_active_session_count(payload: dict[str, Any] | None) -> int:
|
||||
"""Return active session count from /health runtime metadata."""
|
||||
if payload is None:
|
||||
|
|
@ -2942,19 +3027,24 @@ def copilot(
|
|||
env = os.environ.copy()
|
||||
openai_api_url: str | None = None
|
||||
copilot_proxy_token: str | None = None
|
||||
subscription_resolution = None
|
||||
if _should_use_copilot_oauth(
|
||||
backend=effective_backend,
|
||||
provider_type=provider_type,
|
||||
env=env,
|
||||
force_subscription=subscription,
|
||||
):
|
||||
client_bearer = (
|
||||
resolve_subscription_bearer_token() if subscription else resolve_client_bearer_token()
|
||||
)
|
||||
if subscription:
|
||||
subscription_resolution = resolve_subscription_bearer_token_details()
|
||||
client_bearer = (
|
||||
subscription_resolution.token if subscription_resolution is not None else None
|
||||
)
|
||||
else:
|
||||
client_bearer = resolve_client_bearer_token()
|
||||
if not client_bearer:
|
||||
raise click.ClickException(
|
||||
"GitHub Copilot subscription mode requires a reusable GitHub/Copilot bearer "
|
||||
"token, but none could be resolved. Run `copilot auth login` first, or set "
|
||||
"token, but none could be resolved. Run `headroom copilot-auth login` first, or set "
|
||||
"GITHUB_COPILOT_TOKEN / GITHUB_COPILOT_GITHUB_TOKEN."
|
||||
)
|
||||
|
||||
|
|
@ -2992,16 +3082,15 @@ def copilot(
|
|||
else "COPILOT_AUTH_MODE=github-oauth"
|
||||
),
|
||||
]
|
||||
# Resolve the Copilot API host: an explicit GITHUB_COPILOT_API_URL wins,
|
||||
# otherwise the generic public host (api.githubcopilot.com). This is the
|
||||
# same policy for --subscription and the implicit OAuth path. The
|
||||
# account-specific endpoints.api advertised by /copilot_internal/user is
|
||||
# deliberately NOT used to route — it returns a segmented host (e.g.
|
||||
# api.individual.githubcopilot.com) that does not serve newer models on
|
||||
# the responses API (#610), and it is not the host the official Copilot
|
||||
# client routes with. Accounts that require a dedicated host (enterprise /
|
||||
# data residency) set GITHUB_COPILOT_API_URL explicitly.
|
||||
openai_api_url = resolve_copilot_api_url(client_bearer)
|
||||
# Non-subscription OAuth keeps upstream's generic-host policy from
|
||||
# #610. Subscription mode can use the endpoint returned by the Copilot
|
||||
# token exchange, which is how Business accounts advertise their API
|
||||
# host without requiring users to configure it manually.
|
||||
openai_api_url = (
|
||||
subscription_resolution.api_url
|
||||
if subscription_resolution is not None
|
||||
else resolve_copilot_api_url(client_bearer)
|
||||
)
|
||||
env["GITHUB_COPILOT_API_URL"] = openai_api_url
|
||||
env["OPENAI_TARGET_API_URL"] = openai_api_url
|
||||
env_vars_display.append(f"COPILOT_PROVIDER_API_URL={openai_api_url}")
|
||||
|
|
@ -4310,6 +4399,6 @@ def unwrap_codex(port: int, no_stop_proxy: bool) -> None:
|
|||
|
||||
click.echo()
|
||||
click.echo("✓ Codex is no longer routed through the Headroom proxy.")
|
||||
if not no_stop_proxy:
|
||||
if not no_stop_proxy and status != "noop":
|
||||
_echo_unwrap_proxy_stop_status(_stop_local_proxy_for_unwrap(port), port)
|
||||
click.echo()
|
||||
|
|
|
|||
|
|
@ -58,9 +58,10 @@ from __future__ import annotations
|
|||
|
||||
import logging
|
||||
import threading
|
||||
from dataclasses import dataclass, field
|
||||
from dataclasses import dataclass, field, replace
|
||||
from typing import Any
|
||||
|
||||
from .agent_savings import apply_agent_savings_profile
|
||||
from .observability import get_otel_metrics
|
||||
from .pipeline import PipelineExtensionManager, PipelineStage, summarize_routing_markers
|
||||
from .utils import extract_user_query as _extract_user_query
|
||||
|
|
@ -133,6 +134,9 @@ class CompressConfig:
|
|||
Set to 'disabled' to skip ML compression entirely
|
||||
(only SmartCrusher + CacheAligner will run)."""
|
||||
|
||||
savings_profile: str | None = None
|
||||
"""Named high-savings profile, e.g. 'agent-90' for Codex/Claude/Cursor."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class CompressResult:
|
||||
|
|
@ -204,6 +208,9 @@ def compress(
|
|||
for key, value in kwargs.items():
|
||||
if key in config_fields:
|
||||
setattr(cfg, key, value)
|
||||
if cfg.savings_profile:
|
||||
cfg = replace(cfg)
|
||||
apply_agent_savings_profile(cfg, cfg.savings_profile)
|
||||
|
||||
pipeline = _get_pipeline()
|
||||
pipeline_extensions = PipelineExtensionManager(hooks=hooks, discover=False)
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ from __future__ import annotations
|
|||
|
||||
import asyncio
|
||||
import ctypes
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
|
|
@ -18,6 +19,7 @@ from urllib import error as urllib_error
|
|||
from urllib import request as urllib_request
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from headroom import paths
|
||||
from headroom.copilot_linux_secret import read_copilot_oauth_token as read_linux_secret_token
|
||||
from headroom.copilot_macos_keychain import read_copilot_oauth_token as read_macos_keychain_token
|
||||
|
||||
|
|
@ -27,10 +29,13 @@ DEFAULT_API_URL = "https://api.githubcopilot.com"
|
|||
DEFAULT_TOKEN_EXCHANGE_URL = "https://api.github.com/copilot_internal/v2/token"
|
||||
DEFAULT_USER_INFO_URL = "https://api.github.com/copilot_internal/user"
|
||||
DEFAULT_GITHUB_HOST = "github.com"
|
||||
COPILOT_CHAT_OAUTH_CLIENT_ID = "Iv1.b507a08c87ecfe98"
|
||||
_TOKEN_EXPIRY_BUFFER_S = 60
|
||||
_DEFAULT_INTEGRATION_ID = "vscode-chat"
|
||||
_DEFAULT_EDITOR_VERSION = "vscode/1.104.1"
|
||||
_DEFAULT_USER_AGENT = "GitHubCopilotChat/0.1"
|
||||
_DEFAULT_EDITOR_VERSION = "vscode/1.107.0"
|
||||
_DEFAULT_USER_AGENT = "GitHubCopilotChat/0.35.0"
|
||||
_DEFAULT_EDITOR_PLUGIN_VERSION = "copilot-chat/0.35.0"
|
||||
_DEFAULT_COPILOT_INTEGRATION_ID = "vscode-chat"
|
||||
_DEVICE_CODE_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:device_code"
|
||||
|
||||
_API_TOKEN_ENV_VARS = (
|
||||
"GITHUB_COPILOT_API_TOKEN",
|
||||
|
|
@ -80,16 +85,137 @@ class CopilotTokenCandidate:
|
|||
validate_for_subscription: bool = True
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CopilotSubscriptionTokenResolution:
|
||||
"""A Copilot subscription token plus safe routing metadata."""
|
||||
|
||||
token: str
|
||||
source: str
|
||||
confidence: str
|
||||
api_url: str
|
||||
token_fingerprint: str
|
||||
|
||||
|
||||
def token_fingerprint(token: str) -> str:
|
||||
"""Return a stable non-secret fingerprint for comparing token handoffs."""
|
||||
|
||||
digest = hashlib.sha256(token.encode("utf-8", errors="ignore")).hexdigest()
|
||||
return f"sha256:{digest[:12]}"
|
||||
|
||||
|
||||
def _github_host() -> str:
|
||||
return (os.environ.get("GITHUB_COPILOT_HOST") or DEFAULT_GITHUB_HOST).strip().lower()
|
||||
|
||||
|
||||
def headroom_copilot_auth_path() -> Path:
|
||||
"""Return the path where Headroom stores its Copilot OAuth token."""
|
||||
|
||||
override = os.environ.get("HEADROOM_COPILOT_AUTH_FILE", "").strip()
|
||||
if override:
|
||||
return Path(override).expanduser()
|
||||
return paths.workspace_dir() / "copilot_auth.json"
|
||||
|
||||
|
||||
def normalize_copilot_enterprise_url(enterprise_url: str) -> str:
|
||||
"""Normalize a GitHub Enterprise URL or domain."""
|
||||
|
||||
return enterprise_url.strip().replace("https://", "").replace("http://", "").rstrip("/")
|
||||
|
||||
|
||||
def _enterprise_hostname(enterprise_url: str) -> str:
|
||||
normalized = normalize_copilot_enterprise_url(enterprise_url)
|
||||
if not normalized:
|
||||
return ""
|
||||
parsed = urlparse(f"https://{normalized}")
|
||||
return (parsed.hostname or normalized.split("/", 1)[0]).lower()
|
||||
|
||||
|
||||
def _copilot_subdomain_enterprise_host(enterprise_url: str) -> str | None:
|
||||
"""Return a host that supports api.<host> and copilot-api.<host> URLs.
|
||||
|
||||
GitHub.com Enterprise Cloud URLs such as ``github.com/enterprises/acme``
|
||||
identify an account, not an API hostname.
|
||||
"""
|
||||
|
||||
host = _enterprise_hostname(enterprise_url)
|
||||
for prefix in ("copilot-api.", "api."):
|
||||
if host.startswith(prefix):
|
||||
host = host[len(prefix) :]
|
||||
break
|
||||
if not host or host in {"github.com", "www.github.com", "api.github.com"}:
|
||||
return None
|
||||
return host
|
||||
|
||||
|
||||
def copilot_api_url_from_enterprise_url(enterprise_url: str) -> str:
|
||||
"""Return a Copilot API base for GitHub Enterprise Server/custom domains."""
|
||||
|
||||
host = _copilot_subdomain_enterprise_host(enterprise_url)
|
||||
if host is None:
|
||||
return DEFAULT_API_URL
|
||||
return f"https://copilot-api.{host}"
|
||||
|
||||
|
||||
def _configured_enterprise_domain() -> str | None:
|
||||
enterprise_url = (
|
||||
os.environ.get("GITHUB_COPILOT_ENTERPRISE_URL", "").strip()
|
||||
or os.environ.get("GITHUB_COPILOT_ENTERPRISE_DOMAIN", "").strip()
|
||||
)
|
||||
if not enterprise_url:
|
||||
return None
|
||||
return _copilot_subdomain_enterprise_host(enterprise_url)
|
||||
|
||||
|
||||
def _configured_api_url() -> str:
|
||||
api_url = os.environ.get("GITHUB_COPILOT_API_URL", "").strip()
|
||||
if api_url:
|
||||
return api_url.rstrip("/")
|
||||
|
||||
enterprise_domain = _configured_enterprise_domain()
|
||||
if enterprise_domain:
|
||||
return copilot_api_url_from_enterprise_url(enterprise_domain).rstrip("/")
|
||||
|
||||
return DEFAULT_API_URL
|
||||
|
||||
|
||||
def _github_oauth_domain(domain: str | None = None) -> str:
|
||||
raw = (domain or DEFAULT_GITHUB_HOST).strip()
|
||||
if not raw:
|
||||
return DEFAULT_GITHUB_HOST
|
||||
host = _enterprise_hostname(raw)
|
||||
return host or DEFAULT_GITHUB_HOST
|
||||
|
||||
|
||||
def _github_oauth_urls(domain: str) -> dict[str, str]:
|
||||
normalized = _github_oauth_domain(domain)
|
||||
return {
|
||||
"device_code": f"https://{normalized}/login/device/code",
|
||||
"access_token": f"https://{normalized}/login/oauth/access_token",
|
||||
}
|
||||
|
||||
|
||||
def _token_exchange_url() -> str:
|
||||
return os.environ.get("GITHUB_COPILOT_TOKEN_EXCHANGE_URL", DEFAULT_TOKEN_EXCHANGE_URL).strip()
|
||||
override = os.environ.get("GITHUB_COPILOT_TOKEN_EXCHANGE_URL", "").strip()
|
||||
if override:
|
||||
return override
|
||||
|
||||
enterprise_domain = _configured_enterprise_domain()
|
||||
if enterprise_domain:
|
||||
return f"https://api.{enterprise_domain}/copilot_internal/v2/token"
|
||||
|
||||
return DEFAULT_TOKEN_EXCHANGE_URL
|
||||
|
||||
|
||||
def _user_info_url() -> str:
|
||||
return os.environ.get("GITHUB_COPILOT_USER_INFO_URL", DEFAULT_USER_INFO_URL).strip()
|
||||
override = os.environ.get("GITHUB_COPILOT_USER_INFO_URL", "").strip()
|
||||
if override:
|
||||
return override
|
||||
|
||||
enterprise_domain = _configured_enterprise_domain()
|
||||
if enterprise_domain:
|
||||
return f"https://api.{enterprise_domain}/copilot_internal/user"
|
||||
|
||||
return DEFAULT_USER_INFO_URL
|
||||
|
||||
|
||||
def _should_exchange_oauth_token() -> bool:
|
||||
|
|
@ -273,6 +399,143 @@ def _entry_expired(entry: dict[str, Any]) -> bool:
|
|||
return False
|
||||
|
||||
|
||||
def read_headroom_copilot_oauth_token() -> str | None:
|
||||
"""Return Headroom's saved Copilot OAuth token, if one is available."""
|
||||
|
||||
try:
|
||||
payload = json.loads(headroom_copilot_auth_path().read_text(encoding="utf-8"))
|
||||
except FileNotFoundError:
|
||||
return None
|
||||
except Exception as exc:
|
||||
logger.debug("Unable to read Headroom Copilot auth file: %s", exc)
|
||||
return None
|
||||
|
||||
if not isinstance(payload, dict) or payload.get("type") != "oauth":
|
||||
return None
|
||||
token = payload.get("refresh")
|
||||
return token.strip() if isinstance(token, str) and token.strip() else None
|
||||
|
||||
|
||||
def save_headroom_copilot_oauth_token(
|
||||
token: str,
|
||||
*,
|
||||
domain: str = DEFAULT_GITHUB_HOST,
|
||||
) -> Path:
|
||||
"""Persist the Copilot OAuth token returned by GitHub device login."""
|
||||
|
||||
token = token.strip()
|
||||
if not token:
|
||||
raise ValueError("Copilot OAuth token must not be empty.")
|
||||
|
||||
path = headroom_copilot_auth_path()
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
body: dict[str, Any] = {
|
||||
"type": "oauth",
|
||||
"provider": "github-copilot",
|
||||
"refresh": token,
|
||||
"domain": _github_oauth_domain(domain),
|
||||
"created_at": int(time.time()),
|
||||
}
|
||||
path.write_text(json.dumps(body, indent=2, sort_keys=True) + "\n", encoding="utf-8")
|
||||
try:
|
||||
path.chmod(0o600)
|
||||
except OSError:
|
||||
pass
|
||||
return path
|
||||
|
||||
|
||||
def start_copilot_device_authorization(
|
||||
*,
|
||||
domain: str = DEFAULT_GITHUB_HOST,
|
||||
timeout: float = 10.0,
|
||||
) -> dict[str, Any]:
|
||||
"""Start the GitHub Copilot OAuth device-code flow."""
|
||||
|
||||
urls = _github_oauth_urls(domain)
|
||||
body = json.dumps(
|
||||
{
|
||||
"client_id": COPILOT_CHAT_OAUTH_CLIENT_ID,
|
||||
"scope": "read:user",
|
||||
},
|
||||
separators=(",", ":"),
|
||||
).encode("utf-8")
|
||||
request = urllib_request.Request(
|
||||
urls["device_code"],
|
||||
data=body,
|
||||
headers={
|
||||
"Accept": "application/json",
|
||||
"Content-Type": "application/json",
|
||||
"User-Agent": _DEFAULT_USER_AGENT,
|
||||
},
|
||||
method="POST",
|
||||
)
|
||||
with urllib_request.urlopen(request, timeout=timeout) as response:
|
||||
payload = json.loads(response.read().decode("utf-8", errors="replace"))
|
||||
if not isinstance(payload, dict):
|
||||
raise RuntimeError("GitHub device authorization returned an invalid response.")
|
||||
return payload
|
||||
|
||||
|
||||
def poll_copilot_device_authorization(
|
||||
device_code: str,
|
||||
*,
|
||||
domain: str = DEFAULT_GITHUB_HOST,
|
||||
interval: int = 5,
|
||||
expires_in: int = 900,
|
||||
timeout: float = 10.0,
|
||||
) -> str:
|
||||
"""Poll GitHub until the device-code OAuth flow returns an access token."""
|
||||
|
||||
urls = _github_oauth_urls(domain)
|
||||
deadline = time.time() + max(1, expires_in)
|
||||
poll_interval = max(1, interval)
|
||||
while time.time() < deadline:
|
||||
body = json.dumps(
|
||||
{
|
||||
"client_id": COPILOT_CHAT_OAUTH_CLIENT_ID,
|
||||
"device_code": device_code,
|
||||
"grant_type": _DEVICE_CODE_GRANT_TYPE,
|
||||
},
|
||||
separators=(",", ":"),
|
||||
).encode("utf-8")
|
||||
request = urllib_request.Request(
|
||||
urls["access_token"],
|
||||
data=body,
|
||||
headers={
|
||||
"Accept": "application/json",
|
||||
"Content-Type": "application/json",
|
||||
"User-Agent": _DEFAULT_USER_AGENT,
|
||||
},
|
||||
method="POST",
|
||||
)
|
||||
with urllib_request.urlopen(request, timeout=timeout) as response:
|
||||
payload = json.loads(response.read().decode("utf-8", errors="replace"))
|
||||
if not isinstance(payload, dict):
|
||||
raise RuntimeError("GitHub device authorization returned an invalid response.")
|
||||
|
||||
access_token = payload.get("access_token")
|
||||
if isinstance(access_token, str) and access_token.strip():
|
||||
return access_token.strip()
|
||||
|
||||
error = str(payload.get("error") or "").strip()
|
||||
if error == "authorization_pending":
|
||||
time.sleep(poll_interval)
|
||||
continue
|
||||
if error == "slow_down":
|
||||
poll_interval += 5
|
||||
time.sleep(poll_interval)
|
||||
continue
|
||||
if error == "expired_token":
|
||||
raise RuntimeError("GitHub device authorization expired.")
|
||||
if error:
|
||||
description = str(payload.get("error_description") or error).strip()
|
||||
raise RuntimeError(f"GitHub device authorization failed: {description}")
|
||||
|
||||
time.sleep(poll_interval)
|
||||
|
||||
raise RuntimeError("GitHub device authorization expired.")
|
||||
|
||||
|
||||
def _extract_oauth_token(entry: dict[str, Any]) -> str | None:
|
||||
if _entry_expired(entry):
|
||||
return None
|
||||
|
|
@ -318,6 +581,16 @@ def iter_oauth_token_candidates() -> list[CopilotTokenCandidate]:
|
|||
|
||||
candidates: list[CopilotTokenCandidate] = []
|
||||
|
||||
headroom_copilot_token = read_headroom_copilot_oauth_token()
|
||||
if headroom_copilot_token:
|
||||
candidates.append(
|
||||
CopilotTokenCandidate(
|
||||
token=headroom_copilot_token,
|
||||
source=f"headroom-copilot-auth:{headroom_copilot_auth_path()}",
|
||||
confidence="copilot-oauth",
|
||||
)
|
||||
)
|
||||
|
||||
for env_var in _COPILOT_OAUTH_TOKEN_ENV_VARS:
|
||||
token = os.environ.get(env_var, "").strip()
|
||||
if token:
|
||||
|
|
@ -438,28 +711,182 @@ def resolve_client_bearer_token() -> str | None:
|
|||
return read_cached_oauth_token()
|
||||
|
||||
|
||||
def resolve_subscription_bearer_token() -> str | None:
|
||||
"""Return the first discovered token that GitHub accepts for Copilot subscription APIs."""
|
||||
def _copilot_chat_header_defaults() -> dict[str, str]:
|
||||
return {
|
||||
"User-Agent": os.environ.get("GITHUB_COPILOT_USER_AGENT", _DEFAULT_USER_AGENT).strip()
|
||||
or _DEFAULT_USER_AGENT,
|
||||
"Editor-Version": os.environ.get(
|
||||
"GITHUB_COPILOT_EDITOR_VERSION", _DEFAULT_EDITOR_VERSION
|
||||
).strip()
|
||||
or _DEFAULT_EDITOR_VERSION,
|
||||
"Editor-Plugin-Version": os.environ.get(
|
||||
"GITHUB_COPILOT_EDITOR_PLUGIN_VERSION",
|
||||
_DEFAULT_EDITOR_PLUGIN_VERSION,
|
||||
).strip()
|
||||
or _DEFAULT_EDITOR_PLUGIN_VERSION,
|
||||
"Copilot-Integration-Id": os.environ.get(
|
||||
"GITHUB_COPILOT_INTEGRATION_ID",
|
||||
_DEFAULT_COPILOT_INTEGRATION_ID,
|
||||
).strip()
|
||||
or _DEFAULT_COPILOT_INTEGRATION_ID,
|
||||
}
|
||||
|
||||
|
||||
def _set_header_default(headers: dict[str, str], name: str, value: str) -> None:
|
||||
"""Set a header default without duplicating case-insensitive equivalents."""
|
||||
|
||||
name_lower = name.lower()
|
||||
if any(existing.lower() == name_lower for existing in headers):
|
||||
return
|
||||
headers[name] = value
|
||||
|
||||
|
||||
def _copilot_token_exchange_headers(oauth_token: str) -> dict[str, str]:
|
||||
return {
|
||||
"Accept": "application/json",
|
||||
"Authorization": f"Bearer {oauth_token}",
|
||||
**_copilot_chat_header_defaults(),
|
||||
}
|
||||
|
||||
|
||||
def _api_url_from_payload(payload: dict[str, Any] | None) -> str | None:
|
||||
endpoints = payload.get("endpoints") if isinstance(payload, dict) else None
|
||||
api_url = endpoints.get("api") if isinstance(endpoints, dict) else None
|
||||
if isinstance(api_url, str) and api_url.strip():
|
||||
return api_url.strip().rstrip("/")
|
||||
return None
|
||||
|
||||
|
||||
def _subscription_api_url_from_user_info_payload(payload: dict[str, Any] | None) -> str:
|
||||
api_url = _api_url_from_payload(payload)
|
||||
if not api_url:
|
||||
return _configured_api_url()
|
||||
|
||||
host = urlparse(api_url).netloc.lower()
|
||||
if host in {"api.githubcopilot.com", "api.individual.githubcopilot.com"}:
|
||||
return _configured_api_url()
|
||||
if host.endswith(".githubcopilot.com"):
|
||||
return api_url
|
||||
return _configured_api_url()
|
||||
|
||||
|
||||
def _subscription_api_url_from_user_info(oauth_token: str) -> str:
|
||||
return _subscription_api_url_from_user_info_payload(_fetch_copilot_user_info(oauth_token))
|
||||
|
||||
|
||||
def _api_url_from_exchange_payload(payload: dict[str, Any], *, oauth_token: str) -> str:
|
||||
configured = _configured_api_url()
|
||||
if configured != DEFAULT_API_URL:
|
||||
return configured
|
||||
|
||||
api_url = _api_url_from_payload(payload)
|
||||
if api_url:
|
||||
return api_url
|
||||
|
||||
return _subscription_api_url_from_user_info(oauth_token)
|
||||
|
||||
|
||||
def _subscription_resolution(
|
||||
*,
|
||||
token: str,
|
||||
source: str,
|
||||
confidence: str,
|
||||
api_url: str,
|
||||
) -> CopilotSubscriptionTokenResolution:
|
||||
return CopilotSubscriptionTokenResolution(
|
||||
token=token,
|
||||
source=source,
|
||||
confidence=confidence,
|
||||
api_url=api_url,
|
||||
token_fingerprint=token_fingerprint(token),
|
||||
)
|
||||
|
||||
|
||||
def _subscription_resolution_from_token_exchange(
|
||||
candidate: CopilotTokenCandidate,
|
||||
) -> CopilotSubscriptionTokenResolution | None:
|
||||
"""Exchange a reusable GitHub OAuth token for a Copilot API token."""
|
||||
|
||||
try:
|
||||
payload = CopilotTokenProvider._exchange_token_sync(
|
||||
_copilot_token_exchange_headers(candidate.token)
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.debug(
|
||||
"Unable to exchange Copilot OAuth token from %s via %s: %s",
|
||||
candidate.source,
|
||||
_token_exchange_url(),
|
||||
exc,
|
||||
)
|
||||
return None
|
||||
|
||||
token = str(payload.get("token") or "").strip()
|
||||
if not token:
|
||||
logger.debug("Copilot token exchange from %s returned no token", candidate.source)
|
||||
return None
|
||||
|
||||
return _subscription_resolution(
|
||||
token=token,
|
||||
source=f"{candidate.source}:token-exchange",
|
||||
confidence="copilot-token-exchange",
|
||||
api_url=_api_url_from_exchange_payload(payload, oauth_token=candidate.token),
|
||||
)
|
||||
|
||||
|
||||
def resolve_subscription_bearer_token_details() -> CopilotSubscriptionTokenResolution | None:
|
||||
"""Return the first discovered token that GitHub accepts for subscription APIs."""
|
||||
|
||||
for env_var in _API_TOKEN_ENV_VARS:
|
||||
token = os.environ.get(env_var, "").strip()
|
||||
if token and _fetch_copilot_user_info(token) is not None:
|
||||
return token
|
||||
if not token:
|
||||
continue
|
||||
payload = _fetch_copilot_user_info(token)
|
||||
if payload is not None:
|
||||
return _subscription_resolution(
|
||||
token=token,
|
||||
source=f"env:{env_var}",
|
||||
confidence="explicit-api-token",
|
||||
api_url=_subscription_api_url_from_user_info_payload(payload),
|
||||
)
|
||||
|
||||
for candidate in iter_oauth_token_candidates():
|
||||
if not candidate.validate_for_subscription:
|
||||
continue
|
||||
if _fetch_copilot_user_info(candidate.token) is not None:
|
||||
if _is_copilot_api_token(candidate.token):
|
||||
payload = _fetch_copilot_user_info(candidate.token)
|
||||
if payload is not None:
|
||||
logger.debug(
|
||||
"Using Copilot API subscription token from %s (%s)",
|
||||
candidate.source,
|
||||
candidate.confidence,
|
||||
)
|
||||
return _subscription_resolution(
|
||||
token=candidate.token,
|
||||
source=candidate.source,
|
||||
confidence=candidate.confidence,
|
||||
api_url=_subscription_api_url_from_user_info_payload(payload),
|
||||
)
|
||||
continue
|
||||
|
||||
exchanged = _subscription_resolution_from_token_exchange(candidate)
|
||||
if exchanged is not None:
|
||||
logger.debug(
|
||||
"Using Copilot subscription token from %s (%s)",
|
||||
"Using exchanged Copilot subscription token from %s (%s)",
|
||||
candidate.source,
|
||||
candidate.confidence,
|
||||
)
|
||||
return candidate.token
|
||||
return exchanged
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def resolve_subscription_bearer_token() -> str | None:
|
||||
"""Return the first discovered token that GitHub accepts for Copilot subscription APIs."""
|
||||
|
||||
resolution = resolve_subscription_bearer_token_details()
|
||||
return resolution.token if resolution is not None else None
|
||||
|
||||
|
||||
def has_oauth_auth() -> bool:
|
||||
"""Return True when existing Copilot auth can be reused."""
|
||||
|
||||
|
|
@ -473,7 +900,30 @@ def is_copilot_api_url(url: str | None) -> bool:
|
|||
return False
|
||||
parsed = urlparse(url)
|
||||
host = parsed.netloc.lower() or parsed.path.lower()
|
||||
return "githubcopilot.com" in host
|
||||
configured_host = urlparse(_configured_api_url()).netloc.lower()
|
||||
if configured_host and host == configured_host:
|
||||
return True
|
||||
hostname = (parsed.hostname or host.split("/", 1)[0]).lower()
|
||||
return _is_public_copilot_api_host(hostname) or _is_ghe_copilot_api_host(hostname)
|
||||
|
||||
|
||||
def _is_public_copilot_api_host(host: str) -> bool:
|
||||
"""Return True for GitHub-hosted Copilot API domains."""
|
||||
|
||||
return host == "githubcopilot.com" or host.endswith(".githubcopilot.com")
|
||||
|
||||
|
||||
def _is_ghe_copilot_api_host(host: str) -> bool:
|
||||
"""Return True for GitHub Enterprise Copilot API hosts.
|
||||
|
||||
GHE Copilot deployments use hosts like ``copilot-api.<tenant>.ghe.com``.
|
||||
Restrict this to the Copilot API subdomain so unrelated GHE hosts do not
|
||||
receive Copilot auth headers or Copilot-specific path normalization.
|
||||
"""
|
||||
|
||||
return host == "copilot-api.ghe.com" or (
|
||||
host.startswith("copilot-api.") and host.endswith(".ghe.com")
|
||||
)
|
||||
|
||||
|
||||
def build_copilot_upstream_url(base_url: str, path: str) -> str:
|
||||
|
|
@ -506,8 +956,7 @@ def resolve_copilot_api_url(oauth_token: str | None = None) -> str:
|
|||
"""
|
||||
|
||||
del oauth_token # reserved; routing no longer depends on a user-info lookup
|
||||
override = os.environ.get("GITHUB_COPILOT_API_URL", "").strip()
|
||||
return override or DEFAULT_API_URL
|
||||
return _configured_api_url()
|
||||
|
||||
|
||||
def _fetch_copilot_user_info(token: str) -> dict[str, Any] | None:
|
||||
|
|
@ -517,10 +966,7 @@ def _fetch_copilot_user_info(token: str) -> dict[str, Any] | None:
|
|||
if not token:
|
||||
return None
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {token}",
|
||||
"Accept": "application/json",
|
||||
}
|
||||
headers = _copilot_token_exchange_headers(token)
|
||||
request = urllib_request.Request(_user_info_url(), headers=headers, method="GET")
|
||||
try:
|
||||
with urllib_request.urlopen(request, timeout=10.0) as response:
|
||||
|
|
@ -545,8 +991,7 @@ class CopilotTokenProvider:
|
|||
return CopilotAPIToken(
|
||||
token=explicit_api_token,
|
||||
expires_at=time.time() + 3600,
|
||||
api_url=os.environ.get("GITHUB_COPILOT_API_URL", DEFAULT_API_URL).strip()
|
||||
or DEFAULT_API_URL,
|
||||
api_url=_configured_api_url(),
|
||||
)
|
||||
|
||||
cached = self._cached
|
||||
|
|
@ -566,8 +1011,7 @@ class CopilotTokenProvider:
|
|||
direct_token = CopilotAPIToken(
|
||||
token=oauth_token,
|
||||
expires_at=time.time() + 3600,
|
||||
api_url=os.environ.get("GITHUB_COPILOT_API_URL", DEFAULT_API_URL).strip()
|
||||
or DEFAULT_API_URL,
|
||||
api_url=_configured_api_url(),
|
||||
)
|
||||
self._cached = direct_token
|
||||
return direct_token
|
||||
|
|
@ -577,23 +1021,18 @@ class CopilotTokenProvider:
|
|||
return exchanged
|
||||
|
||||
async def _exchange_token(self, oauth_token: str) -> CopilotAPIToken:
|
||||
headers = {
|
||||
"Authorization": f"Bearer {oauth_token}",
|
||||
"Accept": "application/json",
|
||||
"Editor-Version": os.environ.get(
|
||||
"GITHUB_COPILOT_EDITOR_VERSION", _DEFAULT_EDITOR_VERSION
|
||||
),
|
||||
"User-Agent": _DEFAULT_USER_AGENT,
|
||||
}
|
||||
headers = _copilot_token_exchange_headers(oauth_token)
|
||||
payload = await asyncio.to_thread(self._exchange_token_sync, headers)
|
||||
token = str(payload.get("token") or "").strip()
|
||||
if not token:
|
||||
raise RuntimeError("Copilot token exchange returned an empty token.")
|
||||
|
||||
expires_at = _parse_expiry(payload.get("expires_at")) or (time.time() + 1800)
|
||||
raw_endpoints = payload.get("endpoints")
|
||||
endpoints: dict[str, Any] = raw_endpoints if isinstance(raw_endpoints, dict) else {}
|
||||
api_url = str(endpoints.get("api") or DEFAULT_API_URL).strip() or DEFAULT_API_URL
|
||||
api_url = await asyncio.to_thread(
|
||||
_api_url_from_exchange_payload,
|
||||
payload,
|
||||
oauth_token=oauth_token,
|
||||
)
|
||||
refresh_in = payload.get("refresh_in")
|
||||
sku = payload.get("sku")
|
||||
return CopilotAPIToken(
|
||||
|
|
@ -667,15 +1106,8 @@ async def apply_copilot_api_auth(headers: dict[str, str], *, url: str) -> dict[s
|
|||
if not is_copilot_api_url(url):
|
||||
return resolved
|
||||
|
||||
lower_keys = {k.lower() for k in resolved}
|
||||
if "copilot-integration-id" not in lower_keys:
|
||||
resolved["Copilot-Integration-Id"] = os.environ.get(
|
||||
"GITHUB_COPILOT_INTEGRATION_ID", _DEFAULT_INTEGRATION_ID
|
||||
)
|
||||
if "editor-version" not in lower_keys:
|
||||
resolved["editor-version"] = os.environ.get(
|
||||
"GITHUB_COPILOT_EDITOR_VERSION", _DEFAULT_EDITOR_VERSION
|
||||
)
|
||||
for name, value in _copilot_chat_header_defaults().items():
|
||||
_set_header_default(resolved, name, value)
|
||||
|
||||
incoming_auth = next((v for k, v in resolved.items() if k.lower() == "authorization"), None)
|
||||
if incoming_auth:
|
||||
|
|
@ -685,6 +1117,9 @@ async def apply_copilot_api_auth(headers: dict[str, str], *, url: str) -> dict[s
|
|||
"apply_copilot_api_auth: passing through client token kind=%s",
|
||||
_token_kind(raw_token),
|
||||
)
|
||||
for key in list(resolved):
|
||||
if key.lower() == "x-api-key":
|
||||
resolved.pop(key)
|
||||
return resolved
|
||||
logger.info(
|
||||
"apply_copilot_api_auth: incoming token not suitable (kind=%s), will replace",
|
||||
|
|
@ -693,7 +1128,7 @@ async def apply_copilot_api_auth(headers: dict[str, str], *, url: str) -> dict[s
|
|||
|
||||
token = await get_copilot_token_provider().get_api_token()
|
||||
for key in list(resolved):
|
||||
if key.lower() == "authorization":
|
||||
if key.lower() in {"authorization", "x-api-key"}:
|
||||
resolved.pop(key)
|
||||
resolved["Authorization"] = f"Bearer {token.token}"
|
||||
return resolved
|
||||
|
|
|
|||
|
|
@ -133,6 +133,18 @@
|
|||
<span class="text-xs text-amber-400 font-medium">Anon Telemetry</span>
|
||||
</div>
|
||||
</template>
|
||||
<template x-if="stats.config && stats.config.savings_profile">
|
||||
<div class="inline-flex items-center gap-1.5 rounded-full border border-cyan-500/40 bg-cyan-500/10 px-2.5 py-1"
|
||||
:title="'Current proxy profile: ' + stats.config.savings_profile">
|
||||
<span class="w-1.5 h-1.5 rounded-full bg-cyan-400"></span>
|
||||
<span class="text-xs text-cyan-100">
|
||||
<span x-text="stats.config.savings_profile"></span>
|
||||
<template x-if="stats.config.target_savings_percent !== null">
|
||||
<span x-text="' · target ' + stats.config.target_savings_percent + '%'"></span>
|
||||
</template>
|
||||
</span>
|
||||
</div>
|
||||
</template>
|
||||
<div class="flex items-center gap-2">
|
||||
<span class="text-xs text-gray-500">Status</span>
|
||||
<span class="flex items-center gap-1.5">
|
||||
|
|
@ -229,6 +241,106 @@
|
|||
</div>
|
||||
</div>
|
||||
|
||||
<!-- Agent Usage -->
|
||||
<div class="bg-surface rounded-lg border border-border overflow-hidden mb-6">
|
||||
<div class="px-4 py-3 border-b border-border flex flex-col gap-2 lg:flex-row lg:items-center lg:justify-between">
|
||||
<div>
|
||||
<div class="text-sm font-medium text-gray-300">Agent Usage</div>
|
||||
<div class="text-xs text-gray-500">Before and after token usage by detected client</div>
|
||||
</div>
|
||||
<div class="flex flex-wrap items-center gap-3 text-xs">
|
||||
<span class="text-gray-500" x-text="'Coverage: ' + agentCoverageLabel"></span>
|
||||
<span class="px-2 py-0.5 rounded border border-border bg-[#141414] font-mono text-gray-400"
|
||||
x-text="formatNumber(stats.agent_usage?.totals?.requests || 0) + ' requests'"></span>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div class="p-4">
|
||||
<div class="grid grid-cols-1 md:grid-cols-4 gap-3 mb-5">
|
||||
<div class="rounded-lg border border-border bg-[#141414] p-3">
|
||||
<div class="text-xs text-gray-500 uppercase tracking-wide mb-1">Before</div>
|
||||
<div class="text-2xl font-light tabular-nums" x-text="formatNumber(stats.agent_usage?.totals?.before_tokens || 0)"></div>
|
||||
</div>
|
||||
<div class="rounded-lg border border-border bg-[#141414] p-3">
|
||||
<div class="text-xs text-gray-500 uppercase tracking-wide mb-1">After</div>
|
||||
<div class="text-2xl font-light tabular-nums text-gray-200" x-text="formatNumber(stats.agent_usage?.totals?.after_tokens || 0)"></div>
|
||||
</div>
|
||||
<div class="rounded-lg border border-border bg-[#141414] p-3">
|
||||
<div class="text-xs text-gray-500 uppercase tracking-wide mb-1">Saved</div>
|
||||
<div class="text-2xl font-light tabular-nums text-accent" x-text="formatNumber(stats.agent_usage?.totals?.tokens_saved || 0)"></div>
|
||||
</div>
|
||||
<div class="rounded-lg border border-border bg-[#141414] p-3">
|
||||
<div class="text-xs text-gray-500 uppercase tracking-wide mb-1">Savings</div>
|
||||
<div class="text-2xl font-light tabular-nums text-emerald-400" x-text="(stats.agent_usage?.totals?.savings_percent || 0).toFixed(1) + '%'"></div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<template x-if="agentRows.length > 0">
|
||||
<div class="space-y-3">
|
||||
<template x-for="agent in agentRows" :key="agent.agent">
|
||||
<div class="rounded-lg border border-border bg-[#141414] p-3">
|
||||
<div class="grid grid-cols-1 gap-3 lg:grid-cols-[minmax(160px,0.9fr)_minmax(260px,1.5fr)_minmax(260px,1.2fr)] lg:items-center">
|
||||
<div class="min-w-0">
|
||||
<div class="flex items-center gap-2">
|
||||
<span class="h-2.5 w-2.5 rounded-full" :class="agentDotClass(agent.agent)"></span>
|
||||
<span class="text-sm font-medium text-gray-200 truncate" x-text="agent.label"></span>
|
||||
</div>
|
||||
<div class="mt-1 text-xs text-gray-500">
|
||||
<span x-text="formatNumber(agent.requests || 0) + ' requests'"></span>
|
||||
<span class="mx-1 text-gray-700">/</span>
|
||||
<span x-text="agent.source"></span>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div>
|
||||
<div class="flex items-center justify-between text-xs mb-1">
|
||||
<span class="text-gray-500">Token flow</span>
|
||||
<span class="font-mono text-emerald-400" x-text="(agent.savings_percent || 0).toFixed(1) + '% saved'"></span>
|
||||
</div>
|
||||
<div class="h-3 w-full rounded-full bg-border overflow-hidden flex">
|
||||
<div class="h-full bg-emerald-500 transition-all duration-500"
|
||||
:style="'width:' + agentSavedWidth(agent) + '%'"></div>
|
||||
<div class="h-full bg-accent/60 transition-all duration-500"
|
||||
:style="'width:' + agentAfterWidth(agent) + '%'"></div>
|
||||
</div>
|
||||
<div class="mt-1 flex flex-wrap gap-3 text-xs text-gray-500">
|
||||
<span class="inline-flex items-center gap-1"><span class="h-2 w-2 rounded-full bg-emerald-500"></span>Saved</span>
|
||||
<span class="inline-flex items-center gap-1"><span class="h-2 w-2 rounded-full bg-accent/60"></span>Sent</span>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div class="grid grid-cols-2 sm:grid-cols-4 gap-2 text-right">
|
||||
<div>
|
||||
<div class="text-[11px] uppercase tracking-wide text-gray-500">Before</div>
|
||||
<div class="font-mono text-sm" x-text="formatNumber(agent.before_tokens || 0)"></div>
|
||||
</div>
|
||||
<div>
|
||||
<div class="text-[11px] uppercase tracking-wide text-gray-500">After</div>
|
||||
<div class="font-mono text-sm" x-text="formatNumber(agent.after_tokens || 0)"></div>
|
||||
</div>
|
||||
<div>
|
||||
<div class="text-[11px] uppercase tracking-wide text-gray-500">Saved</div>
|
||||
<div class="font-mono text-sm text-accent" x-text="formatNumber(agent.tokens_saved || 0)"></div>
|
||||
</div>
|
||||
<div>
|
||||
<div class="text-[11px] uppercase tracking-wide text-gray-500">Share</div>
|
||||
<div class="font-mono text-sm text-gray-300" x-text="(agent.share_of_saved_percent || 0).toFixed(1) + '%'"></div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</template>
|
||||
</div>
|
||||
</template>
|
||||
|
||||
<template x-if="agentRows.length === 0">
|
||||
<div class="rounded-lg border border-dashed border-border p-6 text-center text-sm text-gray-500">
|
||||
Agent usage appears after Cursor, Claude, Codex, or another client sends traffic through this proxy.
|
||||
</div>
|
||||
</template>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- Savings Breakdown -->
|
||||
<template x-if="(stats.cost?.savings_usd || 0) > 0 || (stats.cost?.cache_savings_usd || 0) > 0">
|
||||
<div class="bg-surface rounded-lg p-4 border border-border mb-6">
|
||||
|
|
@ -1763,6 +1875,48 @@
|
|||
.substring(0, 20);
|
||||
},
|
||||
|
||||
// --- Agent Usage ---
|
||||
|
||||
get agentRows() {
|
||||
return this.stats.agent_usage?.agents || [];
|
||||
},
|
||||
|
||||
get agentCoverageLabel() {
|
||||
const coverage = this.stats.agent_usage?.coverage || {};
|
||||
if (coverage.mode === 'request_logs') {
|
||||
return this.formatNumber(coverage.logged_requests || 0) + ' logged requests';
|
||||
}
|
||||
return 'aggregate fallback';
|
||||
},
|
||||
|
||||
agentSavedWidth(agent) {
|
||||
const before = agent.before_tokens || 0;
|
||||
if (before <= 0) return 0;
|
||||
return Math.min(100, Math.max(0, (agent.tokens_saved || 0) / before * 100)).toFixed(1);
|
||||
},
|
||||
|
||||
agentAfterWidth(agent) {
|
||||
const before = agent.before_tokens || 0;
|
||||
if (before <= 0) return 0;
|
||||
return Math.min(100, Math.max(0, (agent.after_tokens || 0) / before * 100)).toFixed(1);
|
||||
},
|
||||
|
||||
agentDotClass(agent) {
|
||||
const colors = {
|
||||
'claude-code': 'bg-orange-400',
|
||||
claude: 'bg-orange-400',
|
||||
codex: 'bg-emerald-400',
|
||||
cursor: 'bg-cyan-400',
|
||||
copilot: 'bg-violet-400',
|
||||
openai: 'bg-sky-400',
|
||||
anthropic: 'bg-orange-400',
|
||||
gemini: 'bg-rose-400',
|
||||
aider: 'bg-amber-400',
|
||||
unknown: 'bg-gray-500',
|
||||
};
|
||||
return colors[agent] || 'bg-gray-400';
|
||||
},
|
||||
|
||||
// --- Historical View ---
|
||||
|
||||
get historyGranularityOptions() {
|
||||
|
|
|
|||
|
|
@ -48,6 +48,22 @@ def compute_hash(text: str) -> str:
|
|||
return hashlib.md5(text.encode()).hexdigest()[:16] # nosec B324
|
||||
|
||||
|
||||
def _canonical_call_key(name: str, arguments: Any) -> str:
|
||||
"""Canonical identity for a tool invocation: name + arguments with JSON
|
||||
key order normalized, so semantically identical calls hash equal even
|
||||
when the provider serializes arguments differently."""
|
||||
if isinstance(arguments, str):
|
||||
try:
|
||||
arguments = json.loads(arguments)
|
||||
except (ValueError, TypeError):
|
||||
pass
|
||||
if isinstance(arguments, (dict, list)):
|
||||
canon = json.dumps(arguments, sort_keys=True, separators=(",", ":"), default=str)
|
||||
else:
|
||||
canon = str(arguments)
|
||||
return compute_hash(f"{name}\x00{canon}")
|
||||
|
||||
|
||||
def _extract_tool_result_text(payload: dict[str, Any]) -> str:
|
||||
"""Extract text from a tool result payload.
|
||||
|
||||
|
|
@ -157,6 +173,7 @@ def parse_message_to_blocks(
|
|||
content = message.get("content")
|
||||
if content:
|
||||
tool_result_parts: list[dict[str, Any]] = []
|
||||
tool_use_parts: list[dict[str, Any]] = []
|
||||
if isinstance(content, str):
|
||||
text = content
|
||||
elif isinstance(content, list):
|
||||
|
|
@ -172,6 +189,13 @@ def parse_message_to_blocks(
|
|||
elif isinstance(part, dict) and "toolResult" in part:
|
||||
# Strands/Bedrock converse format; same treatment.
|
||||
tool_result_parts.append(part)
|
||||
elif isinstance(part, dict) and part.get("type") == "tool_use":
|
||||
# Anthropic Messages format: call side of the tool unit;
|
||||
# collect for dedicated tool_call blocks below.
|
||||
tool_use_parts.append(part)
|
||||
elif isinstance(part, dict) and "toolUse" in part:
|
||||
# Strands/Bedrock converse format; same treatment.
|
||||
tool_use_parts.append(part)
|
||||
elif isinstance(part, str):
|
||||
text_parts.append(part)
|
||||
text = "\n".join(text_parts)
|
||||
|
|
@ -242,6 +266,33 @@ def parse_message_to_blocks(
|
|||
)
|
||||
blocks.extend(tr_blocks)
|
||||
|
||||
for part in tool_use_parts:
|
||||
payload = part["toolUse"] if "toolUse" in part else part
|
||||
if not isinstance(payload, dict):
|
||||
continue
|
||||
tu_name = payload.get("name") or "unknown"
|
||||
tu_args = payload.get("input", {})
|
||||
tu_id = payload.get("toolUseId") if "toolUse" in part else payload.get("id")
|
||||
try:
|
||||
tu_args_text = json.dumps(tu_args, sort_keys=True, default=str)
|
||||
except (TypeError, ValueError):
|
||||
tu_args_text = str(tu_args)
|
||||
tu_text = f"{tu_name}({tu_args_text})"
|
||||
blocks.append(
|
||||
Block(
|
||||
kind="tool_call",
|
||||
text=tu_text,
|
||||
tokens_est=tokenizer.count_text(tu_text) + 10,
|
||||
content_hash=compute_hash(tu_text),
|
||||
source_index=index,
|
||||
flags={
|
||||
"tool_call_id": tu_id,
|
||||
"function_name": tu_name,
|
||||
"call_key": _canonical_call_key(tu_name, tu_args),
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
# Handle tool calls (assistant messages with tool_calls)
|
||||
tool_calls = message.get("tool_calls")
|
||||
if tool_calls:
|
||||
|
|
@ -259,6 +310,9 @@ def parse_message_to_blocks(
|
|||
flags={
|
||||
"tool_call_id": tc.get("id"),
|
||||
"function_name": func.get("name"),
|
||||
"call_key": _canonical_call_key(
|
||||
func.get("name") or "unknown", func.get("arguments", "")
|
||||
),
|
||||
},
|
||||
)
|
||||
)
|
||||
|
|
@ -315,6 +369,7 @@ def parse_messages(
|
|||
# at more than one position means the agent re-fetched something already
|
||||
# in context — an over-compression signal (#853). The first serve is
|
||||
# free; every repeat is counted as waste.
|
||||
counted_results: set[int] = set()
|
||||
reread_groups: dict[str, list[Block]] = {}
|
||||
for block in all_blocks:
|
||||
if block.kind == "tool_result" and block.tokens_est >= REREAD_MIN_TOKENS:
|
||||
|
|
@ -337,6 +392,44 @@ def parse_messages(
|
|||
prev_index = block.source_index
|
||||
if not is_polling:
|
||||
total_waste.reread_tokens += block.tokens_est
|
||||
counted_results.add(id(block))
|
||||
|
||||
# Re-issued-call detection: the agent invoking the same tool with the
|
||||
# same arguments again is a re-fetch even when the result bytes differ
|
||||
# (timestamps, mtimes, ordering defeat the content-hash pass above).
|
||||
# Same polling guard and size floor as above, applied to the repeat
|
||||
# invocation's result; results the content-hash pass already counted
|
||||
# are skipped so identical-content repeats are never counted twice.
|
||||
results_by_call_id: dict[str, Block] = {}
|
||||
for block in all_blocks:
|
||||
if block.kind == "tool_result":
|
||||
tc_id = block.flags.get("tool_call_id")
|
||||
if tc_id and tc_id not in results_by_call_id:
|
||||
results_by_call_id[tc_id] = block
|
||||
|
||||
call_groups: dict[str, list[Block]] = {}
|
||||
for block in all_blocks:
|
||||
if block.kind == "tool_call":
|
||||
call_key = block.flags.get("call_key")
|
||||
if call_key:
|
||||
call_groups.setdefault(call_key, []).append(block)
|
||||
|
||||
for group in call_groups.values():
|
||||
prev_index = group[0].source_index
|
||||
for block in group:
|
||||
if block.source_index == prev_index:
|
||||
continue
|
||||
is_polling = block.source_index - prev_index <= REREAD_ADJACENT_GAP
|
||||
prev_index = block.source_index
|
||||
if is_polling:
|
||||
continue
|
||||
result = results_by_call_id.get(block.flags.get("tool_call_id") or "")
|
||||
if result is None or result.tokens_est < REREAD_MIN_TOKENS:
|
||||
continue
|
||||
if id(result) in counted_results:
|
||||
continue
|
||||
total_waste.reread_tokens += result.tokens_est
|
||||
counted_results.add(id(result))
|
||||
|
||||
# Compute block breakdown
|
||||
breakdown: dict[str, int] = {}
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ Anthropic), not the full input price. This prevents overstating dollar savings.
|
|||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from datetime import datetime, timedelta
|
||||
|
|
@ -134,6 +135,7 @@ class PerfRecord:
|
|||
timestamp: str
|
||||
request_id: str
|
||||
model: str = ""
|
||||
client: str = ""
|
||||
num_messages: int = 0
|
||||
tokens_before: int = 0
|
||||
tokens_after: int = 0
|
||||
|
|
@ -143,7 +145,6 @@ class PerfRecord:
|
|||
cache_hit_pct: int = 0
|
||||
optimization_ms: float = 0
|
||||
transforms: list[str] = field(default_factory=list)
|
||||
client: str = ""
|
||||
|
||||
|
||||
@dataclass
|
||||
|
|
@ -234,7 +235,7 @@ def parse_log_files(last_n_hours: float = 168.0) -> PerfReport:
|
|||
report = PerfReport()
|
||||
report.requested_hours = last_n_hours
|
||||
|
||||
log_dir = _paths.log_dir()
|
||||
log_dir = _paths.log_dir() if os.environ.get("HEADROOM_WORKSPACE_DIR") else LOG_DIR
|
||||
if not log_dir.exists():
|
||||
return report
|
||||
|
||||
|
|
@ -299,6 +300,7 @@ def parse_log_files(last_n_hours: float = 168.0) -> PerfReport:
|
|||
timestamp=ts,
|
||||
request_id=m.group("rid"),
|
||||
model=kv.get("model", ""),
|
||||
client=kv.get("client", ""),
|
||||
num_messages=int(kv.get("msgs", 0)),
|
||||
tokens_before=int(kv.get("tok_before", 0)),
|
||||
tokens_after=int(kv.get("tok_after", 0)),
|
||||
|
|
@ -308,7 +310,6 @@ def parse_log_files(last_n_hours: float = 168.0) -> PerfReport:
|
|||
cache_hit_pct=int(kv.get("cache_hit_pct", 0)),
|
||||
optimization_ms=float(kv.get("opt_ms", 0)),
|
||||
transforms=transforms,
|
||||
client=kv.get("client", ""),
|
||||
)
|
||||
)
|
||||
continue
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
|
|
@ -30,6 +31,31 @@ _ORPHAN_HEADROOM_TABLE = re.compile(
|
|||
)
|
||||
|
||||
|
||||
def codex_uses_chatgpt_auth(auth_path: Path) -> bool:
|
||||
"""Whether Codex authenticated via ChatGPT OAuth (vs an OpenAI API key).
|
||||
|
||||
The account menu (profile/email/plan/usage) only renders when the active
|
||||
provider carries ``requires_openai_auth = true``, but that flag forces codex
|
||||
to demand an OpenAI OAuth login (#406) and would break API-key users. So we
|
||||
emit it only in ChatGPT-OAuth mode, read from the sibling ``auth.json``.
|
||||
"""
|
||||
try:
|
||||
data = json.loads(auth_path.read_text(encoding="utf-8"))
|
||||
except (OSError, ValueError):
|
||||
return False
|
||||
if not isinstance(data, dict):
|
||||
return False
|
||||
mode = data.get("auth_mode")
|
||||
if isinstance(mode, str):
|
||||
return mode.lower() == "chatgpt"
|
||||
# Older auth.json files predate `auth_mode`: infer from an OAuth account id.
|
||||
tokens = data.get("tokens")
|
||||
if isinstance(tokens, dict):
|
||||
account_id = tokens.get("account_id")
|
||||
return isinstance(account_id, str) and bool(account_id.strip())
|
||||
return False
|
||||
|
||||
|
||||
def build_provider_section(
|
||||
*,
|
||||
port: int,
|
||||
|
|
@ -37,12 +63,14 @@ def build_provider_section(
|
|||
marker_start: str = _CODEX_MARKER_START,
|
||||
marker_end: str = _CODEX_MARKER_END,
|
||||
include_markers: bool = True,
|
||||
requires_openai_auth: bool = False,
|
||||
) -> str:
|
||||
"""Build a managed Codex provider block (without requires_openai_auth).
|
||||
"""Build a managed Codex provider block.
|
||||
|
||||
Bug 3 (#406): requires_openai_auth must NOT appear on custom provider
|
||||
blocks — it forces codex to demand OpenAI OAuth login for local-proxy
|
||||
traffic. The built-in openai provider carries this flag; headroom does not.
|
||||
``requires_openai_auth`` is emitted only for ChatGPT-OAuth users: the flag
|
||||
is what makes codex render the account menu, but it also forces codex to
|
||||
demand an OpenAI OAuth login (#406), which breaks API-key users. Callers
|
||||
pass the result of :func:`codex_uses_chatgpt_auth`; it defaults to ``False``.
|
||||
"""
|
||||
body = (
|
||||
"[model_providers.headroom]\n"
|
||||
|
|
@ -50,6 +78,8 @@ def build_provider_section(
|
|||
f'base_url = "{proxy_base_url(port)}"\n'
|
||||
"supports_websockets = true\n"
|
||||
)
|
||||
if requires_openai_auth:
|
||||
body += "requires_openai_auth = true\n"
|
||||
if not include_markers:
|
||||
return body
|
||||
return f"{marker_start}\n{body}{marker_end}\n"
|
||||
|
|
@ -76,6 +106,7 @@ def apply_provider_scope(manifest: DeploymentManifest) -> ManagedMutation | None
|
|||
port=manifest.port,
|
||||
name="Headroom persistent proxy",
|
||||
include_markers=False,
|
||||
requires_openai_auth=codex_uses_chatgpt_auth(path.parent / "auth.json"),
|
||||
)
|
||||
+ f"{_CODEX_MARKER_END}\n"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -217,7 +217,7 @@ CLIENT_UA_MAP: tuple[tuple[str, str], ...] = (
|
|||
)
|
||||
|
||||
|
||||
def classify_client(headers: Mapping[str, Any] | Any) -> str | None:
|
||||
def classify_client(headers: Mapping[str, Any] | Any, *, default: str | None = None) -> str | None:
|
||||
"""Identify the client harness (Codex / Claude Code / aider / etc).
|
||||
|
||||
Decision order:
|
||||
|
|
@ -250,7 +250,7 @@ def classify_client(headers: Mapping[str, Any] | Any) -> str | None:
|
|||
for needle, name in CLIENT_UA_MAP:
|
||||
if needle in ua_lower:
|
||||
return name
|
||||
return None
|
||||
return default
|
||||
|
||||
|
||||
__all__ = [
|
||||
|
|
|
|||
|
|
@ -23,6 +23,7 @@ if TYPE_CHECKING:
|
|||
|
||||
import httpx
|
||||
|
||||
from headroom.agent_savings import proxy_pipeline_kwargs
|
||||
from headroom.copilot_auth import build_copilot_upstream_url
|
||||
from headroom.pipeline import PipelineStage, summarize_routing_markers
|
||||
from headroom.proxy.auth_mode import classify_auth_mode, classify_client
|
||||
|
|
@ -663,7 +664,7 @@ class AnthropicHandlerMixin:
|
|||
# Identify the harness (codex / claude-code / aider / etc.)
|
||||
# from User-Agent or X-Client. Surfaced via the funnel into
|
||||
# PERF logs and RequestLog.tags — see RequestOutcome.client.
|
||||
client = classify_client(headers)
|
||||
client = classify_client(headers, default="claude")
|
||||
# PR-A5 (P5-49): strip internal x-headroom-* from upstream-bound
|
||||
# headers AFTER `_extract_tags` reads them. Inbound bypass gating
|
||||
# uses `request.headers.get(...)` directly above; memory user-id
|
||||
|
|
@ -1066,6 +1067,7 @@ class AnthropicHandlerMixin:
|
|||
biases=biases,
|
||||
request_id=request_id,
|
||||
compression_policy=compression_policy,
|
||||
**proxy_pipeline_kwargs(self.config),
|
||||
),
|
||||
timeout=COMPRESSION_TIMEOUT_SECONDS,
|
||||
)
|
||||
|
|
@ -1106,6 +1108,7 @@ class AnthropicHandlerMixin:
|
|||
biases=biases,
|
||||
request_id=request_id,
|
||||
compression_policy=compression_policy,
|
||||
**proxy_pipeline_kwargs(self.config),
|
||||
),
|
||||
timeout=COMPRESSION_TIMEOUT_SECONDS,
|
||||
)
|
||||
|
|
@ -1137,6 +1140,7 @@ class AnthropicHandlerMixin:
|
|||
biases=biases,
|
||||
request_id=request_id,
|
||||
compression_policy=compression_policy,
|
||||
**proxy_pipeline_kwargs(self.config),
|
||||
),
|
||||
timeout=COMPRESSION_TIMEOUT_SECONDS,
|
||||
)
|
||||
|
|
@ -2550,7 +2554,7 @@ class AnthropicHandlerMixin:
|
|||
headers = dict(request.headers.items())
|
||||
headers.pop("host", None)
|
||||
headers.pop("content-length", None)
|
||||
client = classify_client(headers)
|
||||
client = classify_client(headers, default="claude")
|
||||
tags = extract_tags(headers)
|
||||
# PR-A5 (P5-49): strip internal x-headroom-* before forwarding upstream.
|
||||
from headroom.proxy.helpers import _strip_internal_headers, log_outbound_headers
|
||||
|
|
@ -2802,7 +2806,7 @@ class AnthropicHandlerMixin:
|
|||
|
||||
headers = dict(request.headers.items())
|
||||
headers.pop("host", None)
|
||||
client = classify_client(headers)
|
||||
client = classify_client(headers, default="claude")
|
||||
tags = extract_tags(headers)
|
||||
# PR-A5 (P5-49): strip internal x-headroom-* before forwarding upstream.
|
||||
from headroom.proxy.helpers import _strip_internal_headers, log_outbound_headers
|
||||
|
|
@ -2925,7 +2929,7 @@ class AnthropicHandlerMixin:
|
|||
|
||||
headers = dict(request.headers.items())
|
||||
headers.pop("host", None)
|
||||
client = classify_client(headers)
|
||||
client = classify_client(headers, default="claude")
|
||||
tags = extract_tags(headers)
|
||||
# PR-A5 (P5-49): strip internal x-headroom-* before forwarding upstream.
|
||||
from headroom.proxy.helpers import _strip_internal_headers, log_outbound_headers
|
||||
|
|
|
|||
|
|
@ -108,7 +108,11 @@ class GeminiHandlerMixin:
|
|||
return result
|
||||
|
||||
def _gemini_contents_to_messages(
|
||||
self, contents: list[dict], system_instruction: dict | None = None
|
||||
self,
|
||||
contents: list[dict],
|
||||
system_instruction: dict | None = None,
|
||||
*,
|
||||
include_function_responses: bool = False,
|
||||
) -> tuple[list[dict], set[int]]:
|
||||
"""Convert Gemini contents[] format to OpenAI messages[] format for optimization.
|
||||
|
||||
|
|
@ -119,6 +123,12 @@ class GeminiHandlerMixin:
|
|||
OpenAI format:
|
||||
messages: [{"role": "user", "content": "..."}]
|
||||
|
||||
When include_function_responses is True, functionResponse payloads are
|
||||
additionally emitted as ``role="tool"`` messages so waste-signal
|
||||
detection can see tool output (#819). That richer list is telemetry-only:
|
||||
entries with non-text parts stay in preserved_indices and are restored
|
||||
verbatim, so it must never be used as the compression input.
|
||||
|
||||
Returns:
|
||||
Tuple of (messages, preserved_indices) where preserved_indices contains
|
||||
the indices of content entries that have non-text parts (images, function
|
||||
|
|
@ -151,8 +161,29 @@ class GeminiHandlerMixin:
|
|||
if text_parts:
|
||||
messages.append({"role": role, "content": "\n".join(text_parts)})
|
||||
|
||||
if include_function_responses:
|
||||
for part in parts:
|
||||
if "functionResponse" not in part:
|
||||
continue
|
||||
payload = self._function_response_text(part["functionResponse"])
|
||||
if payload:
|
||||
messages.append({"role": "tool", "content": payload})
|
||||
|
||||
return messages, preserved_indices
|
||||
|
||||
@staticmethod
|
||||
def _function_response_text(function_response: dict) -> str:
|
||||
"""Serialize a functionResponse payload for waste-signal parsing."""
|
||||
response = function_response.get("response")
|
||||
if response is None:
|
||||
return ""
|
||||
if isinstance(response, str):
|
||||
return response
|
||||
try:
|
||||
return json.dumps(response, ensure_ascii=False, default=str)
|
||||
except (TypeError, ValueError):
|
||||
return str(response)
|
||||
|
||||
def _messages_to_gemini_contents(self, messages: list[dict]) -> tuple[list[dict], dict | None]:
|
||||
"""Convert OpenAI messages[] format back to Gemini contents[] format.
|
||||
|
||||
|
|
@ -446,11 +477,17 @@ class GeminiHandlerMixin:
|
|||
try:
|
||||
# Use OpenAI pipeline (similar message format)
|
||||
context_limit = self.openai_provider.get_context_limit(model)
|
||||
# Richer conversion incl. functionResponse payloads so tool
|
||||
# output reaches waste-signal detection (#819); telemetry-only.
|
||||
waste_messages, _ = self._gemini_contents_to_messages(
|
||||
contents, system_instruction, include_function_responses=True
|
||||
)
|
||||
result = self.openai_pipeline.apply(
|
||||
messages=messages,
|
||||
model=model,
|
||||
model_limit=context_limit,
|
||||
context=extract_user_query(messages),
|
||||
waste_messages=waste_messages,
|
||||
)
|
||||
if result.messages != messages:
|
||||
optimized_messages = result.messages
|
||||
|
|
@ -792,11 +829,17 @@ class GeminiHandlerMixin:
|
|||
if _decision.should_compress:
|
||||
try:
|
||||
context_limit = self.openai_provider.get_context_limit(model)
|
||||
# Richer conversion incl. functionResponse payloads so tool
|
||||
# output reaches waste-signal detection (#819); telemetry-only.
|
||||
waste_messages, _ = self._gemini_contents_to_messages(
|
||||
contents, system_instruction, include_function_responses=True
|
||||
)
|
||||
result = self.openai_pipeline.apply(
|
||||
messages=messages,
|
||||
model=model,
|
||||
model_limit=context_limit,
|
||||
context=extract_user_query(messages),
|
||||
waste_messages=waste_messages,
|
||||
)
|
||||
if result.messages != messages:
|
||||
optimized_messages = result.messages
|
||||
|
|
|
|||
|
|
@ -21,6 +21,7 @@ from concurrent.futures import ThreadPoolExecutor, as_completed
|
|||
from dataclasses import replace
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from headroom.proxy.helpers import (
|
||||
COMPRESSION_TIMEOUT_SECONDS,
|
||||
|
|
@ -41,6 +42,7 @@ if TYPE_CHECKING:
|
|||
|
||||
import httpx
|
||||
|
||||
from headroom.agent_savings import proxy_pipeline_kwargs
|
||||
from headroom.copilot_auth import apply_copilot_api_auth, build_copilot_upstream_url
|
||||
from headroom.pipeline import PipelineStage, summarize_routing_markers
|
||||
from headroom.proxy.auth_mode import classify_auth_mode, classify_client
|
||||
|
|
@ -139,7 +141,12 @@ def _openai_responses_unit_executor() -> ThreadPoolExecutor:
|
|||
return _OPENAI_RESPONSES_UNIT_EXECUTOR
|
||||
|
||||
|
||||
def _openai_responses_unit_cache_key(unit: Any, *, model: str) -> str:
|
||||
def _openai_responses_unit_cache_key(
|
||||
unit: Any,
|
||||
*,
|
||||
model: str,
|
||||
target_ratio: float | None = None,
|
||||
) -> str:
|
||||
text_hash = hashlib.sha256(unit.text.encode("utf-8", errors="replace")).hexdigest()
|
||||
key_payload = {
|
||||
"version": _OPENAI_RESPONSES_UNIT_CACHE_VERSION,
|
||||
|
|
@ -155,6 +162,7 @@ def _openai_responses_unit_cache_key(unit: Any, *, model: str) -> str:
|
|||
"question": unit.question,
|
||||
"bias": unit.bias,
|
||||
"metadata": unit.metadata,
|
||||
"target_ratio": target_ratio,
|
||||
"text_sha256": text_hash,
|
||||
}
|
||||
serialized = json.dumps(key_payload, sort_keys=True, separators=(",", ":"), default=str)
|
||||
|
|
@ -332,6 +340,69 @@ def _responses_input_item_text_bytes(item: Any) -> int:
|
|||
return _json_byte_len(item)
|
||||
|
||||
|
||||
_RESPONSES_OUTPUT_ITEM_TYPES = frozenset(
|
||||
{
|
||||
"custom_tool_call_output",
|
||||
"function_call_output",
|
||||
"local_shell_call_output",
|
||||
"apply_patch_call_output",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _responses_part_text(value: Any) -> str:
|
||||
"""Best-effort text from a Responses item field (string or part list)."""
|
||||
if isinstance(value, str):
|
||||
return value
|
||||
if isinstance(value, list):
|
||||
texts = []
|
||||
for part in value:
|
||||
if isinstance(part, str):
|
||||
texts.append(part)
|
||||
elif isinstance(part, dict) and isinstance(part.get("text"), str):
|
||||
texts.append(part["text"])
|
||||
return "\n".join(t for t in texts if t)
|
||||
return ""
|
||||
|
||||
|
||||
def _responses_input_to_waste_messages(instructions: Any, input_data: Any) -> list[dict[str, Any]]:
|
||||
"""Convert a Responses payload to OpenAI-style messages for waste parsing (#820).
|
||||
|
||||
Telemetry-only — never used as a compression input. Tool output items
|
||||
become ``role="tool"`` messages so tool results (where most waste lives)
|
||||
reach ``parse_messages``; ``message`` items keep their role and joined
|
||||
part text.
|
||||
"""
|
||||
messages: list[dict[str, Any]] = []
|
||||
if isinstance(instructions, str) and instructions:
|
||||
messages.append({"role": "system", "content": instructions})
|
||||
if isinstance(input_data, str):
|
||||
if input_data:
|
||||
messages.append({"role": "user", "content": input_data})
|
||||
return messages
|
||||
if not isinstance(input_data, list):
|
||||
return messages
|
||||
for item in input_data:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
if item.get("type") in _RESPONSES_OUTPUT_ITEM_TYPES:
|
||||
text = _responses_part_text(item.get("output"))
|
||||
if text:
|
||||
message: dict[str, Any] = {"role": "tool", "content": text}
|
||||
call_id = item.get("call_id")
|
||||
if isinstance(call_id, str) and call_id:
|
||||
message["tool_call_id"] = call_id
|
||||
messages.append(message)
|
||||
continue
|
||||
text = _responses_part_text(item.get("content"))
|
||||
if text:
|
||||
role = item.get("role")
|
||||
messages.append(
|
||||
{"role": role if isinstance(role, str) and role else "user", "content": text}
|
||||
)
|
||||
return messages
|
||||
|
||||
|
||||
def _openai_responses_context_budget(payload: dict[str, Any]) -> dict[str, Any]:
|
||||
payload_bytes = _json_byte_len(payload)
|
||||
buckets: dict[str, int] = {}
|
||||
|
|
@ -509,16 +580,21 @@ def _resolve_codex_routing_headers(headers: dict[str, str]) -> tuple[dict[str, s
|
|||
return resolved, False
|
||||
|
||||
|
||||
def _prefers_http1_passthrough(base_url: str) -> bool:
|
||||
"""Whether passthrough to this host must use HTTP/1.1.
|
||||
|
||||
ChatGPT's Cloudflare edge issues a managed challenge to our HTTP/2
|
||||
fingerprint on sensitive account endpoints; HTTP/1.1 is accepted.
|
||||
"""
|
||||
host = (urlparse(base_url).hostname or "").lower()
|
||||
return host == "chatgpt.com" or host.endswith(".chatgpt.com")
|
||||
|
||||
|
||||
class OpenAIHandlerMixin:
|
||||
"""Mixin providing OpenAI API handler methods for HeadroomProxy."""
|
||||
|
||||
OPENAI_RESPONSES_ROUTER_MIN_BYTES = 512
|
||||
OPENAI_RESPONSES_OUTPUT_TYPES = {
|
||||
"custom_tool_call_output",
|
||||
"function_call_output",
|
||||
"local_shell_call_output",
|
||||
"apply_patch_call_output",
|
||||
}
|
||||
OPENAI_RESPONSES_OUTPUT_TYPES = _RESPONSES_OUTPUT_ITEM_TYPES
|
||||
|
||||
def _openai_responses_unit_cache(self) -> tuple[Any, OrderedDict[str, Any]]:
|
||||
with _OPENAI_RESPONSES_UNIT_CACHE_INIT_LOCK:
|
||||
|
|
@ -646,6 +722,10 @@ class OpenAIHandlerMixin:
|
|||
if router is None:
|
||||
logger.debug("[%s] OpenAI Responses ContentRouter unavailable", request_id)
|
||||
return payload, False, 0, [], {}, [], 0
|
||||
profile_kwargs = proxy_pipeline_kwargs(getattr(self, "config", None))
|
||||
unit_target_ratio = profile_kwargs.get("target_ratio")
|
||||
if unit_target_ratio is not None:
|
||||
unit_target_ratio = float(unit_target_ratio)
|
||||
|
||||
try:
|
||||
tokenizer = self.openai_provider.get_token_counter(model)
|
||||
|
|
@ -877,7 +957,12 @@ class OpenAIHandlerMixin:
|
|||
# `elapsed_ms=60000+` in production logs even though they did
|
||||
# no work. With the semaphore deleted, this timer is honest.
|
||||
unit_started = time.perf_counter()
|
||||
result = compress_unit_with_router(routed.unit, router=router, tokenizer=tokenizer)
|
||||
result = compress_unit_with_router(
|
||||
routed.unit,
|
||||
router=router,
|
||||
tokenizer=tokenizer,
|
||||
target_ratio=unit_target_ratio,
|
||||
)
|
||||
elapsed_ms = (time.perf_counter() - unit_started) * 1000.0
|
||||
return routed.slot, result, elapsed_ms
|
||||
|
||||
|
|
@ -886,7 +971,11 @@ class OpenAIHandlerMixin:
|
|||
cache_misses: list[tuple[int, str, RoutedCompressionUnit]] = []
|
||||
cache_miss_followers: dict[str, list[int]] = {}
|
||||
for unit_idx, routed in enumerate(routed_units):
|
||||
cache_key = _openai_responses_unit_cache_key(routed.unit, model=model)
|
||||
cache_key = _openai_responses_unit_cache_key(
|
||||
routed.unit,
|
||||
model=model,
|
||||
target_ratio=unit_target_ratio,
|
||||
)
|
||||
cached = self._get_openai_responses_cached_unit(cache_key)
|
||||
if cached is not None:
|
||||
routed_results[unit_idx] = (routed.slot, cached, 0.0)
|
||||
|
|
@ -2975,6 +3064,8 @@ class OpenAIHandlerMixin:
|
|||
)
|
||||
|
||||
headers, is_chatgpt_auth = _resolve_codex_routing_headers(headers)
|
||||
if is_chatgpt_auth:
|
||||
client = "codex"
|
||||
|
||||
# Route to correct endpoint based on auth mode.
|
||||
# ChatGPT session auth (codex login) uses chatgpt.com, not api.openai.com.
|
||||
|
|
@ -3097,6 +3188,24 @@ class OpenAIHandlerMixin:
|
|||
},
|
||||
)
|
||||
|
||||
# Waste-signal detection for the Responses path (#820). The transform
|
||||
# pipeline never runs here (compression goes through CompressionUnits),
|
||||
# so parse a telemetry-only message conversion directly, behind the
|
||||
# same >100 saved-token gate as TransformPipeline.apply.
|
||||
waste_signals_dict: dict[str, int] | None = None
|
||||
if tokens_saved > 100:
|
||||
try:
|
||||
from headroom.parser import parse_messages
|
||||
|
||||
_, _, _waste = parse_messages(
|
||||
_responses_input_to_waste_messages(instructions, input_data),
|
||||
tokenizer,
|
||||
)
|
||||
if _waste.total() > 0:
|
||||
waste_signals_dict = _waste.to_dict()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
try:
|
||||
if stream:
|
||||
# Streaming for Responses API uses semantic events
|
||||
|
|
@ -3115,6 +3224,7 @@ class OpenAIHandlerMixin:
|
|||
optimization_latency,
|
||||
memory_user_id=memory_user_id,
|
||||
memory_request_ctx=memory_request_ctx,
|
||||
waste_signals=waste_signals_dict,
|
||||
)
|
||||
else:
|
||||
headers = await apply_copilot_api_auth(headers, url=url)
|
||||
|
|
@ -3302,6 +3412,7 @@ class OpenAIHandlerMixin:
|
|||
total_latency_ms=total_latency,
|
||||
overhead_ms=optimization_latency,
|
||||
transforms_applied=tuple(transforms_applied),
|
||||
waste_signals=waste_signals_dict,
|
||||
num_messages=len(messages) if isinstance(messages, list) else 0,
|
||||
tags=_resp_log_tags,
|
||||
turn_id=compute_turn_id(model, body.get("instructions"), messages),
|
||||
|
|
@ -3692,6 +3803,17 @@ class OpenAIHandlerMixin:
|
|||
|
||||
with contextlib.suppress(Exception):
|
||||
get_codex_rate_limit_state().update_from_headers(dict(_codex_handshake))
|
||||
|
||||
# Current Codex no longer ships x-codex-* on the handshake, so the
|
||||
# block above is usually a no-op. Pull the live subscription window
|
||||
# from the dedicated usage endpoint instead (throttled, scoped to
|
||||
# ChatGPT-session traffic, fire-and-forget so accept isn't blocked).
|
||||
with contextlib.suppress(Exception):
|
||||
from headroom.subscription.codex_rate_limits import (
|
||||
maybe_schedule_usage_poll,
|
||||
)
|
||||
|
||||
maybe_schedule_usage_poll(ws_headers)
|
||||
async with stage_timer.measure("accept"):
|
||||
await websocket.accept(
|
||||
subprotocol=client_subprotocols[0] if client_subprotocols else None,
|
||||
|
|
@ -4876,7 +4998,8 @@ class OpenAIHandlerMixin:
|
|||
f"cache_write={_perf_cache_write} "
|
||||
f"cache_hit_pct={_perf_cache_hit_pct} "
|
||||
f"opt_ms={overhead_delta_ms:.0f} "
|
||||
f"transforms={_summarize_transforms(transforms_applied)}"
|
||||
f"transforms={_summarize_transforms(transforms_applied)} "
|
||||
f"client={client or ''}"
|
||||
)
|
||||
|
||||
ws_recorded_input_tokens_total = ws_input_tokens_total
|
||||
|
|
@ -5816,7 +5939,10 @@ class OpenAIHandlerMixin:
|
|||
protect_recent = compress_config.get("protect_recent")
|
||||
protect_analysis_context = compress_config.get("protect_analysis_context")
|
||||
|
||||
pipeline_kwargs: dict = {"model_limit": context_limit}
|
||||
pipeline_kwargs: dict = {
|
||||
"model_limit": context_limit,
|
||||
**proxy_pipeline_kwargs(self.config),
|
||||
}
|
||||
if compress_user_messages:
|
||||
pipeline_kwargs["compress_user_messages"] = True
|
||||
if target_ratio is not None:
|
||||
|
|
@ -5912,8 +6038,17 @@ class OpenAIHandlerMixin:
|
|||
body = await request.body()
|
||||
|
||||
headers = await apply_copilot_api_auth(headers, url=url)
|
||||
# Cloudflare bot-management challenges our HTTP/2 fingerprint on
|
||||
# ChatGPT's sensitive account endpoints (/backend-api/me,
|
||||
# /backend-api/accounts/check), returning a 403 challenge page instead
|
||||
# of JSON and collapsing the Codex account menu to just "Settings".
|
||||
# Those endpoints answer fine over HTTP/1.1, so forward ChatGPT
|
||||
# passthrough on the HTTP/1.1 client. Other hosts keep HTTP/2.
|
||||
passthrough_client = self.http_client
|
||||
if _prefers_http1_passthrough(base_url):
|
||||
passthrough_client = self.http_client_h1 or self.http_client
|
||||
try:
|
||||
response = await self.http_client.request( # type: ignore[union-attr]
|
||||
response = await passthrough_client.request( # type: ignore[union-attr]
|
||||
method=request.method,
|
||||
url=url,
|
||||
headers=headers,
|
||||
|
|
|
|||
|
|
@ -663,6 +663,7 @@ class StreamingMixin:
|
|||
full_sse_data: str = "",
|
||||
parsed_response: dict[str, Any] | None = None,
|
||||
client: str | None = None,
|
||||
waste_signals: dict[str, int] | None = None,
|
||||
) -> None:
|
||||
from headroom.proxy.outcome import RequestOutcome
|
||||
|
||||
|
|
@ -786,6 +787,7 @@ class StreamingMixin:
|
|||
ttfb_ms=stream_state["ttfb_ms"] or total_latency,
|
||||
pipeline_timing=pipeline_timing,
|
||||
original_messages=original_messages,
|
||||
waste_signals=waste_signals,
|
||||
)
|
||||
await self._record_request_outcome(outcome)
|
||||
|
||||
|
|
@ -813,6 +815,7 @@ class StreamingMixin:
|
|||
mutation_reasons: list[str] | None = None,
|
||||
memory_request_ctx: Any | None = None,
|
||||
outcome_provider: str | None = None,
|
||||
waste_signals: dict[str, int] | None = None,
|
||||
) -> Response | StreamingResponse:
|
||||
"""Stream response with metrics tracking and memory tool handling.
|
||||
|
||||
|
|
@ -1064,6 +1067,7 @@ class StreamingMixin:
|
|||
prefix_tracker=prefix_tracker,
|
||||
original_messages=original_messages,
|
||||
client=client,
|
||||
waste_signals=waste_signals,
|
||||
)
|
||||
return Response(
|
||||
content=error_content,
|
||||
|
|
@ -1328,6 +1332,7 @@ class StreamingMixin:
|
|||
full_sse_data=_final_full_sse_data,
|
||||
parsed_response=parsed_response,
|
||||
client=client,
|
||||
waste_signals=waste_signals,
|
||||
)
|
||||
|
||||
return StreamingResponse(
|
||||
|
|
|
|||
|
|
@ -191,6 +191,240 @@ _build_session_summary = build_session_summary
|
|||
_merge_cost_stats = merge_cost_stats
|
||||
|
||||
|
||||
_AGENT_LABELS: dict[str, str] = {
|
||||
"claude": "Claude",
|
||||
"claude-code": "Claude",
|
||||
"claude_cli": "Claude",
|
||||
"claude-code-cli": "Claude",
|
||||
"codex": "Codex",
|
||||
"codex-cli": "Codex",
|
||||
"cursor": "Cursor",
|
||||
"copilot": "GitHub Copilot",
|
||||
"github-copilot": "GitHub Copilot",
|
||||
"aider": "Aider",
|
||||
"zed": "Zed",
|
||||
"opencode": "OpenCode",
|
||||
"openclaw": "OpenClaw",
|
||||
"gemini": "Gemini",
|
||||
"google": "Gemini",
|
||||
"vertex:google": "Gemini",
|
||||
"anthropic": "Claude",
|
||||
"openai": "OpenAI",
|
||||
"unknown": "Unidentified",
|
||||
}
|
||||
|
||||
_AGENT_SOURCE_PRIORITY: dict[str, int] = {
|
||||
"unknown": 0,
|
||||
"provider": 1,
|
||||
"model": 2,
|
||||
"stack": 3,
|
||||
"client": 4,
|
||||
}
|
||||
|
||||
|
||||
def _normalize_agent_key(raw: Any) -> str | None:
|
||||
if raw is None:
|
||||
return None
|
||||
value = str(raw).strip().lower()
|
||||
if not value:
|
||||
return None
|
||||
value = value.replace(" ", "-").replace("_", "-")
|
||||
if value.startswith("wrap-"):
|
||||
value = value.removeprefix("wrap-")
|
||||
if value in {"claude-cli", "claude-code", "claude-code-cli"}:
|
||||
return "claude-code"
|
||||
if value in {"codex-cli", "codex"}:
|
||||
return "codex"
|
||||
if value in {"github-copilot", "copilot"}:
|
||||
return "copilot"
|
||||
if value in {"google", "vertex-google", "vertex:google"}:
|
||||
return "gemini"
|
||||
return value
|
||||
|
||||
|
||||
def _agent_label(agent_key: str) -> str:
|
||||
if agent_key in _AGENT_LABELS:
|
||||
return _AGENT_LABELS[agent_key]
|
||||
return agent_key.replace("-", " ").replace("_", " ").title()
|
||||
|
||||
|
||||
def _classify_agent_from_log(entry: dict[str, Any]) -> tuple[str, str, str]:
|
||||
raw_tags = entry.get("tags")
|
||||
tags = raw_tags if isinstance(raw_tags, dict) else {}
|
||||
for source, candidate in (
|
||||
("client", tags.get("client")),
|
||||
("stack", tags.get("stack") or tags.get("headroom-stack")),
|
||||
):
|
||||
key = _normalize_agent_key(candidate)
|
||||
if key:
|
||||
return key, _agent_label(key), source
|
||||
|
||||
model = str(entry.get("model") or "").lower()
|
||||
if "codex" in model:
|
||||
return "codex", _agent_label("codex"), "model"
|
||||
if "claude" in model:
|
||||
return "claude-code", _agent_label("claude-code"), "model"
|
||||
if "gemini" in model:
|
||||
return "gemini", _agent_label("gemini"), "model"
|
||||
|
||||
key = _normalize_agent_key(entry.get("provider"))
|
||||
if key:
|
||||
return key, _agent_label(key), "provider"
|
||||
|
||||
return "unknown", _agent_label("unknown"), "unknown"
|
||||
|
||||
|
||||
def _build_agent_usage_summary(
|
||||
logs: list[dict[str, Any]],
|
||||
*,
|
||||
requests_by_provider: dict[str, int],
|
||||
requests_by_model: dict[str, int],
|
||||
global_before_tokens: int,
|
||||
global_after_tokens: int,
|
||||
global_tokens_saved: int,
|
||||
global_output_tokens: int,
|
||||
) -> dict[str, Any]:
|
||||
agents: dict[str, dict[str, Any]] = {}
|
||||
|
||||
def _agent_row(agent_key: str, label: str, source: str) -> dict[str, Any]:
|
||||
row = agents.setdefault(
|
||||
agent_key,
|
||||
{
|
||||
"agent": agent_key,
|
||||
"label": label,
|
||||
"source": source,
|
||||
"requests": 0,
|
||||
"before_tokens": 0,
|
||||
"after_tokens": 0,
|
||||
"output_tokens": 0,
|
||||
"tokens_saved": 0,
|
||||
"models": {},
|
||||
"providers": {},
|
||||
"has_exact_tokens": False,
|
||||
},
|
||||
)
|
||||
if _AGENT_SOURCE_PRIORITY.get(source, 0) > _AGENT_SOURCE_PRIORITY.get(
|
||||
str(row.get("source") or "unknown"), 0
|
||||
):
|
||||
row["source"] = source
|
||||
return row
|
||||
|
||||
for entry in logs:
|
||||
agent_key, label, source = _classify_agent_from_log(entry)
|
||||
row = _agent_row(agent_key, label, source)
|
||||
before = max(0, int(entry.get("input_tokens_original") or 0))
|
||||
after = max(0, int(entry.get("input_tokens_optimized") or 0))
|
||||
saved = max(0, int(entry.get("tokens_saved") or 0))
|
||||
output = max(0, int(entry.get("output_tokens") or 0))
|
||||
provider = str(entry.get("provider") or "unknown")
|
||||
model = str(entry.get("model") or "unknown")
|
||||
|
||||
row["requests"] += 1
|
||||
row["before_tokens"] += before
|
||||
row["after_tokens"] += after
|
||||
row["output_tokens"] += output
|
||||
row["tokens_saved"] += saved
|
||||
row["providers"][provider] = int(row["providers"].get(provider, 0)) + 1
|
||||
row["models"][model] = int(row["models"].get(model, 0)) + 1
|
||||
if before > 0 or after > 0 or saved > 0:
|
||||
row["has_exact_tokens"] = True
|
||||
|
||||
if not agents:
|
||||
inferred_model_counts: dict[str, int] = {}
|
||||
for model, count in requests_by_model.items():
|
||||
model_lower = str(model).lower()
|
||||
if "codex" in model_lower:
|
||||
key = "codex"
|
||||
elif "claude" in model_lower:
|
||||
key = "claude-code"
|
||||
elif "gemini" in model_lower:
|
||||
key = "gemini"
|
||||
else:
|
||||
continue
|
||||
inferred_model_counts[str(model)] = int(count)
|
||||
|
||||
provider_request_count = sum(max(0, int(count)) for count in requests_by_provider.values())
|
||||
inferred_request_count = sum(max(0, count) for count in inferred_model_counts.values())
|
||||
use_model_fallback = (
|
||||
inferred_request_count > 0 and inferred_request_count == provider_request_count
|
||||
)
|
||||
|
||||
if not use_model_fallback:
|
||||
for provider, count in requests_by_provider.items():
|
||||
key = _normalize_agent_key(provider) or "unknown"
|
||||
row = _agent_row(key, _agent_label(key), "provider")
|
||||
row["requests"] += int(count)
|
||||
row["providers"][provider] = int(row["providers"].get(provider, 0)) + int(count)
|
||||
for model, count in requests_by_model.items():
|
||||
model_lower = str(model).lower()
|
||||
if "codex" in model_lower:
|
||||
key = "codex"
|
||||
elif "claude" in model_lower:
|
||||
key = "claude-code"
|
||||
elif "gemini" in model_lower:
|
||||
key = "gemini"
|
||||
else:
|
||||
continue
|
||||
if not use_model_fallback:
|
||||
continue
|
||||
row = _agent_row(key, _agent_label(key), "model")
|
||||
row["requests"] += int(count)
|
||||
row["models"][str(model)] = int(row["models"].get(str(model), 0)) + int(count)
|
||||
|
||||
rows: list[dict[str, Any]] = []
|
||||
for row in agents.values():
|
||||
before = int(row["before_tokens"])
|
||||
saved = int(row["tokens_saved"])
|
||||
after = int(row["after_tokens"])
|
||||
if before == 0 and (after > 0 or saved > 0):
|
||||
before = after + saved
|
||||
savings_percent = round((saved / before) * 100.0, 2) if before else 0.0
|
||||
row["before_tokens"] = before
|
||||
row["savings_percent"] = savings_percent
|
||||
row["after_percent"] = round((after / before) * 100.0, 2) if before else 0.0
|
||||
row["share_of_saved_percent"] = (
|
||||
round((saved / global_tokens_saved) * 100.0, 2) if global_tokens_saved else 0.0
|
||||
)
|
||||
row["share_of_requests_percent"] = 0.0
|
||||
rows.append(row)
|
||||
|
||||
total_requests = sum(int(row["requests"]) for row in rows)
|
||||
for row in rows:
|
||||
row["share_of_requests_percent"] = (
|
||||
round((int(row["requests"]) / total_requests) * 100.0, 2) if total_requests else 0.0
|
||||
)
|
||||
|
||||
rows.sort(
|
||||
key=lambda row: (
|
||||
int(row.get("tokens_saved", 0)),
|
||||
int(row.get("before_tokens", 0)),
|
||||
int(row.get("requests", 0)),
|
||||
),
|
||||
reverse=True,
|
||||
)
|
||||
|
||||
return {
|
||||
"agents": rows,
|
||||
"totals": {
|
||||
"requests": total_requests,
|
||||
"before_tokens": global_before_tokens,
|
||||
"after_tokens": global_after_tokens,
|
||||
"output_tokens": global_output_tokens,
|
||||
"tokens_saved": global_tokens_saved,
|
||||
"savings_percent": (
|
||||
round((global_tokens_saved / global_before_tokens) * 100.0, 2)
|
||||
if global_before_tokens
|
||||
else 0.0
|
||||
),
|
||||
},
|
||||
"coverage": {
|
||||
"logged_requests": len(logs),
|
||||
"exact_token_rows": sum(1 for row in rows if row.get("has_exact_tokens")),
|
||||
"mode": "request_logs" if logs else "aggregate_fallback",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
logging.basicConfig(
|
||||
level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
|
||||
)
|
||||
|
|
@ -372,7 +606,6 @@ class HeadroomProxy(
|
|||
enable_code_aware=config.code_aware_enabled,
|
||||
tool_profiles=config.tool_profiles,
|
||||
read_lifecycle=ReadLifecycleConfig(enabled=config.read_lifecycle),
|
||||
ccr_inject_marker=config.ccr_inject_marker,
|
||||
smart_crusher_max_items_after_crush=cast(
|
||||
int | None,
|
||||
profile_kwargs.get("max_items_after_crush"),
|
||||
|
|
@ -381,6 +614,7 @@ class HeadroomProxy(
|
|||
bool,
|
||||
profile_kwargs.get("smart_crusher_with_compaction", True),
|
||||
),
|
||||
ccr_inject_marker=config.ccr_inject_marker,
|
||||
)
|
||||
if config.disable_kompress:
|
||||
router_config.enable_kompress = False
|
||||
|
|
@ -469,6 +703,9 @@ class HeadroomProxy(
|
|||
|
||||
# HTTP client
|
||||
self.http_client: httpx.AsyncClient | None = None
|
||||
# HTTP/1.1-only client for ChatGPT passthrough (Cloudflare challenges
|
||||
# our HTTP/2 fingerprint on its sensitive account endpoints).
|
||||
self.http_client_h1: httpx.AsyncClient | None = None
|
||||
|
||||
# Shared cold-start warmup registry (populated by startup()).
|
||||
# Holds typed slots with loaded / loading / null / error status for
|
||||
|
|
@ -901,19 +1138,26 @@ class HeadroomProxy(
|
|||
metadata={"port": self.config.port, "host": self.config.host},
|
||||
)
|
||||
_ca_bundle = find_ca_bundle()
|
||||
self.http_client = httpx.AsyncClient(
|
||||
timeout=httpx.Timeout(
|
||||
_client_kwargs: dict[str, Any] = {
|
||||
"timeout": httpx.Timeout(
|
||||
connect=self.config.connect_timeout_seconds,
|
||||
read=self.config.request_timeout_seconds,
|
||||
write=self.config.request_timeout_seconds,
|
||||
pool=self.config.connect_timeout_seconds,
|
||||
),
|
||||
limits=httpx.Limits(
|
||||
"limits": httpx.Limits(
|
||||
max_connections=self.config.max_connections,
|
||||
max_keepalive_connections=self.config.max_keepalive_connections,
|
||||
),
|
||||
http2=self.config.http2,
|
||||
verify=_ca_bundle if _ca_bundle is not None else True,
|
||||
"verify": _ca_bundle if _ca_bundle is not None else True,
|
||||
}
|
||||
self.http_client = httpx.AsyncClient(http2=self.config.http2, **_client_kwargs)
|
||||
# Reuse the primary client when HTTP/2 is already off; otherwise keep a
|
||||
# dedicated HTTP/1.1 client for ChatGPT passthrough.
|
||||
self.http_client_h1 = (
|
||||
self.http_client
|
||||
if not self.config.http2
|
||||
else httpx.AsyncClient(http2=False, **_client_kwargs)
|
||||
)
|
||||
logger.info("Headroom Proxy started")
|
||||
logger.info(f"Optimization: {'ENABLED' if self.config.optimize else 'DISABLED'}")
|
||||
|
|
@ -1135,6 +1379,9 @@ class HeadroomProxy(
|
|||
|
||||
async def shutdown(self):
|
||||
"""Cleanup async resources."""
|
||||
if self.http_client_h1 and self.http_client_h1 is not self.http_client:
|
||||
await self.http_client_h1.aclose()
|
||||
self.http_client_h1 = None
|
||||
if self.http_client:
|
||||
await self.http_client.aclose()
|
||||
self.http_client = None
|
||||
|
|
@ -2020,6 +2267,35 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
_stats_snapshot_lock = asyncio.Lock()
|
||||
_stats_snapshot: dict[str, Any] = {"expires_at": 0.0, "value": None}
|
||||
|
||||
RECENT_REQUEST_LOG_WINDOW = 100
|
||||
|
||||
def _build_recent_request_payload(limit: int = RECENT_REQUEST_LOG_WINDOW) -> dict[str, Any]:
|
||||
recent_request_logs = proxy.logger.get_recent(limit) if proxy.logger else []
|
||||
dashboard_recent_requests = [
|
||||
{
|
||||
"request_id": log.get("request_id"),
|
||||
"timestamp": log.get("timestamp"),
|
||||
"provider": log.get("provider"),
|
||||
"model": log.get("model"),
|
||||
"input_tokens_original": log.get("input_tokens_original"),
|
||||
"input_tokens_optimized": log.get("input_tokens_optimized"),
|
||||
"output_tokens": log.get("output_tokens"),
|
||||
"tokens_saved": log.get("tokens_saved"),
|
||||
"savings_percent": log.get("savings_percent"),
|
||||
"optimization_latency_ms": log.get("optimization_latency_ms"),
|
||||
"total_latency_ms": log.get("total_latency_ms"),
|
||||
"transforms_applied": log.get("transforms_applied", []),
|
||||
"waste_signals": log.get("waste_signals"),
|
||||
}
|
||||
for log in recent_request_logs
|
||||
if log.get("input_tokens_original") is not None
|
||||
and log.get("input_tokens_optimized") is not None
|
||||
][-10:]
|
||||
return {
|
||||
"request_logs": recent_request_logs[-10:],
|
||||
"recent_requests": dashboard_recent_requests,
|
||||
}
|
||||
|
||||
async def _build_stats_payload() -> dict[str, Any]:
|
||||
"""Build the full `/stats` response payload.
|
||||
|
||||
|
|
@ -2167,9 +2443,21 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
total_tokens_all_layers = all_layers_tokens_saved
|
||||
persistent_savings = m.savings_tracker.stats_preview()
|
||||
display_session = persistent_savings.get("display_session", {})
|
||||
recent_request_logs = proxy.logger.get_recent(10_000) if proxy.logger else []
|
||||
recent_request_payload = _build_recent_request_payload()
|
||||
agent_usage = _build_agent_usage_summary(
|
||||
recent_request_logs,
|
||||
requests_by_provider=dict(m.requests_by_provider),
|
||||
requests_by_model=dict(m.requests_by_model),
|
||||
global_before_tokens=proxy_total_before_compression,
|
||||
global_after_tokens=m.tokens_input_total,
|
||||
global_tokens_saved=proxy_compression_tokens,
|
||||
global_output_tokens=m.tokens_output_total,
|
||||
)
|
||||
|
||||
return {
|
||||
"summary": summary,
|
||||
"agent_usage": agent_usage,
|
||||
"savings": {
|
||||
"total_tokens": total_tokens_all_layers,
|
||||
"per_project": persistent_savings.get("projects", {}),
|
||||
|
|
@ -2434,11 +2722,43 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
"proxy_inbound": proxy.metrics.inbound_snapshot(),
|
||||
"cache": await proxy.cache.stats() if proxy.cache else None,
|
||||
"rate_limiter": await proxy.rate_limiter.stats() if proxy.rate_limiter else None,
|
||||
"recent_requests": proxy.logger.get_recent(10) if proxy.logger else [],
|
||||
**recent_request_payload,
|
||||
"log_full_messages": proxy.config.log_full_messages if proxy else False,
|
||||
**get_quota_registry().get_all_stats(),
|
||||
}
|
||||
|
||||
def _dashboard_config_payload() -> dict[str, Any]:
|
||||
profile_kwargs = proxy_pipeline_kwargs(config)
|
||||
target_ratio = profile_kwargs.get("target_ratio", config.target_ratio)
|
||||
target_savings_percent = None
|
||||
if isinstance(target_ratio, (int, float)):
|
||||
target_savings_percent = round(max(0.0, min(1.0, 1.0 - float(target_ratio))) * 100, 1)
|
||||
return {
|
||||
"savings_profile": config.savings_profile,
|
||||
"target_ratio": target_ratio,
|
||||
"target_savings_percent": target_savings_percent,
|
||||
"compress_user_messages": bool(
|
||||
profile_kwargs.get("compress_user_messages", config.compress_user_messages)
|
||||
),
|
||||
"compress_system_messages": bool(
|
||||
profile_kwargs.get("compress_system_messages", config.compress_system_messages)
|
||||
),
|
||||
"protect_recent": profile_kwargs.get("read_protection_window", config.protect_recent),
|
||||
"protect_analysis_context": config.protect_analysis_context,
|
||||
"min_tokens_to_crush": profile_kwargs.get(
|
||||
"min_tokens_to_compress", config.min_tokens_to_crush
|
||||
),
|
||||
"max_items_after_crush": profile_kwargs.get(
|
||||
"max_items_after_crush", config.max_items_after_crush
|
||||
),
|
||||
"smart_crusher_with_compaction": profile_kwargs.get(
|
||||
"smart_crusher_with_compaction",
|
||||
config.smart_crusher_with_compaction,
|
||||
),
|
||||
"force_kompress": bool(profile_kwargs.get("force_kompress", False)),
|
||||
"accuracy_guard": config.accuracy_guard,
|
||||
}
|
||||
|
||||
async def _get_cached_stats_payload() -> dict[str, Any]:
|
||||
"""Return a short-TTL cached `/stats` snapshot for dashboard polling."""
|
||||
now = time.monotonic()
|
||||
|
|
@ -2474,8 +2794,13 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
snapshot to avoid rebuilding the full payload on every UI poll.
|
||||
"""
|
||||
if cached:
|
||||
return await _get_cached_stats_payload()
|
||||
return await _build_stats_payload()
|
||||
payload = dict(await _get_cached_stats_payload())
|
||||
payload.update(_build_recent_request_payload())
|
||||
payload["config"] = _dashboard_config_payload()
|
||||
return payload
|
||||
payload = await _build_stats_payload()
|
||||
payload["config"] = _dashboard_config_payload()
|
||||
return payload
|
||||
|
||||
@app.post("/stats/reset", dependencies=[Depends(_require_loopback)])
|
||||
async def stats_reset():
|
||||
|
|
@ -3610,6 +3935,11 @@ if __name__ == "__main__":
|
|||
optimize=optimize,
|
||||
min_tokens_to_crush=_get_env_int("HEADROOM_MIN_TOKENS", args.min_tokens),
|
||||
max_items_after_crush=_get_env_int("HEADROOM_MAX_ITEMS", args.max_items),
|
||||
smart_crusher_with_compaction=(
|
||||
_get_env_bool("HEADROOM_SMART_CRUSHER_COMPACTION", False)
|
||||
if "HEADROOM_SMART_CRUSHER_COMPACTION" in os.environ
|
||||
else None
|
||||
),
|
||||
cache_enabled=cache_enabled,
|
||||
cache_ttl_seconds=_get_env_int("HEADROOM_CACHE_TTL", args.cache_ttl),
|
||||
rate_limit_enabled=rate_limit_enabled,
|
||||
|
|
@ -3632,6 +3962,28 @@ if __name__ == "__main__":
|
|||
mode=normalize_proxy_mode(_get_env_str("HEADROOM_MODE", PROXY_MODE_TOKEN)),
|
||||
compress_user_messages=args.compress_user_messages
|
||||
or _get_env_bool("HEADROOM_COMPRESS_USER_MESSAGES", False),
|
||||
savings_profile=os.environ.get("HEADROOM_SAVINGS_PROFILE") or None,
|
||||
target_ratio=(
|
||||
float(os.environ["HEADROOM_TARGET_RATIO"])
|
||||
if os.environ.get("HEADROOM_TARGET_RATIO")
|
||||
else None
|
||||
),
|
||||
compress_system_messages=(
|
||||
_get_env_bool("HEADROOM_COMPRESS_SYSTEM_MESSAGES", False)
|
||||
if "HEADROOM_COMPRESS_SYSTEM_MESSAGES" in os.environ
|
||||
else None
|
||||
),
|
||||
protect_recent=(
|
||||
int(os.environ["HEADROOM_PROTECT_RECENT"])
|
||||
if os.environ.get("HEADROOM_PROTECT_RECENT")
|
||||
else None
|
||||
),
|
||||
protect_analysis_context=(
|
||||
_get_env_bool("HEADROOM_PROTECT_ANALYSIS_CONTEXT", False)
|
||||
if "HEADROOM_PROTECT_ANALYSIS_CONTEXT" in os.environ
|
||||
else None
|
||||
),
|
||||
accuracy_guard=os.environ.get("HEADROOM_ACCURACY_GUARD") or None,
|
||||
)
|
||||
|
||||
# Get worker and concurrency settings
|
||||
|
|
|
|||
|
|
@ -1,10 +1,19 @@
|
|||
"""Passive tracking of OpenAI Codex rate-limit window data from response headers.
|
||||
"""Tracking of OpenAI Codex rate-limit window data.
|
||||
|
||||
Codex (OpenAI) embeds rate-limit data directly in API response headers
|
||||
(``x-codex-primary-used-percent``, ``x-codex-primary-window-minutes``, etc.)
|
||||
rather than exposing a dedicated usage endpoint. This module captures those
|
||||
headers from responses that headroom proxies and makes them available in
|
||||
``/stats`` and the dashboard.
|
||||
Historically Codex embedded rate-limit data in API *response headers*
|
||||
(``x-codex-primary-used-percent`` etc.) and headroom captured those headers
|
||||
from proxied responses (:meth:`CodexRateLimitState.update_from_headers`).
|
||||
Current Codex (codex_exec / TUI on the ChatGPT WebSocket transport) no longer
|
||||
emits those headers on the ``/responses`` handshake or stream -- the window is
|
||||
served from a dedicated endpoint instead:
|
||||
|
||||
GET https://chatgpt.com/backend-api/wham/usage (ChatGPT OAuth/session)
|
||||
|
||||
So this module also exposes :func:`maybe_schedule_usage_poll`, a throttled
|
||||
fire-and-forget GET against that endpoint using the client's own bearer token
|
||||
and ``ChatGPT-Account-Id``. The header-capture path is kept intact: if OpenAI
|
||||
ever returns ``x-codex-*`` again it still works, and API-key requests (no
|
||||
account id) simply never trigger a poll.
|
||||
|
||||
Header schema (parsed by codex-rs ``rate_limits.rs``):
|
||||
x-codex-primary-used-percent float 0-100
|
||||
|
|
@ -18,16 +27,47 @@ Header schema (parsed by codex-rs ``rate_limits.rs``):
|
|||
x-codex-credits-balance str e.g. "$5.00"
|
||||
x-codex-promo-message str server announcement
|
||||
x-codex-limit-name str e.g. "gpt-5.2-codex-sonic"
|
||||
|
||||
``GET /wham/usage`` JSON schema (mapped by :func:`parse_codex_usage_payload`):
|
||||
plan_type str
|
||||
rate_limit.primary_window.used_percent float 0-100
|
||||
rate_limit.primary_window.limit_window_seconds int
|
||||
rate_limit.primary_window.reset_at int Unix timestamp (seconds)
|
||||
rate_limit.secondary_window.* same shape (optional)
|
||||
credits.has_credits / unlimited / balance bool / bool / str
|
||||
rate_limit_reached_type str | null
|
||||
promo obj | str | null
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from threading import Lock
|
||||
|
||||
import httpx
|
||||
|
||||
from headroom.subscription.base import QuotaTracker
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Dedicated Codex usage endpoint for ChatGPT OAuth/session auth. Overridable
|
||||
# for tests / self-hosted gateways via env.
|
||||
CODEX_USAGE_URL = (
|
||||
os.environ.get("HEADROOM_CODEX_USAGE_URL", "https://chatgpt.com/backend-api/wham/usage").strip()
|
||||
or "https://chatgpt.com/backend-api/wham/usage"
|
||||
)
|
||||
|
||||
# Minimum seconds between live usage polls. Codex turns can arrive in bursts;
|
||||
# one GET per minute is plenty to keep the gauge fresh without hammering.
|
||||
USAGE_POLL_MIN_INTERVAL_S = 60.0
|
||||
|
||||
# Bound the usage GET so a slow upstream never wedges the fire-and-forget task.
|
||||
_USAGE_POLL_TIMEOUT_S = 10.0
|
||||
|
||||
|
||||
@dataclass
|
||||
class CodexRateLimitWindow:
|
||||
|
|
@ -192,6 +232,94 @@ def parse_codex_rate_limits(headers: dict[str, str]) -> CodexRateLimitSnapshot |
|
|||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Usage-endpoint (GET /wham/usage) JSON parsing
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _window_from_usage_json(win: object) -> CodexRateLimitWindow | None:
|
||||
"""Map one ``rate_limit.{primary,secondary}_window`` object to a window."""
|
||||
if not isinstance(win, dict):
|
||||
return None
|
||||
used = win.get("used_percent")
|
||||
try:
|
||||
used_f = float(used) # type: ignore[arg-type]
|
||||
except (ValueError, TypeError):
|
||||
return None
|
||||
if used_f != used_f: # NaN guard
|
||||
return None
|
||||
|
||||
window_minutes: int | None = None
|
||||
secs = win.get("limit_window_seconds")
|
||||
if isinstance(secs, (int, float)) and secs > 0:
|
||||
# Round up, matching codex-rs window_minutes_from_seconds.
|
||||
window_minutes = (int(secs) + 59) // 60
|
||||
|
||||
resets_at = win.get("reset_at")
|
||||
resets_at = int(resets_at) if isinstance(resets_at, (int, float)) else None
|
||||
|
||||
return CodexRateLimitWindow(
|
||||
used_percent=used_f,
|
||||
window_minutes=window_minutes,
|
||||
resets_at=resets_at,
|
||||
)
|
||||
|
||||
|
||||
def parse_codex_usage_payload(payload: object) -> CodexRateLimitSnapshot | None:
|
||||
"""Parse a snapshot from a ``GET /wham/usage`` JSON body.
|
||||
|
||||
Returns ``None`` when the body carries no usable rate-limit data.
|
||||
"""
|
||||
if not isinstance(payload, dict):
|
||||
return None
|
||||
|
||||
rate_limit = payload.get("rate_limit")
|
||||
rate_limit = rate_limit if isinstance(rate_limit, dict) else {}
|
||||
primary = _window_from_usage_json(rate_limit.get("primary_window"))
|
||||
secondary = _window_from_usage_json(rate_limit.get("secondary_window"))
|
||||
|
||||
credits: CodexCreditsSnapshot | None = None
|
||||
cred = payload.get("credits")
|
||||
if isinstance(cred, dict) and cred.get("has_credits") is not None:
|
||||
has = bool(cred.get("has_credits"))
|
||||
raw_balance = cred.get("balance")
|
||||
credits = CodexCreditsSnapshot(
|
||||
has_credits=has,
|
||||
unlimited=bool(cred.get("unlimited")),
|
||||
# Only surface a balance when the account actually has credits;
|
||||
# a "0" balance on a no-credits plan is noise to the gauge.
|
||||
balance=(str(raw_balance) if has and raw_balance not in (None, "") else None),
|
||||
)
|
||||
|
||||
promo = payload.get("promo")
|
||||
if isinstance(promo, dict):
|
||||
promo_message = promo.get("message")
|
||||
elif isinstance(promo, str):
|
||||
promo_message = promo
|
||||
else:
|
||||
promo_message = None
|
||||
promo_message = (promo_message or "").strip() or None
|
||||
|
||||
raw_limit_name = payload.get("rate_limit_reached_type")
|
||||
limit_name = (
|
||||
raw_limit_name.strip()
|
||||
if isinstance(raw_limit_name, str) and raw_limit_name.strip()
|
||||
else None
|
||||
)
|
||||
|
||||
if primary is None and secondary is None and credits is None and promo_message is None:
|
||||
return None
|
||||
|
||||
return CodexRateLimitSnapshot(
|
||||
limit_id="codex",
|
||||
limit_name=limit_name,
|
||||
primary=primary,
|
||||
secondary=secondary,
|
||||
credits=credits,
|
||||
promo_message=promo_message,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Singleton state store
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -214,6 +342,8 @@ class CodexRateLimitState(QuotaTracker):
|
|||
def __init__(self) -> None:
|
||||
self._lock = Lock()
|
||||
self._latest: CodexRateLimitSnapshot | None = None
|
||||
self._last_poll_monotonic: float = 0.0
|
||||
self._poll_inflight: bool = False
|
||||
|
||||
def update_from_headers(self, headers: dict[str, str]) -> None:
|
||||
"""Update state from a response header dict (no-op if no Codex headers)."""
|
||||
|
|
@ -223,6 +353,38 @@ class CodexRateLimitState(QuotaTracker):
|
|||
with self._lock:
|
||||
self._latest = snapshot
|
||||
|
||||
def update_from_usage_payload(self, payload: object) -> bool:
|
||||
"""Update state from a ``GET /wham/usage`` JSON body.
|
||||
|
||||
Returns ``True`` when a snapshot was stored.
|
||||
"""
|
||||
snapshot = parse_codex_usage_payload(payload)
|
||||
if snapshot is None:
|
||||
return False
|
||||
with self._lock:
|
||||
self._latest = snapshot
|
||||
return True
|
||||
|
||||
def _try_begin_poll(self, min_interval_s: float) -> bool:
|
||||
"""Atomically claim a usage-poll slot.
|
||||
|
||||
Returns ``False`` when a poll is already in flight or one ran within
|
||||
``min_interval_s``. On ``True`` the caller MUST call :meth:`_end_poll`.
|
||||
"""
|
||||
now = time.monotonic()
|
||||
with self._lock:
|
||||
if self._poll_inflight:
|
||||
return False
|
||||
if (now - self._last_poll_monotonic) < min_interval_s:
|
||||
return False
|
||||
self._poll_inflight = True
|
||||
self._last_poll_monotonic = now
|
||||
return True
|
||||
|
||||
def _end_poll(self) -> None:
|
||||
with self._lock:
|
||||
self._poll_inflight = False
|
||||
|
||||
@property
|
||||
def latest(self) -> CodexRateLimitSnapshot | None:
|
||||
with self._lock:
|
||||
|
|
@ -245,3 +407,80 @@ def get_codex_rate_limit_state() -> CodexRateLimitState:
|
|||
if _state is None:
|
||||
_state = CodexRateLimitState()
|
||||
return _state
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Live usage poll (GET /wham/usage)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _build_usage_headers(request_headers: dict[str, str]) -> dict[str, str] | None:
|
||||
"""Build outbound /wham/usage headers from a client's request headers.
|
||||
|
||||
Returns ``None`` unless the request carries a bearer token *and* a
|
||||
``ChatGPT-Account-Id`` -- the latter scopes the poll to ChatGPT OAuth
|
||||
sessions (Codex), so API-key and non-Codex OAuth traffic never triggers it.
|
||||
"""
|
||||
lower = {str(k).lower(): v for k, v in request_headers.items()}
|
||||
auth = str(lower.get("authorization", ""))
|
||||
if not auth.startswith("Bearer ") or not auth[len("Bearer ") :].strip():
|
||||
return None
|
||||
account_id = lower.get("chatgpt-account-id")
|
||||
if not account_id:
|
||||
return None
|
||||
|
||||
headers = {
|
||||
"Authorization": auth,
|
||||
"ChatGPT-Account-Id": str(account_id),
|
||||
"Accept": "application/json",
|
||||
}
|
||||
# Mirror the client's own UA/originator so the request looks like Codex.
|
||||
for src, dst in (("user-agent", "User-Agent"), ("originator", "originator")):
|
||||
val = lower.get(src)
|
||||
if val:
|
||||
headers[dst] = str(val)
|
||||
return headers
|
||||
|
||||
|
||||
async def _fetch_and_store_usage(url: str, headers: dict[str, str]) -> None:
|
||||
state = get_codex_rate_limit_state()
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=_USAGE_POLL_TIMEOUT_S) as client:
|
||||
resp = await client.get(url, headers=headers)
|
||||
if resp.status_code == 200:
|
||||
if state.update_from_usage_payload(resp.json()):
|
||||
logger.debug("codex usage poll: refreshed rate-limit window")
|
||||
else:
|
||||
logger.debug("codex usage poll: 200 but no usable rate-limit data")
|
||||
else:
|
||||
logger.debug("codex usage poll: HTTP %s", resp.status_code)
|
||||
except Exception as exc: # pragma: no cover - network/JSON defensive
|
||||
logger.debug("codex usage poll failed: %s", exc)
|
||||
finally:
|
||||
state._end_poll()
|
||||
|
||||
|
||||
def maybe_schedule_usage_poll(
|
||||
request_headers: dict[str, str],
|
||||
*,
|
||||
url: str = CODEX_USAGE_URL,
|
||||
min_interval_s: float = USAGE_POLL_MIN_INTERVAL_S,
|
||||
) -> bool:
|
||||
"""Fire-and-forget a throttled ``GET /wham/usage`` to refresh the window.
|
||||
|
||||
Safe to call on every Codex request: scoped to ChatGPT-session traffic via
|
||||
:func:`_build_usage_headers` and internally throttled to at most one live
|
||||
poll per ``min_interval_s``. Returns ``True`` when a poll was scheduled.
|
||||
"""
|
||||
headers = _build_usage_headers(request_headers)
|
||||
if headers is None:
|
||||
return False
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
except RuntimeError:
|
||||
return False
|
||||
state = get_codex_rate_limit_state()
|
||||
if not state._try_begin_poll(min_interval_s):
|
||||
return False
|
||||
loop.create_task(_fetch_and_store_usage(url, headers))
|
||||
return True
|
||||
|
|
|
|||
|
|
@ -137,7 +137,12 @@ class CompressionPolicy:
|
|||
Mirrors ``CompressionPolicy::net_mutation_gain`` in the Rust
|
||||
crate (source of truth — see its docstring for the derivation)::
|
||||
|
||||
gain = dT * (w + r*(R - 1)) - P_alive * (w - r) * S
|
||||
gain = dT * (w + r*(R - 1)) - P_alive * (w - r) * (S + dT)
|
||||
|
||||
The warm-case penalty covers ``S + dT``: with a live cache the
|
||||
``dT`` tokens are already cache-written, so keeping them costs
|
||||
only reads — a mutation avoids at most ``dT*r*R``, not a fresh
|
||||
write.
|
||||
|
||||
Inputs are clamped: ``delta_t``/``suffix_tokens`` to ``>= 0``
|
||||
(the Rust signature takes ``u32``), ``expected_reads`` to
|
||||
|
|
@ -152,7 +157,7 @@ class CompressionPolicy:
|
|||
# f32::max in the Rust source of truth — guard explicitly.
|
||||
reads = 0.0 if math.isnan(expected_reads) else max(expected_reads, 0.0)
|
||||
alive = 1.0 if math.isnan(p_alive) else min(max(p_alive, 0.0), 1.0)
|
||||
return float(dt) * (w + r * (reads - 1.0)) - alive * (w - r) * float(suffix)
|
||||
return float(dt) * (w + r * (reads - 1.0)) - alive * (w - r) * float(suffix + dt)
|
||||
|
||||
def should_mutate_deep(
|
||||
self,
|
||||
|
|
@ -169,7 +174,10 @@ class CompressionPolicy:
|
|||
"""Remaining-read count at which a warm-cache (``p_alive=1``)
|
||||
mutation breaks even::
|
||||
|
||||
R = ((w - r) / r) * (S/dT - 1) ~= 11.5 * S/dT for S >> dT
|
||||
R = ((w - r) / r) * (S/dT) = 11.5 * S/dT (Anthropic 5-min)
|
||||
|
||||
With the corrected penalty this reproduces the #856 anchors
|
||||
exactly: 2K/50K -> 287.5, 50K/10K -> 2.3.
|
||||
|
||||
Returns 0 when ``delta_t`` is ``<= 0`` (no savings — callers
|
||||
gate on ``delta_t > 0``; the Rust signature takes ``u32``).
|
||||
|
|
@ -179,7 +187,7 @@ class CompressionPolicy:
|
|||
return 0.0
|
||||
w = CACHE_WRITE_MULTIPLIER
|
||||
r = CACHE_READ_MULTIPLIER
|
||||
return ((w - r) / r) * (float(max(0, suffix_tokens)) / float(delta_t) - 1.0)
|
||||
return ((w - r) / r) * (float(max(0, suffix_tokens)) / float(delta_t))
|
||||
|
||||
|
||||
def policy_for_mode(mode: AuthMode) -> CompressionPolicy:
|
||||
|
|
|
|||
|
|
@ -212,11 +212,14 @@ class TransformPipeline:
|
|||
- output_buffer: Output buffer override.
|
||||
- tool_profiles: Per-tool compression profiles.
|
||||
- request_id: Optional request ID for diff artifact.
|
||||
- waste_messages: Optional richer conversion of the same request
|
||||
used for waste-signal detection only (never transformed).
|
||||
|
||||
Returns:
|
||||
Combined TransformResult.
|
||||
"""
|
||||
record_metrics = kwargs.pop("record_metrics", True)
|
||||
waste_messages = kwargs.pop("waste_messages", None)
|
||||
tokenizer = self._get_tokenizer(model)
|
||||
provider_name = self._provider_name()
|
||||
|
||||
|
|
@ -430,13 +433,17 @@ class TransformPipeline:
|
|||
transforms=transform_diffs,
|
||||
)
|
||||
|
||||
# Detect waste signals in original messages (only when significant compression)
|
||||
# Detect waste signals in original messages (only when significant
|
||||
# compression). Handlers whose wire format carries tool output the
|
||||
# message conversion drops (e.g. Gemini functionResponse parts, #819)
|
||||
# pass a richer waste_messages list that is parsed instead — it is
|
||||
# telemetry-only and never transformed.
|
||||
waste_signals: WasteSignals | None = None
|
||||
if tokens_before > tokens_after and (tokens_before - tokens_after) > 100:
|
||||
try:
|
||||
from ..parser import parse_messages
|
||||
|
||||
_, _, waste_signals = parse_messages(messages, tokenizer)
|
||||
_, _, waste_signals = parse_messages(waste_messages or messages, tokenizer)
|
||||
if waste_signals.total() == 0:
|
||||
waste_signals = None
|
||||
except Exception:
|
||||
|
|
|
|||
|
|
@ -1,7 +1,8 @@
|
|||
#!/usr/bin/env bash
|
||||
# Install git hooks for the Headroom repo:
|
||||
# 1. pre-commit — repo pre-commit checks (ruff, mypy, sync-plugin-versions)
|
||||
# 2. pre-push — full ci-precheck (cargo fmt/clippy/test + python suite)
|
||||
# 2. commit-msg — conventional-commit enforcement via commitlint
|
||||
# 3. pre-push — full ci-precheck (cargo fmt/clippy/test + python suite)
|
||||
#
|
||||
# Why pre-push was added: the 2026-04-27 push hit five CI failures that could
|
||||
# all have been caught locally — cargo fmt drift, an x86_64-apple-darwin wheel
|
||||
|
|
@ -24,6 +25,11 @@ if [[ ! -d .git/hooks ]]; then
|
|||
exit 1
|
||||
fi
|
||||
|
||||
if ! command -v npx &>/dev/null; then
|
||||
echo "error: npx not found — install Node 18+ before installing Headroom's git hooks." >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
HOOK_PATH=".git/hooks/pre-push"
|
||||
|
||||
cat > "$HOOK_PATH" <<'HOOK_EOF'
|
||||
|
|
@ -89,7 +95,9 @@ fi
|
|||
|
||||
if [[ -n "$PRE_COMMIT_BIN" ]]; then
|
||||
"$PRE_COMMIT_BIN" install
|
||||
"$PRE_COMMIT_BIN" install --hook-type commit-msg
|
||||
echo "✅ installed: .git/hooks/pre-commit (repo pre-commit checks via pre-commit)"
|
||||
echo "✅ installed: .git/hooks/commit-msg (conventional commit enforcement via commitlint)"
|
||||
else
|
||||
echo "error: pre-commit not found — run 'pip install -e .[dev]' first, then re-run this script." >&2
|
||||
exit 1
|
||||
|
|
|
|||
297
scripts/pr-governance.py
Normal file
297
scripts/pr-governance.py
Normal file
|
|
@ -0,0 +1,297 @@
|
|||
#!/usr/bin/env python3
|
||||
"""Validate Headroom PR template compliance for GitHub Actions."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
COMMENT_MARKER = "<!-- headroom-pr-governance -->"
|
||||
READY_LABEL = "status: ready for review"
|
||||
AUTHOR_ACTION_LABEL = "status: needs author action"
|
||||
|
||||
REQUIRED_SECTIONS = (
|
||||
"Description",
|
||||
"Type of Change",
|
||||
"Changes Made",
|
||||
"Testing",
|
||||
"Real Behavior Proof",
|
||||
"Review Readiness",
|
||||
)
|
||||
PROOF_FIELDS = (
|
||||
"Environment",
|
||||
"Exact command / steps",
|
||||
"Observed result",
|
||||
"Not tested",
|
||||
)
|
||||
|
||||
SECTION_RE = re.compile(r"^##\s+(.+?)\s*$", re.MULTILINE)
|
||||
CHECKBOX_RE = re.compile(r"^- \[(?P<checked>[ xX])\] (?P<label>.+)$", re.MULTILINE)
|
||||
HTML_COMMENT_RE = re.compile(r"<!--.*?-->", re.DOTALL)
|
||||
CODE_BLOCK_RE = re.compile(r"```(?:[\w.+-]+)?\n(?P<content>.*?)```", re.DOTALL)
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class GovernanceReport:
|
||||
"""Serializable PR governance result."""
|
||||
|
||||
comment_marker: str
|
||||
valid: bool
|
||||
is_draft: bool
|
||||
is_bot_pr: bool
|
||||
ready_for_review: bool
|
||||
needs_author_action: bool
|
||||
problems: list[str] = field(default_factory=list)
|
||||
labels_to_add: list[str] = field(default_factory=list)
|
||||
labels_to_remove: list[str] = field(default_factory=list)
|
||||
comment_markdown: str = ""
|
||||
summary_markdown: str = ""
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return asdict(self)
|
||||
|
||||
|
||||
def load_event(path: Path) -> dict[str, Any]:
|
||||
return json.loads(path.read_text(encoding="utf-8"))
|
||||
|
||||
|
||||
def extract_sections(body: str) -> dict[str, str]:
|
||||
matches = list(SECTION_RE.finditer(body))
|
||||
sections: dict[str, str] = {}
|
||||
for index, match in enumerate(matches):
|
||||
start = match.end()
|
||||
end = matches[index + 1].start() if index + 1 < len(matches) else len(body)
|
||||
sections[match.group(1).strip()] = body[start:end].strip()
|
||||
return sections
|
||||
|
||||
|
||||
def strip_html_comments(text: str) -> str:
|
||||
return HTML_COMMENT_RE.sub("", text).strip()
|
||||
|
||||
|
||||
def non_empty_lines(text: str) -> list[str]:
|
||||
return [line.strip() for line in strip_html_comments(text).splitlines() if line.strip()]
|
||||
|
||||
|
||||
def checked_items(section: str) -> list[str]:
|
||||
return [
|
||||
match.group("label").strip()
|
||||
for match in CHECKBOX_RE.finditer(section)
|
||||
if match.group("checked").lower() == "x"
|
||||
]
|
||||
|
||||
|
||||
def has_descriptive_text(section: str) -> bool:
|
||||
ignored_prefixes = ("closes #", "fixes #", "resolves #", "related to #")
|
||||
for line in non_empty_lines(section):
|
||||
lowered = line.lower()
|
||||
if line.startswith("#"):
|
||||
continue
|
||||
if lowered.startswith(ignored_prefixes):
|
||||
continue
|
||||
if len(line) >= 10:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def has_non_placeholder_bullets(section: str) -> bool:
|
||||
placeholders = {"change 1", "change 2", "change 3"}
|
||||
for line in non_empty_lines(section):
|
||||
if not line.startswith("- "):
|
||||
continue
|
||||
bullet = line[2:].strip().lower()
|
||||
if bullet and bullet not in placeholders:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def has_test_output(section: str) -> bool:
|
||||
for match in CODE_BLOCK_RE.finditer(section):
|
||||
content = strip_html_comments(match.group("content")).strip()
|
||||
if not content:
|
||||
continue
|
||||
if "paste relevant command output or artifact links here" in content.lower():
|
||||
continue
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def proof_field_values(section: str) -> dict[str, str]:
|
||||
values: dict[str, str] = {}
|
||||
for line in non_empty_lines(section):
|
||||
if not line.startswith("- ") or ":" not in line:
|
||||
continue
|
||||
label, value = line[2:].split(":", 1)
|
||||
values[label.strip()] = value.strip()
|
||||
return values
|
||||
|
||||
|
||||
def normalize_checkbox_map(items: list[str]) -> set[str]:
|
||||
return {item.lower() for item in items}
|
||||
|
||||
|
||||
def validate_pull_request(event: dict[str, Any]) -> GovernanceReport:
|
||||
pull_request = event["pull_request"]
|
||||
author = pull_request["user"]["login"]
|
||||
is_draft = bool(pull_request.get("draft", False))
|
||||
is_bot_pr = author.endswith("[bot]")
|
||||
body = pull_request.get("body") or ""
|
||||
|
||||
if is_bot_pr:
|
||||
summary = "### PR governance\n\nBot-authored PR detected; template enforcement is skipped."
|
||||
return GovernanceReport(
|
||||
comment_marker=COMMENT_MARKER,
|
||||
valid=True,
|
||||
is_draft=is_draft,
|
||||
is_bot_pr=True,
|
||||
ready_for_review=False,
|
||||
needs_author_action=False,
|
||||
comment_markdown=summary,
|
||||
summary_markdown=summary,
|
||||
)
|
||||
|
||||
sections = extract_sections(body)
|
||||
problems: list[str] = []
|
||||
|
||||
for section_name in REQUIRED_SECTIONS:
|
||||
if section_name not in sections:
|
||||
problems.append(f"Missing required section `{section_name}`.")
|
||||
|
||||
description = sections.get("Description", "")
|
||||
if description and not has_descriptive_text(description):
|
||||
problems.append("Fill in `Description` with a real summary of the change.")
|
||||
|
||||
changes_made = sections.get("Changes Made", "")
|
||||
if changes_made and not has_non_placeholder_bullets(changes_made):
|
||||
problems.append(
|
||||
"Replace the placeholder bullets in `Changes Made` with the actual changes."
|
||||
)
|
||||
|
||||
type_of_change_checked = checked_items(sections.get("Type of Change", ""))
|
||||
if sections.get("Type of Change") and not type_of_change_checked:
|
||||
problems.append("Check at least one box in `Type of Change`.")
|
||||
|
||||
testing_section = sections.get("Testing", "")
|
||||
testing_checked = checked_items(testing_section)
|
||||
if testing_section and not testing_checked:
|
||||
problems.append("Check at least one verification item in `Testing`.")
|
||||
if testing_section and not has_test_output(testing_section):
|
||||
problems.append("Paste real command output or artifact links in `Testing` → `Test Output`.")
|
||||
|
||||
proof_section = sections.get("Real Behavior Proof", "")
|
||||
proof_values = proof_field_values(proof_section)
|
||||
for field_name in PROOF_FIELDS:
|
||||
if proof_section and not proof_values.get(field_name):
|
||||
problems.append(f"Fill in `Real Behavior Proof` → `{field_name}`.")
|
||||
|
||||
readiness_checked = normalize_checkbox_map(checked_items(sections.get("Review Readiness", "")))
|
||||
has_self_review = "i have performed a self-review" in readiness_checked
|
||||
has_ready_checkbox = "this pr is ready for human review" in readiness_checked
|
||||
if not is_draft:
|
||||
if not has_self_review:
|
||||
problems.append(
|
||||
"Check `I have performed a self-review` before requesting human review."
|
||||
)
|
||||
if not has_ready_checkbox:
|
||||
problems.append(
|
||||
"Check `This PR is ready for human review` or convert the PR back to draft."
|
||||
)
|
||||
|
||||
valid = not problems
|
||||
ready_for_review = valid and not is_draft and has_ready_checkbox and has_self_review
|
||||
needs_author_action = not valid
|
||||
|
||||
if valid and ready_for_review:
|
||||
status_lines = [
|
||||
"### PR governance",
|
||||
"",
|
||||
"This PR follows the template and is marked ready for human review.",
|
||||
]
|
||||
elif valid:
|
||||
status_lines = [
|
||||
"### PR governance",
|
||||
"",
|
||||
"This draft PR follows the template so far. Keep it in draft until it is ready for human review.",
|
||||
]
|
||||
else:
|
||||
status_lines = [
|
||||
"### PR governance",
|
||||
"",
|
||||
"This PR does not yet satisfy the required template fields:",
|
||||
"",
|
||||
*[f"- {problem}" for problem in problems],
|
||||
"",
|
||||
"Please update the PR body, or move the PR back to draft while it is still in progress.",
|
||||
]
|
||||
|
||||
labels_to_add: list[str] = []
|
||||
labels_to_remove: list[str] = []
|
||||
if needs_author_action:
|
||||
labels_to_add.append(AUTHOR_ACTION_LABEL)
|
||||
labels_to_remove.append(READY_LABEL)
|
||||
else:
|
||||
labels_to_remove.append(AUTHOR_ACTION_LABEL)
|
||||
if ready_for_review:
|
||||
labels_to_add.append(READY_LABEL)
|
||||
else:
|
||||
labels_to_remove.append(READY_LABEL)
|
||||
|
||||
comment_markdown = "\n".join(status_lines)
|
||||
return GovernanceReport(
|
||||
comment_marker=COMMENT_MARKER,
|
||||
valid=valid,
|
||||
is_draft=is_draft,
|
||||
is_bot_pr=False,
|
||||
ready_for_review=ready_for_review,
|
||||
needs_author_action=needs_author_action,
|
||||
problems=problems,
|
||||
labels_to_add=labels_to_add,
|
||||
labels_to_remove=labels_to_remove,
|
||||
comment_markdown=comment_markdown,
|
||||
summary_markdown=comment_markdown,
|
||||
)
|
||||
|
||||
|
||||
def emit_outputs(report: GovernanceReport) -> None:
|
||||
output_path = os.environ.get("GITHUB_OUTPUT")
|
||||
lines = [
|
||||
f"valid={str(report.valid).lower()}",
|
||||
f"ready_for_review={str(report.ready_for_review).lower()}",
|
||||
f"needs_author_action={str(report.needs_author_action).lower()}",
|
||||
f"is_bot_pr={str(report.is_bot_pr).lower()}",
|
||||
]
|
||||
if not output_path:
|
||||
for line in lines:
|
||||
print(line)
|
||||
return
|
||||
|
||||
with Path(output_path).open("a", encoding="utf-8") as output_file:
|
||||
for line in lines:
|
||||
output_file.write(f"{line}\n")
|
||||
|
||||
|
||||
def parse_args(argv: list[str]) -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument(
|
||||
"--event", type=Path, required=True, help="Path to the GitHub event payload JSON."
|
||||
)
|
||||
parser.add_argument("--report", type=Path, required=True, help="Path to write the JSON report.")
|
||||
return parser.parse_args(argv)
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
args = parse_args(argv or sys.argv[1:])
|
||||
report = validate_pull_request(load_event(args.event))
|
||||
args.report.write_text(json.dumps(report.to_dict(), indent=2), encoding="utf-8")
|
||||
emit_outputs(report)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
163
scripts/tests/test_pr_governance.py
Normal file
163
scripts/tests/test_pr_governance.py
Normal file
|
|
@ -0,0 +1,163 @@
|
|||
"""Tests for pr-governance.py."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def _load_module():
|
||||
script = Path(__file__).parent.parent / "pr-governance.py"
|
||||
spec = importlib.util.spec_from_file_location("pr_governance", script)
|
||||
assert spec is not None
|
||||
assert spec.loader is not None
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[spec.name] = module
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
def _event(body: str, *, draft: bool = False, login: str = "octocat") -> dict[str, object]:
|
||||
return {
|
||||
"pull_request": {
|
||||
"number": 42,
|
||||
"draft": draft,
|
||||
"body": body,
|
||||
"user": {"login": login},
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
VALID_BODY = """## Description
|
||||
|
||||
Add a required PR-governance gate for template validation and review readiness.
|
||||
|
||||
Closes #123
|
||||
|
||||
## Type of Change
|
||||
|
||||
- [x] New feature (non-breaking change that adds functionality)
|
||||
- [ ] Documentation update
|
||||
|
||||
## Changes Made
|
||||
|
||||
- Added a workflow-backed PR template validator.
|
||||
- Added local commit message linting in the commit-msg hook.
|
||||
|
||||
## Testing
|
||||
|
||||
- [x] Unit tests pass (`pytest`)
|
||||
- [x] Manual testing performed
|
||||
|
||||
### Test Output
|
||||
|
||||
```text
|
||||
pytest scripts/tests/test_pr_governance.py -q
|
||||
```
|
||||
|
||||
## Real Behavior Proof
|
||||
|
||||
- Environment: Ubuntu runner, Python 3.12
|
||||
- Exact command / steps: Open a PR, remove the ready checkbox, re-run the workflow.
|
||||
- Observed result: The governance check fails and the PR gets a needs-author-action label.
|
||||
- Not tested: Automatic Copilot review rulesets in repository settings.
|
||||
|
||||
## Review Readiness
|
||||
|
||||
- [x] I have performed a self-review
|
||||
- [x] This PR is ready for human review
|
||||
|
||||
## Additional Notes
|
||||
|
||||
- Maintainers can optionally enable Copilot code review from repository rulesets.
|
||||
"""
|
||||
|
||||
|
||||
def test_validate_pull_request_marks_ready_pr_valid() -> None:
|
||||
module = _load_module()
|
||||
|
||||
report = module.validate_pull_request(_event(VALID_BODY))
|
||||
|
||||
assert report.valid is True
|
||||
assert report.ready_for_review is True
|
||||
assert report.needs_author_action is False
|
||||
assert report.problems == []
|
||||
assert report.labels_to_add == [module.READY_LABEL]
|
||||
assert module.AUTHOR_ACTION_LABEL in report.labels_to_remove
|
||||
|
||||
|
||||
def test_validate_pull_request_allows_draft_without_ready_checkboxes() -> None:
|
||||
module = _load_module()
|
||||
body = VALID_BODY.replace(
|
||||
"- [x] I have performed a self-review", "- [ ] I have performed a self-review"
|
||||
)
|
||||
body = body.replace(
|
||||
"- [x] This PR is ready for human review",
|
||||
"- [ ] This PR is ready for human review",
|
||||
)
|
||||
|
||||
report = module.validate_pull_request(_event(body, draft=True))
|
||||
|
||||
assert report.valid is True
|
||||
assert report.ready_for_review is False
|
||||
assert report.needs_author_action is False
|
||||
assert report.labels_to_add == []
|
||||
assert module.READY_LABEL in report.labels_to_remove
|
||||
|
||||
|
||||
def test_validate_pull_request_fails_on_missing_required_content() -> None:
|
||||
module = _load_module()
|
||||
body = """## Description
|
||||
|
||||
Fixes #123
|
||||
|
||||
## Type of Change
|
||||
|
||||
- [ ] New feature (non-breaking change that adds functionality)
|
||||
|
||||
## Changes Made
|
||||
|
||||
- Change 1
|
||||
|
||||
## Testing
|
||||
|
||||
### Test Output
|
||||
|
||||
```text
|
||||
# Paste relevant command output or artifact links here
|
||||
```
|
||||
|
||||
## Real Behavior Proof
|
||||
|
||||
- Environment:
|
||||
- Exact command / steps:
|
||||
- Observed result:
|
||||
- Not tested:
|
||||
|
||||
## Review Readiness
|
||||
|
||||
- [ ] I have performed a self-review
|
||||
- [ ] This PR is ready for human review
|
||||
"""
|
||||
|
||||
report = module.validate_pull_request(_event(body))
|
||||
|
||||
assert report.valid is False
|
||||
assert report.needs_author_action is True
|
||||
assert module.AUTHOR_ACTION_LABEL in report.labels_to_add
|
||||
assert any("Description" in problem for problem in report.problems)
|
||||
assert any("Type of Change" in problem for problem in report.problems)
|
||||
assert any("Test Output" in problem for problem in report.problems)
|
||||
assert any("Real Behavior Proof" in problem for problem in report.problems)
|
||||
|
||||
|
||||
def test_validate_pull_request_skips_bot_authored_prs() -> None:
|
||||
module = _load_module()
|
||||
|
||||
report = module.validate_pull_request(_event("", login="dependabot[bot]"))
|
||||
|
||||
assert report.valid is True
|
||||
assert report.is_bot_pr is True
|
||||
assert report.needs_author_action is False
|
||||
assert report.labels_to_add == []
|
||||
|
|
@ -32,4 +32,6 @@ run_act act workflow_dispatch -W .github/workflows/release.yml -e .github/act/dr
|
|||
# validation step exercises the same code path CI actually fires on.
|
||||
run_act act release -W .github/workflows/release.yml -e .github/act/release-published.json -n
|
||||
run_act act push -W .github/workflows/release-please.yml -e .github/act/push-feat.json -n
|
||||
run_act act pull_request_target -W .github/workflows/pr-health.yml -e .github/act/pr-governance-invalid.json -n
|
||||
run_act act pull_request_target -W .github/workflows/pr-health.yml -e .github/act/pr-governance-valid.json -n
|
||||
run_act act workflow_dispatch -W .github/workflows/docker.yml -e .github/act/docker-version.json -n
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ from headroom.agent_savings import (
|
|||
proxy_pipeline_kwargs,
|
||||
with_target_savings,
|
||||
)
|
||||
from headroom.cli import wrap as wrap_module
|
||||
from headroom.cli.main import main
|
||||
from headroom.compress import CompressConfig, compress
|
||||
from headroom.proxy.models import ProxyConfig
|
||||
|
|
@ -137,6 +138,110 @@ def test_compress_applies_agent_savings_profile_to_pipeline(monkeypatch) -> None
|
|||
assert captured["min_tokens_to_compress"] == 120
|
||||
|
||||
|
||||
def test_compress_savings_profile_does_not_mutate_supplied_config(monkeypatch) -> None:
|
||||
captured: dict[str, object] = {}
|
||||
messages = [{"role": "user", "content": "x" * 500}]
|
||||
config = CompressConfig(
|
||||
compress_user_messages=False,
|
||||
compress_system_messages=False,
|
||||
protect_recent=9,
|
||||
protect_analysis_context=False,
|
||||
target_ratio=None,
|
||||
min_tokens_to_compress=999,
|
||||
)
|
||||
|
||||
class Pipeline:
|
||||
def apply(self, **kwargs):
|
||||
captured.update(kwargs)
|
||||
return SimpleNamespace(
|
||||
messages=messages,
|
||||
tokens_before=1000,
|
||||
tokens_after=100,
|
||||
transforms_applied=["test"],
|
||||
)
|
||||
|
||||
monkeypatch.setattr(compress_module, "_get_pipeline", lambda: Pipeline())
|
||||
|
||||
compress(messages, config=config, savings_profile=AGENT_90_PROFILE)
|
||||
|
||||
assert captured["target_ratio"] == 0.10
|
||||
assert captured["min_tokens_to_compress"] == 120
|
||||
assert config.compress_user_messages is False
|
||||
assert config.compress_system_messages is False
|
||||
assert config.protect_recent == 9
|
||||
assert config.protect_analysis_context is False
|
||||
assert config.target_ratio is None
|
||||
assert config.min_tokens_to_compress == 999
|
||||
|
||||
|
||||
def test_agent_savings_config_mismatches_returns_specific_labels() -> None:
|
||||
profile = get_agent_savings_profile(AGENT_90_PROFILE)
|
||||
running_config = {
|
||||
"savings_profile": profile.name,
|
||||
"target_ratio": 0.20,
|
||||
"compress_user_messages": profile.compress_user_messages,
|
||||
"compress_system_messages": profile.compress_system_messages,
|
||||
"protect_recent": profile.protect_recent,
|
||||
"protect_analysis_context": profile.protect_analysis_context,
|
||||
"min_tokens_to_crush": profile.min_tokens_to_compress,
|
||||
"max_items_after_crush": profile.max_items_after_crush,
|
||||
"smart_crusher_with_compaction": profile.smart_crusher_with_compaction,
|
||||
"accuracy_guard": profile.accuracy_guard,
|
||||
}
|
||||
|
||||
assert wrap_module._agent_savings_config_mismatches(running_config, "codex") == ["target-ratio"]
|
||||
|
||||
|
||||
def test_agent_savings_config_mismatches_ignores_non_target_agents() -> None:
|
||||
assert wrap_module._agent_savings_config_mismatches({}, "openhands") == []
|
||||
|
||||
|
||||
def test_agent_savings_config_mismatches_accepts_matching_runtime_config() -> None:
|
||||
profile = get_agent_savings_profile(AGENT_90_PROFILE)
|
||||
running_config = {
|
||||
"savings_profile": profile.name,
|
||||
"target_ratio": "0.10",
|
||||
"compress_user_messages": True,
|
||||
"compress_system_messages": True,
|
||||
"protect_recent": "2",
|
||||
"protect_analysis_context": True,
|
||||
"min_tokens_to_crush": "120",
|
||||
"max_items_after_crush": "8",
|
||||
"smart_crusher_with_compaction": False,
|
||||
"accuracy_guard": "strict",
|
||||
}
|
||||
|
||||
assert wrap_module._agent_savings_config_mismatches(running_config, "cursor") == []
|
||||
|
||||
|
||||
def test_agent_savings_config_mismatches_reports_unparseable_values() -> None:
|
||||
running_config = {
|
||||
"savings_profile": None,
|
||||
"target_ratio": "not-a-float",
|
||||
"compress_user_messages": None,
|
||||
"compress_system_messages": None,
|
||||
"protect_recent": "not-an-int",
|
||||
"protect_analysis_context": None,
|
||||
"min_tokens_to_crush": object(),
|
||||
"max_items_after_crush": object(),
|
||||
"smart_crusher_with_compaction": None,
|
||||
"accuracy_guard": None,
|
||||
}
|
||||
|
||||
assert wrap_module._agent_savings_config_mismatches(running_config, "claude") == [
|
||||
"savings-profile",
|
||||
"target-ratio",
|
||||
"compress-user-messages",
|
||||
"compress-system-messages",
|
||||
"protect-recent",
|
||||
"protect-analysis-context",
|
||||
"min-tokens",
|
||||
"max-items",
|
||||
"smart-crusher-compaction",
|
||||
"accuracy-guard",
|
||||
]
|
||||
|
||||
|
||||
def test_agent_90_profile_applies_to_proxy_config_runtime_kwargs() -> None:
|
||||
config = ProxyConfig(savings_profile="agent-90")
|
||||
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ from headroom.proxy.auth_mode import (
|
|||
SUBSCRIPTION_UA_PREFIXES,
|
||||
AuthMode,
|
||||
classify_auth_mode,
|
||||
classify_client,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -186,3 +187,9 @@ def test_classify_under_100us_per_call() -> None:
|
|||
per_call_us = (elapsed / iters) * 1_000_000
|
||||
|
||||
assert per_call_us < 100, f"classify_auth_mode took {per_call_us:.2f} us/call (limit: 100 us)"
|
||||
|
||||
|
||||
def test_classify_client_uses_default_when_no_client_signal():
|
||||
headers = {"user-agent": "anthropic/0.42.0"}
|
||||
|
||||
assert classify_client(headers, default="claude") == "claude"
|
||||
|
|
|
|||
50
tests/test_cli/test_copilot_auth.py
Normal file
50
tests/test_cli/test_copilot_auth.py
Normal file
|
|
@ -0,0 +1,50 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from click.testing import CliRunner
|
||||
|
||||
from headroom.cli import main
|
||||
|
||||
|
||||
def test_copilot_auth_login_saves_token(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
auth_file = tmp_path / "copilot_auth.json"
|
||||
monkeypatch.setenv("HEADROOM_COPILOT_AUTH_FILE", str(auth_file))
|
||||
monkeypatch.setattr(
|
||||
"headroom.cli.copilot_auth.start_copilot_device_authorization",
|
||||
lambda domain: {
|
||||
"verification_uri": "https://github.com/login/device",
|
||||
"user_code": "ABCD-1234",
|
||||
"device_code": "device-code",
|
||||
"interval": 1,
|
||||
"expires_in": 900,
|
||||
},
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"headroom.cli.copilot_auth.poll_copilot_device_authorization",
|
||||
lambda device_code, *, domain, interval, expires_in: "gho-headroom",
|
||||
)
|
||||
|
||||
result = CliRunner().invoke(main, ["copilot-auth", "login"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert "https://github.com/login/device" in result.output
|
||||
assert "ABCD-1234" in result.output
|
||||
assert "gho-headroom" not in result.output
|
||||
assert auth_file.exists()
|
||||
|
||||
|
||||
def test_copilot_auth_status_reports_missing_login(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
monkeypatch.setenv("HEADROOM_COPILOT_AUTH_FILE", str(tmp_path / "missing.json"))
|
||||
|
||||
result = CliRunner().invoke(main, ["copilot-auth", "status"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert "Status: not logged in" in result.output
|
||||
|
|
@ -529,6 +529,30 @@ def test_ensure_codex_provider_replaces_existing_model_provider(
|
|||
assert parsed["features"]["hooks"] is True
|
||||
|
||||
|
||||
def test_ensure_codex_provider_emits_requires_openai_auth_for_chatgpt(
|
||||
monkeypatch, tmp_path: Path
|
||||
) -> None:
|
||||
init_cli, _ = _load_init_module(monkeypatch)
|
||||
path = tmp_path / "config.toml"
|
||||
(tmp_path / "auth.json").write_text('{"auth_mode": "chatgpt"}', encoding="utf-8")
|
||||
|
||||
init_cli._ensure_codex_provider(path, 8787)
|
||||
|
||||
assert "requires_openai_auth = true" in path.read_text(encoding="utf-8")
|
||||
|
||||
|
||||
def test_ensure_codex_provider_omits_requires_openai_auth_for_api_key(
|
||||
monkeypatch, tmp_path: Path
|
||||
) -> None:
|
||||
init_cli, _ = _load_init_module(monkeypatch)
|
||||
path = tmp_path / "config.toml"
|
||||
(tmp_path / "auth.json").write_text('{"auth_mode": "apikey"}', encoding="utf-8")
|
||||
|
||||
init_cli._ensure_codex_provider(path, 8787)
|
||||
|
||||
assert "requires_openai_auth" not in path.read_text(encoding="utf-8")
|
||||
|
||||
|
||||
def test_ensure_codex_feature_flag_replaces_existing_marker(monkeypatch, tmp_path: Path) -> None:
|
||||
init_cli, _ = _load_init_module(monkeypatch)
|
||||
path = tmp_path / "config.toml"
|
||||
|
|
|
|||
|
|
@ -350,6 +350,30 @@ class TestSubscriptionRouting:
|
|||
content = (tmp_path / ".codex" / "config.toml").read_text()
|
||||
assert 'openai_base_url = "http://127.0.0.1:8787/v1"' in content
|
||||
|
||||
def test_inject_emits_requires_openai_auth_for_chatgpt(
|
||||
self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
_set_test_home(monkeypatch, tmp_path)
|
||||
config_dir = tmp_path / ".codex"
|
||||
config_dir.mkdir()
|
||||
(config_dir / "auth.json").write_text('{"auth_mode": "chatgpt"}', encoding="utf-8")
|
||||
|
||||
wrap_mod._inject_codex_provider_config(8787)
|
||||
|
||||
assert "requires_openai_auth = true" in (config_dir / "config.toml").read_text()
|
||||
|
||||
def test_inject_omits_requires_openai_auth_for_api_key(
|
||||
self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
_set_test_home(monkeypatch, tmp_path)
|
||||
config_dir = tmp_path / ".codex"
|
||||
config_dir.mkdir()
|
||||
(config_dir / "auth.json").write_text('{"auth_mode": "apikey"}', encoding="utf-8")
|
||||
|
||||
wrap_mod._inject_codex_provider_config(8787)
|
||||
|
||||
assert "requires_openai_auth" not in (config_dir / "config.toml").read_text()
|
||||
|
||||
def test_openai_base_url_port_updates_on_rewrap(
|
||||
self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
|
|
@ -517,6 +541,67 @@ def test_start_proxy_uses_separate_session_for_signal_isolation(
|
|||
assert popen_kwargs["start_new_session"] == (wrap_mod.os.name == "posix")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("agent_type", ["claude", "codex", "cursor"])
|
||||
def test_start_proxy_applies_agent_90_defaults(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path: Path, agent_type: str
|
||||
) -> None:
|
||||
"""Wrapped coding agents should start the proxy with high-savings defaults."""
|
||||
popen_kwargs: dict[str, object] = {}
|
||||
|
||||
class FakeProc:
|
||||
returncode = None
|
||||
|
||||
def poll(self) -> None:
|
||||
return None
|
||||
|
||||
def fake_popen(*args: object, **kwargs: object) -> FakeProc:
|
||||
popen_kwargs.update(kwargs)
|
||||
return FakeProc()
|
||||
|
||||
monkeypatch.setattr(wrap_mod, "_get_log_path", lambda: tmp_path / "proxy.log")
|
||||
monkeypatch.setattr(wrap_mod, "_check_proxy", lambda port: True)
|
||||
monkeypatch.setattr(wrap_mod.subprocess, "Popen", fake_popen)
|
||||
|
||||
wrap_mod._start_proxy(8787, agent_type=agent_type)
|
||||
|
||||
env = popen_kwargs["env"]
|
||||
assert isinstance(env, dict)
|
||||
assert env["HEADROOM_SAVINGS_PROFILE"] == "agent-90"
|
||||
assert env["HEADROOM_TARGET_RATIO"] == "0.10"
|
||||
assert env["HEADROOM_MAX_ITEMS"] == "8"
|
||||
assert env["HEADROOM_SMART_CRUSHER_COMPACTION"] == "0"
|
||||
|
||||
|
||||
def test_start_proxy_preserves_explicit_savings_overrides(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
"""User-provided savings env vars should override wrapper defaults."""
|
||||
popen_kwargs: dict[str, object] = {}
|
||||
|
||||
class FakeProc:
|
||||
returncode = None
|
||||
|
||||
def poll(self) -> None:
|
||||
return None
|
||||
|
||||
def fake_popen(*args: object, **kwargs: object) -> FakeProc:
|
||||
popen_kwargs.update(kwargs)
|
||||
return FakeProc()
|
||||
|
||||
monkeypatch.setenv("HEADROOM_TARGET_RATIO", "0.20")
|
||||
monkeypatch.setenv("HEADROOM_MAX_ITEMS", "12")
|
||||
monkeypatch.setattr(wrap_mod, "_get_log_path", lambda: tmp_path / "proxy.log")
|
||||
monkeypatch.setattr(wrap_mod, "_check_proxy", lambda port: True)
|
||||
monkeypatch.setattr(wrap_mod.subprocess, "Popen", fake_popen)
|
||||
|
||||
wrap_mod._start_proxy(8787, agent_type="codex")
|
||||
|
||||
env = popen_kwargs["env"]
|
||||
assert isinstance(env, dict)
|
||||
assert env["HEADROOM_TARGET_RATIO"] == "0.20"
|
||||
assert env["HEADROOM_MAX_ITEMS"] == "12"
|
||||
|
||||
|
||||
def test_launch_tool_ignores_sigint_in_wrapper(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
|
|
|
|||
|
|
@ -14,7 +14,7 @@ import click
|
|||
import pytest
|
||||
from click.testing import CliRunner
|
||||
|
||||
from headroom.copilot_auth import DEFAULT_API_URL
|
||||
from headroom.copilot_auth import DEFAULT_API_URL, CopilotSubscriptionTokenResolution
|
||||
|
||||
|
||||
def _expected_project_prefix() -> str:
|
||||
|
|
@ -27,6 +27,22 @@ def runner() -> CliRunner:
|
|||
return CliRunner()
|
||||
|
||||
|
||||
def _subscription_resolution(
|
||||
token: str = "gho-existing",
|
||||
*,
|
||||
api_url: str = DEFAULT_API_URL,
|
||||
source: str = "headroom-copilot-auth:/tmp/copilot_auth.json:token-exchange",
|
||||
confidence: str = "copilot-token-exchange",
|
||||
) -> CopilotSubscriptionTokenResolution:
|
||||
return CopilotSubscriptionTokenResolution(
|
||||
token=token,
|
||||
source=source,
|
||||
confidence=confidence,
|
||||
api_url=api_url,
|
||||
token_fingerprint="sha256:0123456789ab",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def wrap_modules(monkeypatch: pytest.MonkeyPatch) -> tuple[types.ModuleType, click.Group]:
|
||||
headroom_pkg = sys.modules.get("headroom")
|
||||
|
|
@ -247,7 +263,10 @@ def test_wrap_copilot_subscription_uses_github_auth_without_provider_key(
|
|||
|
||||
with (
|
||||
patch("headroom.cli.wrap.shutil.which", return_value="copilot"),
|
||||
patch("headroom.cli.wrap.resolve_subscription_bearer_token", return_value="gho-existing"),
|
||||
patch(
|
||||
"headroom.cli.wrap.resolve_subscription_bearer_token_details",
|
||||
return_value=_subscription_resolution(),
|
||||
),
|
||||
patch("headroom.cli.wrap.has_oauth_auth", return_value=False),
|
||||
patch("headroom.cli.wrap._launch_tool", side_effect=fake_launch_tool),
|
||||
):
|
||||
|
|
@ -284,7 +303,10 @@ def test_wrap_copilot_subscription_defaults_to_responses_for_reasoning_model(
|
|||
|
||||
with (
|
||||
patch("headroom.cli.wrap.shutil.which", return_value="copilot"),
|
||||
patch("headroom.cli.wrap.resolve_subscription_bearer_token", return_value="gho-existing"),
|
||||
patch(
|
||||
"headroom.cli.wrap.resolve_subscription_bearer_token_details",
|
||||
return_value=_subscription_resolution("gho-existing"),
|
||||
),
|
||||
patch("headroom.cli.wrap.has_oauth_auth", return_value=False),
|
||||
patch("headroom.cli.wrap._launch_tool", side_effect=fake_launch_tool),
|
||||
):
|
||||
|
|
@ -321,7 +343,10 @@ def test_wrap_copilot_subscription_keeps_gpt4_on_completions(
|
|||
|
||||
with (
|
||||
patch("headroom.cli.wrap.shutil.which", return_value="copilot"),
|
||||
patch("headroom.cli.wrap.resolve_subscription_bearer_token", return_value="gho-existing"),
|
||||
patch(
|
||||
"headroom.cli.wrap.resolve_subscription_bearer_token_details",
|
||||
return_value=_subscription_resolution("gho-existing"),
|
||||
),
|
||||
patch("headroom.cli.wrap.has_oauth_auth", return_value=False),
|
||||
patch("headroom.cli.wrap._launch_tool", side_effect=fake_launch_tool),
|
||||
):
|
||||
|
|
@ -351,7 +376,10 @@ def test_wrap_copilot_subscription_allows_explicit_responses_wire_api(
|
|||
|
||||
with (
|
||||
patch("headroom.cli.wrap.shutil.which", return_value="copilot"),
|
||||
patch("headroom.cli.wrap.resolve_subscription_bearer_token", return_value="gho-existing"),
|
||||
patch(
|
||||
"headroom.cli.wrap.resolve_subscription_bearer_token_details",
|
||||
return_value=_subscription_resolution("gho-existing"),
|
||||
),
|
||||
patch("headroom.cli.wrap.has_oauth_auth", return_value=False),
|
||||
patch("headroom.cli.wrap._launch_tool", side_effect=fake_launch_tool),
|
||||
):
|
||||
|
|
@ -403,10 +431,9 @@ def test_wrap_copilot_subscription_pins_validated_token_for_proxy(
|
|||
with (
|
||||
patch("headroom.cli.wrap.shutil.which", return_value="copilot"),
|
||||
patch(
|
||||
"headroom.cli.wrap.resolve_subscription_bearer_token",
|
||||
return_value="gho-validated",
|
||||
"headroom.cli.wrap.resolve_subscription_bearer_token_details",
|
||||
return_value=_subscription_resolution("gho-validated", api_url=business_api),
|
||||
),
|
||||
patch("headroom.cli.wrap.resolve_copilot_api_url", return_value=business_api),
|
||||
patch("headroom.cli.wrap.has_oauth_auth", return_value=False),
|
||||
patch("headroom.cli.wrap._launch_tool", side_effect=fake_launch_tool),
|
||||
):
|
||||
|
|
@ -437,12 +464,13 @@ def test_wrap_copilot_subscription_requires_reusable_auth(
|
|||
_wrap_cli, main = wrap_modules
|
||||
with (
|
||||
patch("headroom.cli.wrap.shutil.which", return_value="copilot"),
|
||||
patch("headroom.cli.wrap.resolve_subscription_bearer_token", return_value=None),
|
||||
patch("headroom.cli.wrap.resolve_subscription_bearer_token_details", return_value=None),
|
||||
):
|
||||
result = runner.invoke(main, ["wrap", "copilot", "--subscription", "--no-rtk"])
|
||||
|
||||
assert result.exit_code != 0
|
||||
assert "subscription mode requires a reusable GitHub/Copilot bearer token" in result.output
|
||||
assert "headroom copilot-auth login" in result.output
|
||||
|
||||
|
||||
def test_wrap_copilot_subscription_rejects_translated_backend(
|
||||
|
|
@ -634,6 +662,8 @@ def _clear_copilot_env(monkeypatch: pytest.MonkeyPatch) -> None:
|
|||
"OPENAI_API_KEY",
|
||||
"ANTHROPIC_API_KEY",
|
||||
"GITHUB_COPILOT_API_URL",
|
||||
"GITHUB_COPILOT_ENTERPRISE_URL",
|
||||
"GITHUB_COPILOT_ENTERPRISE_DOMAIN",
|
||||
"GITHUB_COPILOT_TOKEN",
|
||||
"GITHUB_COPILOT_GITHUB_TOKEN",
|
||||
"COPILOT_MODEL",
|
||||
|
|
@ -747,17 +777,15 @@ def test_wrap_copilot_byok_never_resolves_copilot_endpoint(
|
|||
assert env["COPILOT_PROVIDER_TYPE"] == "openai"
|
||||
|
||||
|
||||
def test_wrap_copilot_subscription_uses_generic_endpoint_not_account(
|
||||
def test_wrap_copilot_subscription_uses_resolved_subscription_endpoint(
|
||||
runner: CliRunner,
|
||||
wrap_modules: tuple[types.ModuleType, click.Group],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""#610 (subscription has the same latent bug): --subscription must route to
|
||||
the generic host too, even when /copilot_internal/user advertises an
|
||||
account-specific host. The segmented host does not serve newer models on the
|
||||
responses API, and it is not the host the official Copilot client uses."""
|
||||
"""Subscription mode uses the endpoint returned with the resolved token."""
|
||||
_wrap_cli, main = wrap_modules
|
||||
_clear_copilot_env(monkeypatch)
|
||||
business_api = "https://api.business.githubcopilot.com"
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
def fake_launch_tool(**kwargs): # noqa: ANN003
|
||||
|
|
@ -765,7 +793,10 @@ def test_wrap_copilot_subscription_uses_generic_endpoint_not_account(
|
|||
|
||||
with (
|
||||
patch("headroom.cli.wrap.shutil.which", return_value="copilot"),
|
||||
patch("headroom.cli.wrap.resolve_subscription_bearer_token", return_value="gho-sub"),
|
||||
patch(
|
||||
"headroom.cli.wrap.resolve_subscription_bearer_token_details",
|
||||
return_value=_subscription_resolution("copilot-api", api_url=business_api),
|
||||
),
|
||||
patch("headroom.cli.wrap.has_oauth_auth", return_value=True),
|
||||
patch("headroom.copilot_auth._fetch_copilot_user_info", return_value=_ACCOUNT_USER_INFO),
|
||||
patch("headroom.cli.wrap._launch_tool", side_effect=fake_launch_tool),
|
||||
|
|
@ -778,9 +809,9 @@ def test_wrap_copilot_subscription_uses_generic_endpoint_not_account(
|
|||
assert result.exit_code == 0, result.output
|
||||
env = captured["env"]
|
||||
assert isinstance(env, dict)
|
||||
assert captured["openai_api_url"] == DEFAULT_API_URL
|
||||
assert env["OPENAI_TARGET_API_URL"] == DEFAULT_API_URL
|
||||
assert env["COPILOT_PROVIDER_BEARER_TOKEN"] == "gho-sub"
|
||||
assert captured["openai_api_url"] == business_api
|
||||
assert env["OPENAI_TARGET_API_URL"] == business_api
|
||||
assert env["COPILOT_PROVIDER_BEARER_TOKEN"] == "copilot-api"
|
||||
|
||||
|
||||
def test_wrap_copilot_subscription_honors_api_url_override(
|
||||
|
|
@ -800,7 +831,15 @@ def test_wrap_copilot_subscription_honors_api_url_override(
|
|||
|
||||
with (
|
||||
patch("headroom.cli.wrap.shutil.which", return_value="copilot"),
|
||||
patch("headroom.cli.wrap.resolve_subscription_bearer_token", return_value="gho-sub"),
|
||||
patch(
|
||||
"headroom.cli.wrap.resolve_subscription_bearer_token_details",
|
||||
return_value=_subscription_resolution(
|
||||
"gho-sub",
|
||||
api_url="https://api.enterprise.example.com",
|
||||
source="env:GITHUB_COPILOT_API_TOKEN",
|
||||
confidence="explicit-api-token",
|
||||
),
|
||||
),
|
||||
patch("headroom.cli.wrap.has_oauth_auth", return_value=True),
|
||||
patch("headroom.cli.wrap._launch_tool", side_effect=fake_launch_tool),
|
||||
):
|
||||
|
|
|
|||
|
|
@ -156,6 +156,35 @@ def test_perf_json_raw_is_array(runner, monkeypatch):
|
|||
assert data[0]["request_id"] == "hr_1"
|
||||
|
||||
|
||||
def test_perf_json_raw_preserves_client_field(runner, monkeypatch):
|
||||
report = _sample_report()
|
||||
report.perf_records[0].client = "codex"
|
||||
_patch_report(monkeypatch, report)
|
||||
|
||||
result = runner.invoke(main, ["perf", "--format", "json", "--raw"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
data = json.loads(result.output)
|
||||
assert data[0]["client"] == "codex"
|
||||
|
||||
|
||||
def test_parse_perf_line_preserves_client_field(monkeypatch, tmp_path):
|
||||
log_dir = tmp_path / "logs"
|
||||
log_dir.mkdir()
|
||||
(log_dir / "proxy.log").write_text(
|
||||
"2026-06-10 10:00:00,000 - headroom.proxy - INFO - "
|
||||
"[hr_codex] PERF model=gpt-5 msgs=3 tok_before=1000 "
|
||||
"tok_after=90 tok_saved=910 cache_read=0 cache_write=0 "
|
||||
"cache_hit_pct=0 opt_ms=12 transforms=content_router client=codex\n"
|
||||
)
|
||||
monkeypatch.setattr(analyzer, "LOG_DIR", log_dir)
|
||||
|
||||
report = analyzer.parse_log_files(last_n_hours=0)
|
||||
|
||||
assert len(report.perf_records) == 1
|
||||
assert report.perf_records[0].client == "codex"
|
||||
|
||||
|
||||
def test_perf_csv_by_model(runner, monkeypatch):
|
||||
_patch_report(monkeypatch, _sample_report())
|
||||
result = runner.invoke(main, ["perf", "--format", "csv"])
|
||||
|
|
@ -191,3 +220,22 @@ def test_perf_rejects_unknown_format(runner, monkeypatch):
|
|||
_patch_report(monkeypatch, _sample_report())
|
||||
result = runner.invoke(main, ["perf", "--format", "xml"])
|
||||
assert result.exit_code != 0
|
||||
|
||||
|
||||
def test_parse_perf_line_preserves_blank_client_field(
|
||||
tmp_path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
logs_dir = tmp_path / "logs"
|
||||
logs_dir.mkdir()
|
||||
monkeypatch.setattr(analyzer, "LOG_DIR", logs_dir)
|
||||
(logs_dir / "proxy.log").write_text(
|
||||
"2026-06-10 10:00:00,000 - headroom.proxy - INFO - [req-blank] PERF "
|
||||
"model=gpt-5 msgs=1 tok_before=100 tok_after=50 tok_saved=50 "
|
||||
"cache_read=0 cache_write=0 cache_hit_pct=0 opt_ms=1 transforms=test client=\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
report = analyzer.parse_log_files(last_n_hours=0)
|
||||
|
||||
assert len(report.perf_records) == 1
|
||||
assert report.perf_records[0].client == ""
|
||||
|
|
|
|||
|
|
@ -76,6 +76,36 @@ class TestCLIWrapProxyTimeout:
|
|||
assert sleeps == [1]
|
||||
assert fake_proc.killed is False
|
||||
|
||||
def test_start_proxy_passes_resolved_copilot_api_url_to_proxy(self, monkeypatch, tmp_path):
|
||||
fake_proc = _FakeProxyProcess()
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
monkeypatch.delenv(wrap_mod._WRAP_PROXY_TIMEOUT_ENV, raising=False)
|
||||
monkeypatch.setattr(wrap_mod, "_ml_wrap_extras_detected", lambda: False)
|
||||
monkeypatch.setattr(wrap_mod, "_get_log_path", lambda: tmp_path / "proxy.log")
|
||||
monkeypatch.setattr(wrap_mod, "_check_proxy", lambda _port: True)
|
||||
monkeypatch.setattr(wrap_mod.time, "sleep", lambda _seconds: None)
|
||||
|
||||
def fake_popen(*args, **kwargs): # noqa: ANN002, ANN003
|
||||
captured["args"] = args
|
||||
captured["kwargs"] = kwargs
|
||||
return fake_proc
|
||||
|
||||
monkeypatch.setattr(wrap_mod.subprocess, "Popen", fake_popen)
|
||||
|
||||
proc = wrap_mod._start_proxy(
|
||||
8787,
|
||||
agent_type="copilot",
|
||||
openai_api_url="https://copilot-api.acme.ghe.com",
|
||||
copilot_api_token="copilot-api-token",
|
||||
)
|
||||
|
||||
assert proc is fake_proc
|
||||
env = captured["kwargs"]["env"]
|
||||
assert env["OPENAI_TARGET_API_URL"] == "https://copilot-api.acme.ghe.com"
|
||||
assert env["GITHUB_COPILOT_API_URL"] == "https://copilot-api.acme.ghe.com"
|
||||
assert env["GITHUB_COPILOT_API_TOKEN"] == "copilot-api-token"
|
||||
|
||||
def test_env_timeout_allows_slow_start_proxy_to_succeed(self, monkeypatch, tmp_path):
|
||||
fake_proc = _FakeProxyProcess()
|
||||
sleeps = []
|
||||
|
|
@ -156,6 +186,24 @@ class TestCLIProxyEnvVars:
|
|||
assert result.exit_code == 0, result.output
|
||||
assert captured_config["config"].port == 9797
|
||||
|
||||
def test_headroom_min_tokens_from_env(self, runner):
|
||||
"""HEADROOM_MIN_TOKENS env var should be passed to ProxyConfig."""
|
||||
captured_config = {}
|
||||
|
||||
def mock_run_server(config, **kwargs):
|
||||
captured_config["config"] = config
|
||||
|
||||
with patch("headroom.proxy.server.run_server", mock_run_server):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["proxy"],
|
||||
env={"HEADROOM_MIN_TOKENS": "120"},
|
||||
catch_exceptions=False,
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert captured_config["config"].min_tokens_to_crush == 120
|
||||
|
||||
def test_headroom_budget_from_env(self, runner):
|
||||
"""HEADROOM_BUDGET env var should be passed to ProxyConfig."""
|
||||
captured_config = {}
|
||||
|
|
|
|||
|
|
@ -2,14 +2,48 @@
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
|
||||
import headroom.subscription.codex_rate_limits as crl
|
||||
from headroom.subscription.codex_rate_limits import (
|
||||
CodexRateLimitState,
|
||||
CodexRateLimitWindow,
|
||||
_build_usage_headers,
|
||||
maybe_schedule_usage_poll,
|
||||
parse_codex_rate_limits,
|
||||
parse_codex_usage_payload,
|
||||
)
|
||||
|
||||
# A faithful GET /wham/usage body (shape captured from a live Plus account).
|
||||
USAGE_PAYLOAD = {
|
||||
"plan_type": "plus",
|
||||
"rate_limit": {
|
||||
"allowed": True,
|
||||
"limit_reached": False,
|
||||
"primary_window": {
|
||||
"used_percent": 23,
|
||||
"limit_window_seconds": 18000,
|
||||
"reset_after_seconds": 12266,
|
||||
"reset_at": 1781276043,
|
||||
},
|
||||
"secondary_window": {
|
||||
"used_percent": 6,
|
||||
"limit_window_seconds": 604800,
|
||||
"reset_after_seconds": 359170,
|
||||
"reset_at": 1781622947,
|
||||
},
|
||||
},
|
||||
"additional_rate_limits": None,
|
||||
"credits": {
|
||||
"has_credits": False,
|
||||
"unlimited": False,
|
||||
"balance": "0",
|
||||
},
|
||||
"rate_limit_reached_type": None,
|
||||
"promo": None,
|
||||
}
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CodexRateLimitWindow helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -225,3 +259,163 @@ class TestCodexRateLimitState:
|
|||
assert snap is not None
|
||||
assert snap.primary is not None
|
||||
assert snap.primary.used_percent == 90.0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# parse_codex_usage_payload (GET /wham/usage)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestParseCodexUsagePayload:
|
||||
def test_parses_full_payload(self):
|
||||
snap = parse_codex_usage_payload(USAGE_PAYLOAD)
|
||||
assert snap is not None
|
||||
assert snap.primary is not None
|
||||
assert snap.primary.used_percent == 23.0
|
||||
assert snap.primary.window_minutes == 300 # 18000s rounded up
|
||||
assert snap.primary.resets_at == 1781276043
|
||||
assert snap.secondary is not None
|
||||
assert snap.secondary.used_percent == 6.0
|
||||
assert snap.secondary.window_minutes == 10080 # 604800s
|
||||
|
||||
def test_window_minutes_rounds_up(self):
|
||||
snap = parse_codex_usage_payload(
|
||||
{"rate_limit": {"primary_window": {"used_percent": 1, "limit_window_seconds": 61}}}
|
||||
)
|
||||
assert snap is not None
|
||||
assert snap.primary is not None
|
||||
assert snap.primary.window_minutes == 2
|
||||
|
||||
def test_no_credits_balance_suppressed(self):
|
||||
# has_credits False -> balance must not surface as "0".
|
||||
snap = parse_codex_usage_payload(USAGE_PAYLOAD)
|
||||
assert snap is not None
|
||||
assert snap.credits is not None
|
||||
assert snap.credits.has_credits is False
|
||||
assert snap.credits.balance is None
|
||||
|
||||
def test_credits_balance_kept_when_has_credits(self):
|
||||
payload = {
|
||||
"rate_limit": {"primary_window": {"used_percent": 5}},
|
||||
"credits": {"has_credits": True, "unlimited": False, "balance": "$5.00"},
|
||||
}
|
||||
snap = parse_codex_usage_payload(payload)
|
||||
assert snap is not None
|
||||
assert snap.credits is not None
|
||||
assert snap.credits.balance == "$5.00"
|
||||
|
||||
def test_promo_object_message(self):
|
||||
payload = {
|
||||
"rate_limit": {"primary_window": {"used_percent": 5}},
|
||||
"promo": {"message": "Hello"},
|
||||
}
|
||||
snap = parse_codex_usage_payload(payload)
|
||||
assert snap is not None
|
||||
assert snap.promo_message == "Hello"
|
||||
|
||||
def test_returns_none_for_empty(self):
|
||||
assert parse_codex_usage_payload({}) is None
|
||||
assert parse_codex_usage_payload(None) is None
|
||||
assert parse_codex_usage_payload({"rate_limit": {}}) is None
|
||||
|
||||
def test_missing_used_percent_window_skipped(self):
|
||||
snap = parse_codex_usage_payload(
|
||||
{"rate_limit": {"primary_window": {"limit_window_seconds": 60}}}
|
||||
)
|
||||
assert snap is None
|
||||
|
||||
def test_update_from_usage_payload_stores(self):
|
||||
state = CodexRateLimitState()
|
||||
assert state.update_from_usage_payload(USAGE_PAYLOAD) is True
|
||||
snap = state.latest
|
||||
assert snap is not None
|
||||
assert snap.primary is not None
|
||||
assert snap.primary.used_percent == 23.0
|
||||
|
||||
def test_update_from_usage_payload_noop_returns_false(self):
|
||||
state = CodexRateLimitState()
|
||||
assert state.update_from_usage_payload({}) is False
|
||||
assert state.latest is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Usage poll: header gating + throttle
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestUsagePollGating:
|
||||
def test_build_headers_requires_account_id(self):
|
||||
assert _build_usage_headers({"authorization": "Bearer abc.def.ghi"}) is None
|
||||
|
||||
def test_build_headers_requires_bearer(self):
|
||||
assert _build_usage_headers({"chatgpt-account-id": "acct"}) is None
|
||||
assert (
|
||||
_build_usage_headers({"authorization": "sk-live", "chatgpt-account-id": "acct"}) is None
|
||||
)
|
||||
|
||||
def test_build_headers_happy_path(self):
|
||||
headers = _build_usage_headers(
|
||||
{
|
||||
"Authorization": "Bearer abc.def.ghi",
|
||||
"ChatGPT-Account-Id": "acct-1",
|
||||
"User-Agent": "codex_exec/0.139.0",
|
||||
"originator": "codex_exec",
|
||||
}
|
||||
)
|
||||
assert headers is not None
|
||||
assert headers["Authorization"] == "Bearer abc.def.ghi"
|
||||
assert headers["ChatGPT-Account-Id"] == "acct-1"
|
||||
assert headers["User-Agent"] == "codex_exec/0.139.0"
|
||||
assert headers["originator"] == "codex_exec"
|
||||
|
||||
def test_try_begin_poll_throttles(self):
|
||||
state = CodexRateLimitState()
|
||||
assert state._try_begin_poll(60.0) is True
|
||||
# Second immediate attempt is throttled (within interval).
|
||||
assert state._try_begin_poll(60.0) is False
|
||||
state._end_poll()
|
||||
# Still throttled by time even after the in-flight flag clears.
|
||||
assert state._try_begin_poll(60.0) is False
|
||||
# A zero interval always allows once the in-flight flag is clear.
|
||||
assert state._try_begin_poll(0.0) is True
|
||||
|
||||
def test_maybe_schedule_returns_false_without_loop(self):
|
||||
# No running event loop -> cannot schedule.
|
||||
assert (
|
||||
maybe_schedule_usage_poll(
|
||||
{"authorization": "Bearer a.b.c", "chatgpt-account-id": "acct"}
|
||||
)
|
||||
is False
|
||||
)
|
||||
|
||||
def test_maybe_schedule_skips_non_codex(self):
|
||||
async def run():
|
||||
return maybe_schedule_usage_poll({"authorization": "Bearer a.b.c"})
|
||||
|
||||
assert asyncio.run(run()) is False
|
||||
|
||||
def test_maybe_schedule_creates_task_and_throttles(self, monkeypatch):
|
||||
# Replace the network fetch with a fast no-op coroutine.
|
||||
calls: list[str] = []
|
||||
|
||||
async def fake_fetch(url, headers): # noqa: ANN001
|
||||
calls.append(url)
|
||||
crl.get_codex_rate_limit_state()._end_poll()
|
||||
|
||||
monkeypatch.setattr(crl, "_fetch_and_store_usage", fake_fetch)
|
||||
# Reset the singleton's throttle so this test is deterministic.
|
||||
monkeypatch.setattr(crl, "_state", None)
|
||||
monkeypatch.setattr(crl, "_state_lock", crl.Lock())
|
||||
|
||||
async def run():
|
||||
req = {"authorization": "Bearer a.b.c", "chatgpt-account-id": "acct"}
|
||||
first = maybe_schedule_usage_poll(req, min_interval_s=60.0)
|
||||
second = maybe_schedule_usage_poll(req, min_interval_s=60.0)
|
||||
# Let the scheduled task run.
|
||||
await asyncio.sleep(0)
|
||||
return first, second
|
||||
|
||||
first, second = asyncio.run(run())
|
||||
assert first is True
|
||||
assert second is False # throttled
|
||||
assert calls == [crl.CODEX_USAGE_URL]
|
||||
|
|
|
|||
157
tests/test_codex_responses_waste_signals.py
Normal file
157
tests/test_codex_responses_waste_signals.py
Normal file
|
|
@ -0,0 +1,157 @@
|
|||
"""Codex (OpenAI Responses API) waste-signal visibility (issue #820).
|
||||
|
||||
The /v1/responses path never ran ``parse_messages``: compression goes through
|
||||
CompressionUnits (not TransformPipeline), and the minimal ``messages`` list it
|
||||
synthesises drops list-typed ``input`` entirely — so tool output never reached
|
||||
waste detection and the dashboard "What Headroom Removed" stayed empty for
|
||||
Codex traffic.
|
||||
|
||||
The fix is telemetry-only:
|
||||
|
||||
1. ``_responses_input_to_waste_messages`` converts a Responses payload into
|
||||
OpenAI-style messages — tool output items (``function_call_output`` etc.)
|
||||
become ``role="tool"`` messages, ``message`` items keep their role and
|
||||
joined part text.
|
||||
2. ``handle_openai_responses`` parses that list (behind the same >100
|
||||
saved-token gate as ``TransformPipeline.apply``) and threads the result
|
||||
into both the non-streaming ``RequestOutcome`` and
|
||||
``_stream_response(waste_signals=...)``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("fastapi")
|
||||
pytest.importorskip("httpx")
|
||||
|
||||
from headroom import OpenAIProvider, Tokenizer
|
||||
from headroom.parser import parse_messages
|
||||
from headroom.proxy.handlers.openai import (
|
||||
_RESPONSES_OUTPUT_ITEM_TYPES,
|
||||
OpenAIHandlerMixin,
|
||||
_responses_input_to_waste_messages,
|
||||
_responses_part_text,
|
||||
)
|
||||
|
||||
_provider = OpenAIProvider()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def tokenizer() -> Tokenizer:
|
||||
return Tokenizer(_provider.get_token_counter("gpt-4o"), "gpt-4o")
|
||||
|
||||
|
||||
def _big_output(rows: int = 200) -> str:
|
||||
return json.dumps(
|
||||
[{"id": i, "name": f"item_{i}", "status": "ok", "score": i * 3.14} for i in range(rows)]
|
||||
)
|
||||
|
||||
|
||||
def _fco(output: object, call_id: str = "call_1") -> dict:
|
||||
return {"type": "function_call_output", "call_id": call_id, "output": output}
|
||||
|
||||
|
||||
class TestResponsesPartText:
|
||||
def test_string_passthrough(self):
|
||||
assert _responses_part_text("plain") == "plain"
|
||||
|
||||
def test_part_list_joined(self):
|
||||
parts = [
|
||||
{"type": "output_text", "text": "first"},
|
||||
"second",
|
||||
{"type": "input_text", "text": "third"},
|
||||
{"type": "input_image", "image_url": "ignored"},
|
||||
]
|
||||
assert _responses_part_text(parts) == "first\nsecond\nthird"
|
||||
|
||||
def test_non_text_returns_empty(self):
|
||||
assert _responses_part_text(None) == ""
|
||||
assert _responses_part_text({"text": "not a list"}) == ""
|
||||
|
||||
|
||||
class TestResponsesWasteConversion:
|
||||
def test_string_input_and_instructions(self):
|
||||
messages = _responses_input_to_waste_messages("be terse", "hello")
|
||||
assert messages == [
|
||||
{"role": "system", "content": "be terse"},
|
||||
{"role": "user", "content": "hello"},
|
||||
]
|
||||
|
||||
def test_message_items_keep_role(self):
|
||||
items = [
|
||||
{"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hi"}]},
|
||||
{
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": "hello"}],
|
||||
},
|
||||
]
|
||||
messages = _responses_input_to_waste_messages(None, items)
|
||||
assert messages == [
|
||||
{"role": "user", "content": "hi"},
|
||||
{"role": "assistant", "content": "hello"},
|
||||
]
|
||||
|
||||
def test_function_call_output_becomes_tool_message(self):
|
||||
output = _big_output()
|
||||
messages = _responses_input_to_waste_messages(None, [_fco(output)])
|
||||
assert messages == [{"role": "tool", "content": output, "tool_call_id": "call_1"}]
|
||||
|
||||
def test_output_part_list_joined(self):
|
||||
messages = _responses_input_to_waste_messages(
|
||||
None,
|
||||
[_fco([{"type": "output_text", "text": "a"}, {"type": "output_text", "text": "b"}])],
|
||||
)
|
||||
assert messages[0]["content"] == "a\nb"
|
||||
|
||||
def test_all_output_item_types_covered(self):
|
||||
for item_type in _RESPONSES_OUTPUT_ITEM_TYPES:
|
||||
messages = _responses_input_to_waste_messages(
|
||||
None, [{"type": item_type, "output": "tool output text"}]
|
||||
)
|
||||
assert messages == [{"role": "tool", "content": "tool output text"}], item_type
|
||||
|
||||
def test_skips_unusable_items(self):
|
||||
items = [
|
||||
"not a dict",
|
||||
{"type": "function_call", "name": "f", "arguments": "{}"},
|
||||
{"type": "function_call_output", "call_id": "c", "output": ""},
|
||||
{"type": "message", "role": "user", "content": []},
|
||||
]
|
||||
assert _responses_input_to_waste_messages(None, items) == []
|
||||
|
||||
def test_non_list_non_string_input(self):
|
||||
assert _responses_input_to_waste_messages(None, {"weird": True}) == []
|
||||
|
||||
def test_class_attr_aliases_module_constant(self):
|
||||
assert OpenAIHandlerMixin.OPENAI_RESPONSES_OUTPUT_TYPES is _RESPONSES_OUTPUT_ITEM_TYPES
|
||||
|
||||
|
||||
class TestResponsesWasteParsing:
|
||||
def test_tool_output_reaches_waste_signals(self, tokenizer):
|
||||
items = [
|
||||
{"type": "message", "role": "user", "content": [{"type": "input_text", "text": "go"}]},
|
||||
_fco(_big_output()),
|
||||
]
|
||||
messages = _responses_input_to_waste_messages("be terse", items)
|
||||
blocks, _, waste = parse_messages(messages, tokenizer)
|
||||
assert any(b.kind == "tool_result" for b in blocks)
|
||||
assert waste.json_bloat_tokens > 0
|
||||
|
||||
def test_repeated_tool_output_counts_as_reread(self, tokenizer):
|
||||
output = _big_output()
|
||||
filler = [
|
||||
{
|
||||
"type": "message",
|
||||
"role": "user",
|
||||
"content": [{"type": "input_text", "text": f"step {i}"}],
|
||||
}
|
||||
for i in range(5)
|
||||
]
|
||||
items = [_fco(output, "call_1"), *filler, _fco(output, "call_2")]
|
||||
messages = _responses_input_to_waste_messages(None, items)
|
||||
_, _, waste = parse_messages(messages, tokenizer)
|
||||
assert waste.reread_tokens > 0
|
||||
|
|
@ -169,25 +169,29 @@ class TestNetCostFormula:
|
|||
"""
|
||||
|
||||
def test_small_shave_deep_suffix_is_loss(self):
|
||||
# 2000*(1.25 + 0.1*9) - 1.0*1.15*50000 = 4300 - 57500 = -53200.
|
||||
# 2000*(1.25 + 0.1*9) - 1.0*1.15*52000 = 4300 - 59800 = -55500.
|
||||
p = policy_for_mode(AuthMode.PAYG)
|
||||
gain = p.net_mutation_gain(2_000, 50_000, 10.0, 1.0)
|
||||
assert abs(gain - (-53_200.0)) < 1.0
|
||||
assert abs(gain - (-55_500.0)) < 1.0
|
||||
assert not p.should_mutate_deep(2_000, 50_000, 10.0, 1.0)
|
||||
|
||||
def test_big_shave_shallow_suffix_is_win(self):
|
||||
# 50000*(1.25 + 0.1*2) - 1.0*1.15*10000 = 72500 - 11500 = 61000.
|
||||
# 50000*(1.25 + 0.1*2) - 1.0*1.15*60000 = 72500 - 69000 = 3500.
|
||||
# Tight but positive — consistent with the 2.3-read break-even.
|
||||
p = policy_for_mode(AuthMode.PAYG)
|
||||
gain = p.net_mutation_gain(50_000, 10_000, 3.0, 1.0)
|
||||
assert abs(gain - 61_000.0) < 1.0
|
||||
assert abs(gain - 3_500.0) < 1.0
|
||||
assert p.should_mutate_deep(50_000, 10_000, 3.0, 1.0)
|
||||
|
||||
def test_live_zone_edit_always_profitable(self):
|
||||
# S = 0 derives the existing Subscription live-zone policy as a
|
||||
# special case of the formula.
|
||||
def test_no_suffix_edit_profitable_with_reads_remaining(self):
|
||||
# S = 0: warm-case saving is the avoided rereads, dT*r*R —
|
||||
# positive whenever at least one read remains. At R=0 with a
|
||||
# warm cache the gain is exactly 0 (already written, never read
|
||||
# again): pointless rather than harmful.
|
||||
p = policy_for_mode(AuthMode.SUBSCRIPTION)
|
||||
assert p.should_mutate_deep(1, 0, 0.0, 1.0)
|
||||
assert p.should_mutate_deep(2_000, 0, 0.0, 1.0)
|
||||
assert p.should_mutate_deep(1, 0, 1.0, 1.0)
|
||||
assert p.should_mutate_deep(2_000, 0, 1.0, 1.0)
|
||||
assert abs(p.net_mutation_gain(2_000, 0, 0.0, 1.0)) < 1e-6
|
||||
|
||||
def test_cold_cache_ignores_suffix(self):
|
||||
# P_alive = 0 (TTL lapsed): the idle-timer compaction window.
|
||||
|
|
@ -220,11 +224,11 @@ class TestNetCostFormula:
|
|||
assert p.net_mutation_gain(2_000, -1, 5.0, 1.0) == p.net_mutation_gain(2_000, 0, 5.0, 1.0)
|
||||
|
||||
def test_break_even_reads_matches_research_anchor(self):
|
||||
# R = 11.5*(S/dT - 1): 2K/50K -> 276; 50K/10K -> negative
|
||||
# (profitable from the first read); dT=0 -> 0.
|
||||
# R = 11.5*S/dT, the #856 anchors exactly: 2K/50K -> 287.5;
|
||||
# 50K/10K -> 2.3; dT=0 -> 0.
|
||||
p = policy_for_mode(AuthMode.PAYG)
|
||||
assert abs(p.break_even_reads(2_000, 50_000) - 276.0) < 0.5
|
||||
assert p.break_even_reads(50_000, 10_000) < 0.0
|
||||
assert abs(p.break_even_reads(2_000, 50_000) - 287.5) < 0.5
|
||||
assert abs(p.break_even_reads(50_000, 10_000) - 2.3) < 0.05
|
||||
assert p.break_even_reads(0, 10_000) == 0.0
|
||||
|
||||
def test_constants_match_rust(self):
|
||||
|
|
|
|||
|
|
@ -12,11 +12,49 @@ import pytest
|
|||
from headroom import copilot_auth
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _isolated_copilot_auth(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
|
||||
"""Keep Copilot auth tests away from user secret stores and real auth files."""
|
||||
|
||||
for var in (
|
||||
"GITHUB_COPILOT_API_TOKEN",
|
||||
"COPILOT_PROVIDER_BEARER_TOKEN",
|
||||
"GITHUB_COPILOT_GITHUB_TOKEN",
|
||||
"GITHUB_COPILOT_TOKEN",
|
||||
"COPILOT_GITHUB_TOKEN",
|
||||
"GH_TOKEN",
|
||||
"GITHUB_TOKEN",
|
||||
"GITHUB_COPILOT_API_URL",
|
||||
"GITHUB_COPILOT_ENTERPRISE_URL",
|
||||
"GITHUB_COPILOT_ENTERPRISE_DOMAIN",
|
||||
"GITHUB_COPILOT_TOKEN_EXCHANGE_URL",
|
||||
"GITHUB_COPILOT_USER_INFO_URL",
|
||||
"GITHUB_COPILOT_USER_AGENT",
|
||||
"GITHUB_COPILOT_EDITOR_VERSION",
|
||||
"GITHUB_COPILOT_EDITOR_PLUGIN_VERSION",
|
||||
"GITHUB_COPILOT_INTEGRATION_ID",
|
||||
):
|
||||
monkeypatch.delenv(var, raising=False)
|
||||
monkeypatch.setattr(copilot_auth, "_provider", None)
|
||||
monkeypatch.setenv("HEADROOM_COPILOT_AUTH_FILE", str(tmp_path / "copilot_auth.json"))
|
||||
monkeypatch.setattr(copilot_auth, "read_macos_keychain_token", lambda *, host: None)
|
||||
monkeypatch.setattr(copilot_auth, "read_linux_secret_token", lambda *, host: None)
|
||||
|
||||
|
||||
def test_read_cached_oauth_token_prefers_env(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("GITHUB_COPILOT_TOKEN", "gho-env")
|
||||
assert copilot_auth.read_cached_oauth_token() == "gho-env"
|
||||
|
||||
|
||||
def test_read_cached_oauth_token_prefers_headroom_login(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setenv("GITHUB_COPILOT_TOKEN", "gho-env")
|
||||
copilot_auth.save_headroom_copilot_oauth_token("gho-headroom")
|
||||
|
||||
assert copilot_auth.read_cached_oauth_token() == "gho-headroom"
|
||||
|
||||
|
||||
def test_read_cached_oauth_token_prefers_copilot_cli_before_generic_github_token(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
|
|
@ -56,6 +94,9 @@ def test_resolve_subscription_bearer_token_skips_invalid_generic_token(
|
|||
) -> None:
|
||||
monkeypatch.delenv("GITHUB_COPILOT_API_TOKEN", raising=False)
|
||||
monkeypatch.delenv("COPILOT_PROVIDER_BEARER_TOKEN", raising=False)
|
||||
monkeypatch.setattr(
|
||||
copilot_auth, "_subscription_resolution_from_token_exchange", lambda _: None
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
copilot_auth,
|
||||
"iter_oauth_token_candidates",
|
||||
|
|
@ -65,6 +106,38 @@ def test_resolve_subscription_bearer_token_skips_invalid_generic_token(
|
|||
source="env:GITHUB_TOKEN",
|
||||
confidence="generic-github",
|
||||
),
|
||||
copilot_auth.CopilotTokenCandidate(
|
||||
token="tid_copilot",
|
||||
source="macos-keychain:copilot-cli",
|
||||
confidence="high",
|
||||
),
|
||||
],
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
copilot_auth,
|
||||
"_fetch_copilot_user_info",
|
||||
lambda token: (
|
||||
{"endpoints": {"api": "https://api.individual.githubcopilot.com"}}
|
||||
if token == "tid_copilot"
|
||||
else None
|
||||
),
|
||||
)
|
||||
|
||||
assert copilot_auth.resolve_subscription_bearer_token() == "tid_copilot"
|
||||
|
||||
|
||||
def test_resolve_subscription_bearer_token_does_not_fallback_to_unexchanged_oauth(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.delenv("GITHUB_COPILOT_API_TOKEN", raising=False)
|
||||
monkeypatch.delenv("COPILOT_PROVIDER_BEARER_TOKEN", raising=False)
|
||||
monkeypatch.setattr(
|
||||
copilot_auth, "_subscription_resolution_from_token_exchange", lambda _: None
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
copilot_auth,
|
||||
"iter_oauth_token_candidates",
|
||||
lambda: [
|
||||
copilot_auth.CopilotTokenCandidate(
|
||||
token="gho-copilot",
|
||||
source="macos-keychain:copilot-cli",
|
||||
|
|
@ -82,7 +155,148 @@ def test_resolve_subscription_bearer_token_skips_invalid_generic_token(
|
|||
),
|
||||
)
|
||||
|
||||
assert copilot_auth.resolve_subscription_bearer_token() == "gho-copilot"
|
||||
assert copilot_auth.resolve_subscription_bearer_token() is None
|
||||
|
||||
|
||||
def test_resolve_subscription_bearer_token_details_exchanges_oauth_candidate(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.delenv("GITHUB_COPILOT_API_TOKEN", raising=False)
|
||||
monkeypatch.delenv("COPILOT_PROVIDER_BEARER_TOKEN", raising=False)
|
||||
monkeypatch.delenv("GITHUB_COPILOT_API_URL", raising=False)
|
||||
monkeypatch.delenv("GITHUB_COPILOT_ENTERPRISE_URL", raising=False)
|
||||
monkeypatch.delenv("GITHUB_COPILOT_ENTERPRISE_DOMAIN", raising=False)
|
||||
monkeypatch.delenv("GITHUB_COPILOT_TOKEN_EXCHANGE_URL", raising=False)
|
||||
monkeypatch.setattr(
|
||||
copilot_auth,
|
||||
"iter_oauth_token_candidates",
|
||||
lambda: [
|
||||
copilot_auth.CopilotTokenCandidate(
|
||||
token="gho-oauth",
|
||||
source="headroom-copilot-auth:/tmp/copilot_auth.json",
|
||||
confidence="copilot-oauth",
|
||||
),
|
||||
],
|
||||
)
|
||||
captured: dict[str, str] = {}
|
||||
|
||||
def fake_exchange(headers: dict[str, str]) -> dict[str, object]:
|
||||
captured.update(headers)
|
||||
return {
|
||||
"token": "copilot-api",
|
||||
"expires_at": int(time.time()) + 3600,
|
||||
"endpoints": {"api": "https://api.business.githubcopilot.com"},
|
||||
}
|
||||
|
||||
monkeypatch.setattr(
|
||||
copilot_auth.CopilotTokenProvider,
|
||||
"_exchange_token_sync",
|
||||
staticmethod(fake_exchange),
|
||||
)
|
||||
|
||||
resolution = copilot_auth.resolve_subscription_bearer_token_details()
|
||||
|
||||
assert resolution is not None
|
||||
assert resolution.token == "copilot-api"
|
||||
assert resolution.source == "headroom-copilot-auth:/tmp/copilot_auth.json:token-exchange"
|
||||
assert resolution.confidence == "copilot-token-exchange"
|
||||
assert resolution.api_url == "https://api.business.githubcopilot.com"
|
||||
assert resolution.token_fingerprint == copilot_auth.token_fingerprint("copilot-api")
|
||||
assert captured == {
|
||||
"Accept": "application/json",
|
||||
"Authorization": "Bearer gho-oauth",
|
||||
"User-Agent": "GitHubCopilotChat/0.35.0",
|
||||
"Editor-Version": "vscode/1.107.0",
|
||||
"Editor-Plugin-Version": "copilot-chat/0.35.0",
|
||||
"Copilot-Integration-Id": "vscode-chat",
|
||||
}
|
||||
|
||||
|
||||
def test_resolve_subscription_exchange_uses_cloud_enterprise_advertised_api(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.delenv("GITHUB_COPILOT_API_TOKEN", raising=False)
|
||||
monkeypatch.delenv("COPILOT_PROVIDER_BEARER_TOKEN", raising=False)
|
||||
monkeypatch.delenv("GITHUB_COPILOT_API_URL", raising=False)
|
||||
monkeypatch.delenv("GITHUB_COPILOT_ENTERPRISE_DOMAIN", raising=False)
|
||||
monkeypatch.delenv("GITHUB_COPILOT_TOKEN_EXCHANGE_URL", raising=False)
|
||||
monkeypatch.setenv("GITHUB_COPILOT_ENTERPRISE_URL", "github.com/enterprises/cbcrc")
|
||||
monkeypatch.setattr(
|
||||
copilot_auth,
|
||||
"iter_oauth_token_candidates",
|
||||
lambda: [
|
||||
copilot_auth.CopilotTokenCandidate(
|
||||
token="gho-oauth",
|
||||
source="env:GITHUB_COPILOT_TOKEN",
|
||||
confidence="explicit",
|
||||
),
|
||||
],
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
copilot_auth.CopilotTokenProvider,
|
||||
"_exchange_token_sync",
|
||||
staticmethod(lambda _headers: {"token": "copilot-api"}),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
copilot_auth,
|
||||
"_fetch_copilot_user_info",
|
||||
lambda _token: {"endpoints": {"api": "https://api.business.githubcopilot.com"}},
|
||||
)
|
||||
|
||||
resolution = copilot_auth.resolve_subscription_bearer_token_details()
|
||||
|
||||
assert resolution is not None
|
||||
assert resolution.api_url == "https://api.business.githubcopilot.com"
|
||||
assert copilot_auth._token_exchange_url() == "https://api.github.com/copilot_internal/v2/token"
|
||||
|
||||
|
||||
def test_enterprise_domain_routes_token_exchange_and_user_info_together(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.delenv("GITHUB_COPILOT_TOKEN_EXCHANGE_URL", raising=False)
|
||||
monkeypatch.delenv("GITHUB_COPILOT_USER_INFO_URL", raising=False)
|
||||
monkeypatch.delenv("GITHUB_COPILOT_ENTERPRISE_URL", raising=False)
|
||||
monkeypatch.setenv("GITHUB_COPILOT_ENTERPRISE_DOMAIN", "ghe.example.com")
|
||||
|
||||
assert (
|
||||
copilot_auth._token_exchange_url()
|
||||
== "https://api.ghe.example.com/copilot_internal/v2/token"
|
||||
)
|
||||
assert copilot_auth._user_info_url() == "https://api.ghe.example.com/copilot_internal/user"
|
||||
|
||||
|
||||
def test_user_info_url_override_wins_over_enterprise_domain(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setenv("GITHUB_COPILOT_ENTERPRISE_DOMAIN", "ghe.example.com")
|
||||
monkeypatch.setenv(
|
||||
"GITHUB_COPILOT_USER_INFO_URL",
|
||||
"https://custom.example.com/copilot_internal/user",
|
||||
)
|
||||
|
||||
assert copilot_auth._user_info_url() == "https://custom.example.com/copilot_internal/user"
|
||||
|
||||
|
||||
def test_copilot_api_url_from_enterprise_url_supports_enterprise_server_domain() -> None:
|
||||
assert (
|
||||
copilot_auth.copilot_api_url_from_enterprise_url("https://ghe.example.com/")
|
||||
== "https://copilot-api.ghe.example.com"
|
||||
)
|
||||
assert (
|
||||
copilot_auth.copilot_api_url_from_enterprise_url("https://api.ghe.example.com/")
|
||||
== "https://copilot-api.ghe.example.com"
|
||||
)
|
||||
assert (
|
||||
copilot_auth.copilot_api_url_from_enterprise_url("https://copilot-api.ghe.example.com/")
|
||||
== "https://copilot-api.ghe.example.com"
|
||||
)
|
||||
|
||||
|
||||
def test_copilot_api_url_from_enterprise_url_ignores_github_cloud_enterprise_path() -> None:
|
||||
assert (
|
||||
copilot_auth.copilot_api_url_from_enterprise_url("https://github.com/enterprises/cbcrc/")
|
||||
== copilot_auth.DEFAULT_API_URL
|
||||
)
|
||||
|
||||
|
||||
def test_should_exchange_oauth_token_supports_truthy_values(
|
||||
|
|
@ -292,6 +506,22 @@ def test_is_copilot_api_url_matches_expected_hosts() -> None:
|
|||
assert not copilot_auth.is_copilot_api_url("https://api.openai.com/v1/chat/completions")
|
||||
|
||||
|
||||
def test_is_copilot_api_url_matches_ghe_copilot_hosts() -> None:
|
||||
assert copilot_auth.is_copilot_api_url("https://copilot-api.acme.ghe.com/v1/responses")
|
||||
assert copilot_auth.is_copilot_api_url("https://copilot-api.ghe.com/v1/chat/completions")
|
||||
assert not copilot_auth.is_copilot_api_url("https://api.acme.ghe.com/v1/responses")
|
||||
assert not copilot_auth.is_copilot_api_url("https://not-copilot-api.acme.ghe.com/v1/responses")
|
||||
|
||||
|
||||
def test_is_copilot_api_url_trusts_configured_enterprise_api_url(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setenv("GITHUB_COPILOT_ENTERPRISE_DOMAIN", "ghe.example.com")
|
||||
|
||||
assert copilot_auth.is_copilot_api_url("https://copilot-api.ghe.example.com/v1/responses")
|
||||
assert not copilot_auth.is_copilot_api_url("https://copilot-api.other.example.com/v1/responses")
|
||||
|
||||
|
||||
def test_build_copilot_upstream_url_strips_v1_only_for_copilot_hosts() -> None:
|
||||
assert (
|
||||
copilot_auth.build_copilot_upstream_url(
|
||||
|
|
@ -309,6 +539,37 @@ def test_build_copilot_upstream_url_strips_v1_only_for_copilot_hosts() -> None:
|
|||
)
|
||||
|
||||
|
||||
def test_build_copilot_upstream_url_strips_v1_for_ghe_copilot_hosts() -> None:
|
||||
assert (
|
||||
copilot_auth.build_copilot_upstream_url(
|
||||
"https://copilot-api.acme.ghe.com",
|
||||
"/v1/responses",
|
||||
)
|
||||
== "https://copilot-api.acme.ghe.com/responses"
|
||||
)
|
||||
assert (
|
||||
copilot_auth.build_copilot_upstream_url(
|
||||
"https://api.acme.ghe.com",
|
||||
"/v1/responses",
|
||||
)
|
||||
== "https://api.acme.ghe.com/v1/responses"
|
||||
)
|
||||
|
||||
|
||||
def test_build_copilot_upstream_url_strips_v1_for_configured_enterprise_api_url(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setenv("GITHUB_COPILOT_ENTERPRISE_DOMAIN", "ghe.example.com")
|
||||
|
||||
assert (
|
||||
copilot_auth.build_copilot_upstream_url(
|
||||
"https://copilot-api.ghe.example.com",
|
||||
"/v1/responses",
|
||||
)
|
||||
== "https://copilot-api.ghe.example.com/responses"
|
||||
)
|
||||
|
||||
|
||||
def test_apply_copilot_api_auth_replaces_authorization(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
async def fake_get_api_token() -> copilot_auth.CopilotAPIToken:
|
||||
return copilot_auth.CopilotAPIToken(
|
||||
|
|
@ -325,13 +586,18 @@ def test_apply_copilot_api_auth_replaces_authorization(monkeypatch: pytest.Monke
|
|||
|
||||
headers = asyncio.run(
|
||||
copilot_auth.apply_copilot_api_auth(
|
||||
{"authorization": "Bearer downstream-token"},
|
||||
{"authorization": "Bearer downstream-token", "x-api-key": "sk-downstream"},
|
||||
url="https://api.githubcopilot.com/v1/chat/completions",
|
||||
)
|
||||
)
|
||||
|
||||
assert headers["Authorization"] == "Bearer copilot-session"
|
||||
assert "authorization" not in headers
|
||||
assert "x-api-key" not in headers
|
||||
assert headers["User-Agent"] == "GitHubCopilotChat/0.35.0"
|
||||
assert headers["Editor-Version"] == "vscode/1.107.0"
|
||||
assert headers["Editor-Plugin-Version"] == "copilot-chat/0.35.0"
|
||||
assert headers["Copilot-Integration-Id"] == "vscode-chat"
|
||||
|
||||
|
||||
def test_apply_copilot_api_auth_passes_through_existing_api_token(
|
||||
|
|
@ -348,12 +614,16 @@ def test_apply_copilot_api_auth_passes_through_existing_api_token(
|
|||
|
||||
headers = asyncio.run(
|
||||
copilot_auth.apply_copilot_api_auth(
|
||||
{"authorization": "Bearer tid_existing_copilot_token"},
|
||||
{
|
||||
"authorization": "Bearer tid_existing_copilot_token",
|
||||
"x-api-key": "sk-downstream",
|
||||
},
|
||||
url="https://api.githubcopilot.com/v1/chat/completions",
|
||||
)
|
||||
)
|
||||
|
||||
assert headers["authorization"] == "Bearer tid_existing_copilot_token"
|
||||
assert "x-api-key" not in headers
|
||||
|
||||
|
||||
def test_apply_copilot_api_auth_replaces_github_oauth_bearer(
|
||||
|
|
@ -446,7 +716,8 @@ def test_apply_copilot_api_auth_injects_required_headers(
|
|||
|
||||
assert headers["Authorization"] == "Bearer copilot-session"
|
||||
assert headers["Copilot-Integration-Id"] == "vscode-chat"
|
||||
assert headers["editor-version"] == "vscode/1.104.1"
|
||||
assert headers["Editor-Version"] == "vscode/1.107.0"
|
||||
assert headers["Editor-Plugin-Version"] == "copilot-chat/0.35.0"
|
||||
|
||||
|
||||
def test_apply_copilot_api_auth_preserves_existing_copilot_headers(
|
||||
|
|
@ -483,6 +754,47 @@ def test_apply_copilot_api_auth_preserves_existing_copilot_headers(
|
|||
assert headers["Authorization"] == "Bearer copilot-session"
|
||||
|
||||
|
||||
def test_apply_copilot_api_auth_preserves_existing_headers_case_insensitively(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
async def fake_get_api_token() -> copilot_auth.CopilotAPIToken:
|
||||
return copilot_auth.CopilotAPIToken(
|
||||
token="copilot-session",
|
||||
expires_at=time.time() + 3600,
|
||||
api_url=copilot_auth.DEFAULT_API_URL,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
copilot_auth.get_copilot_token_provider(),
|
||||
"get_api_token",
|
||||
fake_get_api_token,
|
||||
)
|
||||
|
||||
headers = asyncio.run(
|
||||
copilot_auth.apply_copilot_api_auth(
|
||||
{
|
||||
"authorization": "Bearer downstream-token",
|
||||
"user-agent": "custom-agent",
|
||||
"editor-version": "custom-editor",
|
||||
"editor-plugin-version": "custom-plugin",
|
||||
"copilot-integration-id": "custom-integration",
|
||||
},
|
||||
url="https://api.githubcopilot.com/v1/chat/completions",
|
||||
)
|
||||
)
|
||||
|
||||
assert headers["Authorization"] == "Bearer copilot-session"
|
||||
assert "authorization" not in headers
|
||||
assert headers["user-agent"] == "custom-agent"
|
||||
assert headers["editor-version"] == "custom-editor"
|
||||
assert headers["editor-plugin-version"] == "custom-plugin"
|
||||
assert headers["copilot-integration-id"] == "custom-integration"
|
||||
assert "User-Agent" not in headers
|
||||
assert "Editor-Version" not in headers
|
||||
assert "Editor-Plugin-Version" not in headers
|
||||
assert "Copilot-Integration-Id" not in headers
|
||||
|
||||
|
||||
def test_token_provider_reuses_oauth_token_without_exchange(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
|
|
@ -517,9 +829,11 @@ def test_token_provider_can_exchange_when_enabled(monkeypatch: pytest.MonkeyPatc
|
|||
|
||||
provider = copilot_auth.CopilotTokenProvider()
|
||||
calls = {"count": 0}
|
||||
captured: dict[str, str] = {}
|
||||
|
||||
def fake_exchange(headers: dict[str, str]) -> dict[str, object]:
|
||||
calls["count"] += 1
|
||||
captured.update(headers)
|
||||
return {
|
||||
"token": "copilot-api",
|
||||
"expires_at": int(time.time()) + 3600,
|
||||
|
|
@ -536,6 +850,11 @@ def test_token_provider_can_exchange_when_enabled(monkeypatch: pytest.MonkeyPatc
|
|||
assert first.token == "copilot-api"
|
||||
assert second.token == "copilot-api"
|
||||
assert calls["count"] == 1
|
||||
assert captured["Authorization"] == "Bearer gho-oauth"
|
||||
assert captured["User-Agent"] == "GitHubCopilotChat/0.35.0"
|
||||
assert captured["Editor-Version"] == "vscode/1.107.0"
|
||||
assert captured["Editor-Plugin-Version"] == "copilot-chat/0.35.0"
|
||||
assert captured["Copilot-Integration-Id"] == "vscode-chat"
|
||||
|
||||
|
||||
def test_token_provider_prefers_explicit_api_token(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
|
|
|
|||
|
|
@ -5,10 +5,11 @@ The subscription flow has to behave identically on macOS, Linux, and Windows
|
|||
Copilot CLI token from the platform secret store — is impossible to exercise
|
||||
portably. This suite proves the *portable* contract instead:
|
||||
|
||||
1. With an explicit token in the environment, resolution + API-URL discovery
|
||||
succeed on every platform without touching any secret store. This is the
|
||||
universal escape hatch (``GITHUB_COPILOT_TOKEN`` etc.) that makes the
|
||||
feature work anywhere, including headless CI.
|
||||
1. With an explicit Copilot API token in the environment, resolution + API-URL
|
||||
discovery succeed on every platform without touching any secret store. This
|
||||
is the deterministic escape hatch (``GITHUB_COPILOT_API_TOKEN``) for
|
||||
headless CI. OAuth tokens still need successful token exchange before
|
||||
subscription mode can use them.
|
||||
2. Each OS-specific secret reader is inert on a foreign platform — so on any
|
||||
given OS only that OS's reader can fire, and a missing/foreign secret store
|
||||
degrades to ``None`` rather than crashing.
|
||||
|
|
@ -33,6 +34,7 @@ BUSINESS_API = "https://api.business.githubcopilot.com"
|
|||
|
||||
def _stub_all_secret_stores(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Simulate 'no OS secret store / not logged in' on every platform."""
|
||||
monkeypatch.setattr(copilot_auth, "read_headroom_copilot_oauth_token", lambda: None)
|
||||
monkeypatch.setattr(copilot_auth, "_read_windows_copilot_cli_oauth_token", lambda: None)
|
||||
monkeypatch.setattr(copilot_auth, "_read_macos_keychain_oauth_token", lambda: None)
|
||||
monkeypatch.setattr(copilot_auth, "_read_linux_secret_oauth_token", lambda: None)
|
||||
|
|
@ -50,28 +52,31 @@ def _clear_token_env(monkeypatch: pytest.MonkeyPatch) -> None:
|
|||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 1. The env-var path resolves on any platform with no secret store.
|
||||
# 1. The explicit API-token env path resolves on any platform with no secret store.
|
||||
# ---------------------------------------------------------------------------
|
||||
def test_env_token_resolves_subscription_without_secret_store(
|
||||
def test_api_token_env_resolves_subscription_without_secret_store(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
_stub_all_secret_stores(monkeypatch)
|
||||
_clear_token_env(monkeypatch)
|
||||
monkeypatch.setenv("GITHUB_COPILOT_TOKEN", "gho-env-universal")
|
||||
monkeypatch.setenv("GITHUB_COPILOT_API_TOKEN", "tid_env_universal")
|
||||
monkeypatch.setattr(
|
||||
copilot_auth, "_subscription_resolution_from_token_exchange", lambda _: None
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
copilot_auth,
|
||||
"_fetch_copilot_user_info",
|
||||
lambda token: (
|
||||
{"endpoints": {"api": BUSINESS_API}} if token == "gho-env-universal" else None
|
||||
{"endpoints": {"api": BUSINESS_API}} if token == "tid_env_universal" else None
|
||||
),
|
||||
)
|
||||
|
||||
assert copilot_auth.resolve_subscription_bearer_token() == "gho-env-universal"
|
||||
assert copilot_auth.resolve_subscription_bearer_token() == "tid_env_universal"
|
||||
# Routing is override -> generic; the account host advertised by user-info is
|
||||
# NOT used (it regressed newer models on the responses API, #610). With no
|
||||
# GITHUB_COPILOT_API_URL pin set, the generic public host is returned.
|
||||
monkeypatch.delenv("GITHUB_COPILOT_API_URL", raising=False)
|
||||
assert copilot_auth.resolve_copilot_api_url("gho-env-universal") == copilot_auth.DEFAULT_API_URL
|
||||
assert copilot_auth.resolve_copilot_api_url("tid_env_universal") == copilot_auth.DEFAULT_API_URL
|
||||
|
||||
|
||||
def test_api_url_falls_back_to_default_when_user_info_unavailable(
|
||||
|
|
@ -85,13 +90,16 @@ def test_api_url_falls_back_to_default_when_user_info_unavailable(
|
|||
assert copilot_auth.resolve_copilot_api_url("gho-anything") == copilot_auth.DEFAULT_API_URL
|
||||
|
||||
|
||||
def test_subscription_rejects_token_github_does_not_accept(
|
||||
def test_subscription_rejects_generic_token_and_accepts_api_token(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
_stub_all_secret_stores(monkeypatch)
|
||||
_clear_token_env(monkeypatch)
|
||||
# A generic GitHub token is present but GitHub's Copilot API rejects it;
|
||||
# a valid Copilot token is discoverable behind it.
|
||||
monkeypatch.setattr(
|
||||
copilot_auth, "_subscription_resolution_from_token_exchange", lambda _: None
|
||||
)
|
||||
# A generic GitHub token is present but cannot be exchanged for a Copilot
|
||||
# API token; a valid Copilot API token is discoverable behind it.
|
||||
monkeypatch.setattr(
|
||||
copilot_auth,
|
||||
"iter_oauth_token_candidates",
|
||||
|
|
@ -100,7 +108,7 @@ def test_subscription_rejects_token_github_does_not_accept(
|
|||
token="ghp-generic-pat", source="env:GITHUB_TOKEN", confidence="generic-github"
|
||||
),
|
||||
copilot_auth.CopilotTokenCandidate(
|
||||
token="gho-real-copilot",
|
||||
token="tid_real_copilot",
|
||||
source="macos-keychain:copilot-cli",
|
||||
confidence="high",
|
||||
),
|
||||
|
|
@ -109,10 +117,10 @@ def test_subscription_rejects_token_github_does_not_accept(
|
|||
monkeypatch.setattr(
|
||||
copilot_auth,
|
||||
"_fetch_copilot_user_info",
|
||||
lambda token: {"endpoints": {"api": BUSINESS_API}} if token == "gho-real-copilot" else None,
|
||||
lambda token: {"endpoints": {"api": BUSINESS_API}} if token == "tid_real_copilot" else None,
|
||||
)
|
||||
|
||||
assert copilot_auth.resolve_subscription_bearer_token() == "gho-real-copilot"
|
||||
assert copilot_auth.resolve_subscription_bearer_token() == "tid_real_copilot"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -182,6 +190,16 @@ def test_end_to_end_subscription_chain(monkeypatch: pytest.MonkeyPatch) -> None:
|
|||
_clear_token_env(monkeypatch)
|
||||
monkeypatch.setenv("GITHUB_COPILOT_TOKEN", "gho-seat-token")
|
||||
monkeypatch.setenv("GITHUB_COPILOT_API_URL", BUSINESS_API)
|
||||
monkeypatch.setattr(
|
||||
copilot_auth,
|
||||
"_subscription_resolution_from_token_exchange",
|
||||
lambda _candidate: copilot_auth._subscription_resolution(
|
||||
token="tid-seat-token",
|
||||
source="env:GITHUB_COPILOT_TOKEN:token-exchange",
|
||||
confidence="copilot-token-exchange",
|
||||
api_url=BUSINESS_API,
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
copilot_auth,
|
||||
"_fetch_copilot_user_info",
|
||||
|
|
@ -193,7 +211,7 @@ def test_end_to_end_subscription_chain(monkeypatch: pytest.MonkeyPatch) -> None:
|
|||
)
|
||||
resolved_token = copilot_auth.resolve_subscription_bearer_token()
|
||||
resolved_url = copilot_auth.resolve_copilot_api_url(resolved_token)
|
||||
assert resolved_token == "gho-seat-token"
|
||||
assert resolved_token == "tid-seat-token"
|
||||
assert resolved_url == BUSINESS_API # the pin wins; the user-info host is ignored
|
||||
|
||||
# (b) hand-off: the wrapper exports exactly these for the proxy.
|
||||
|
|
|
|||
276
tests/test_dashboard_agent_usage.py
Normal file
276
tests/test_dashboard_agent_usage.py
Normal file
|
|
@ -0,0 +1,276 @@
|
|||
from headroom.proxy.server import (
|
||||
_agent_label,
|
||||
_build_agent_usage_summary,
|
||||
_classify_agent_from_log,
|
||||
_normalize_agent_key,
|
||||
)
|
||||
|
||||
|
||||
def test_agent_usage_groups_exact_logged_requests_by_client() -> None:
|
||||
summary = _build_agent_usage_summary(
|
||||
[
|
||||
{
|
||||
"provider": "openai",
|
||||
"model": "gpt-5.2-codex",
|
||||
"tags": {"client": "codex"},
|
||||
"input_tokens_original": 1000,
|
||||
"input_tokens_optimized": 650,
|
||||
"output_tokens": 100,
|
||||
"tokens_saved": 350,
|
||||
},
|
||||
{
|
||||
"provider": "anthropic",
|
||||
"model": "claude-sonnet-4-6",
|
||||
"tags": {"client": "claude-code"},
|
||||
"input_tokens_original": 800,
|
||||
"input_tokens_optimized": 500,
|
||||
"output_tokens": 80,
|
||||
"tokens_saved": 300,
|
||||
},
|
||||
{
|
||||
"provider": "anthropic",
|
||||
"model": "claude-sonnet-4-6",
|
||||
"tags": {"client": "cursor"},
|
||||
"input_tokens_original": 500,
|
||||
"input_tokens_optimized": 400,
|
||||
"output_tokens": 60,
|
||||
"tokens_saved": 100,
|
||||
},
|
||||
],
|
||||
requests_by_provider={},
|
||||
requests_by_model={},
|
||||
global_before_tokens=2300,
|
||||
global_after_tokens=1550,
|
||||
global_tokens_saved=750,
|
||||
global_output_tokens=240,
|
||||
)
|
||||
|
||||
rows = {row["agent"]: row for row in summary["agents"]}
|
||||
|
||||
assert rows["codex"]["label"] == "Codex"
|
||||
assert rows["codex"]["before_tokens"] == 1000
|
||||
assert rows["codex"]["after_tokens"] == 650
|
||||
assert rows["codex"]["tokens_saved"] == 350
|
||||
assert rows["codex"]["savings_percent"] == 35.0
|
||||
|
||||
assert rows["claude-code"]["label"] == "Claude"
|
||||
assert rows["claude-code"]["savings_percent"] == 37.5
|
||||
|
||||
assert rows["cursor"]["label"] == "Cursor"
|
||||
assert rows["cursor"]["share_of_saved_percent"] == 13.33
|
||||
|
||||
assert summary["coverage"] == {
|
||||
"logged_requests": 3,
|
||||
"exact_token_rows": 3,
|
||||
"mode": "request_logs",
|
||||
}
|
||||
|
||||
|
||||
def test_agent_usage_falls_back_to_inferred_model_counts_when_complete() -> None:
|
||||
summary = _build_agent_usage_summary(
|
||||
[],
|
||||
requests_by_provider={"anthropic": 2, "openai": 3},
|
||||
requests_by_model={"claude-sonnet-4-6": 2, "gpt-5.2-codex": 3},
|
||||
global_before_tokens=1000,
|
||||
global_after_tokens=700,
|
||||
global_tokens_saved=300,
|
||||
global_output_tokens=90,
|
||||
)
|
||||
|
||||
rows = {row["agent"]: row for row in summary["agents"]}
|
||||
|
||||
assert set(rows) == {"claude-code", "codex"}
|
||||
assert rows["claude-code"]["label"] == "Claude"
|
||||
assert rows["claude-code"]["source"] == "model"
|
||||
assert rows["claude-code"]["requests"] == 2
|
||||
assert rows["claude-code"]["models"] == {"claude-sonnet-4-6": 2}
|
||||
assert rows["codex"]["label"] == "Codex"
|
||||
assert rows["codex"]["requests"] == 3
|
||||
assert rows["codex"]["models"] == {"gpt-5.2-codex": 3}
|
||||
assert summary["totals"]["savings_percent"] == 30.0
|
||||
assert summary["coverage"]["mode"] == "aggregate_fallback"
|
||||
|
||||
|
||||
def test_agent_usage_fallback_does_not_duplicate_provider_and_model_rows() -> None:
|
||||
summary = _build_agent_usage_summary(
|
||||
[],
|
||||
requests_by_provider={"anthropic": 2, "openai": 3},
|
||||
requests_by_model={"claude-sonnet-4-6": 2, "gpt-5.2-codex": 3},
|
||||
global_before_tokens=1000,
|
||||
global_after_tokens=700,
|
||||
global_tokens_saved=300,
|
||||
global_output_tokens=90,
|
||||
)
|
||||
|
||||
rows = {row["agent"]: row for row in summary["agents"]}
|
||||
|
||||
assert set(rows) == {"claude-code", "codex"}
|
||||
assert all(row["requests"] > 0 for row in rows.values())
|
||||
assert summary["totals"]["requests"] == 5
|
||||
|
||||
|
||||
def test_agent_usage_skips_partial_model_fallback_counts() -> None:
|
||||
summary = _build_agent_usage_summary(
|
||||
[],
|
||||
requests_by_provider={"anthropic": 2, "openai": 3},
|
||||
requests_by_model={"claude-sonnet-4-6": 2},
|
||||
global_before_tokens=1000,
|
||||
global_after_tokens=700,
|
||||
global_tokens_saved=300,
|
||||
global_output_tokens=90,
|
||||
)
|
||||
|
||||
rows = {row["agent"]: row for row in summary["agents"]}
|
||||
|
||||
assert set(rows) == {"anthropic", "openai"}
|
||||
assert rows["anthropic"]["label"] == "Claude"
|
||||
assert rows["anthropic"]["requests"] == 2
|
||||
assert rows["openai"]["label"] == "OpenAI"
|
||||
assert rows["openai"]["requests"] == 3
|
||||
|
||||
|
||||
def test_agent_classifier_uses_model_before_generic_provider() -> None:
|
||||
agent, label, source = _classify_agent_from_log(
|
||||
{
|
||||
"provider": "openai",
|
||||
"model": "gpt-5.2-codex",
|
||||
"tags": {},
|
||||
}
|
||||
)
|
||||
|
||||
assert (agent, label, source) == ("codex", "Codex", "model")
|
||||
|
||||
|
||||
def test_agent_usage_upgrades_source_when_stronger_evidence_arrives() -> None:
|
||||
summary = _build_agent_usage_summary(
|
||||
[
|
||||
{
|
||||
"provider": "anthropic",
|
||||
"model": "claude-sonnet-4-6",
|
||||
"tags": {},
|
||||
"input_tokens_original": 10,
|
||||
"input_tokens_optimized": 8,
|
||||
"tokens_saved": 2,
|
||||
},
|
||||
{
|
||||
"provider": "anthropic",
|
||||
"model": "claude-sonnet-4-6",
|
||||
"tags": {"client": "claude-code"},
|
||||
"input_tokens_original": 20,
|
||||
"input_tokens_optimized": 12,
|
||||
"tokens_saved": 8,
|
||||
},
|
||||
],
|
||||
requests_by_provider={},
|
||||
requests_by_model={},
|
||||
global_before_tokens=30,
|
||||
global_after_tokens=20,
|
||||
global_tokens_saved=10,
|
||||
global_output_tokens=0,
|
||||
)
|
||||
|
||||
row = summary["agents"][0]
|
||||
|
||||
assert row["agent"] == "claude-code"
|
||||
assert row["source"] == "client"
|
||||
assert row["requests"] == 2
|
||||
|
||||
|
||||
def test_agent_key_normalizes_wrapped_underscore_clients() -> None:
|
||||
assert _normalize_agent_key("wrap_claude_cli") == "claude-code"
|
||||
|
||||
|
||||
def test_agent_key_normalizes_claude_code_cli_alias() -> None:
|
||||
assert _normalize_agent_key("claude-code-cli") == "claude-code"
|
||||
|
||||
|
||||
def test_agent_label_title_cases_unknown_agent_key() -> None:
|
||||
assert _agent_label("custom-agent") == "Custom Agent"
|
||||
|
||||
|
||||
def test_agent_classifier_uses_stack_tag_before_model() -> None:
|
||||
agent, label, source = _classify_agent_from_log(
|
||||
{
|
||||
"provider": "openai",
|
||||
"model": "gpt-5.2-codex",
|
||||
"tags": {"headroom-stack": "openclaw"},
|
||||
}
|
||||
)
|
||||
|
||||
assert (agent, label, source) == ("openclaw", "OpenClaw", "stack")
|
||||
|
||||
|
||||
def test_agent_classifier_falls_back_to_unknown() -> None:
|
||||
agent, label, source = _classify_agent_from_log(
|
||||
{
|
||||
"provider": "",
|
||||
"model": "",
|
||||
"tags": [],
|
||||
}
|
||||
)
|
||||
|
||||
assert (agent, label, source) == ("unknown", "Unidentified", "unknown")
|
||||
|
||||
|
||||
def test_agent_usage_recovers_before_tokens_from_after_and_saved() -> None:
|
||||
summary = _build_agent_usage_summary(
|
||||
[
|
||||
{
|
||||
"provider": "openai",
|
||||
"model": "custom-model",
|
||||
"tags": {"client": "custom-agent"},
|
||||
"input_tokens_original": 0,
|
||||
"input_tokens_optimized": 70,
|
||||
"output_tokens": 5,
|
||||
"tokens_saved": 30,
|
||||
}
|
||||
],
|
||||
requests_by_provider={},
|
||||
requests_by_model={},
|
||||
global_before_tokens=100,
|
||||
global_after_tokens=70,
|
||||
global_tokens_saved=0,
|
||||
global_output_tokens=5,
|
||||
)
|
||||
|
||||
row = summary["agents"][0]
|
||||
|
||||
assert row["agent"] == "custom-agent"
|
||||
assert row["label"] == "Custom Agent"
|
||||
assert row["before_tokens"] == 100
|
||||
assert row["savings_percent"] == 30.0
|
||||
assert row["after_percent"] == 70.0
|
||||
assert row["share_of_saved_percent"] == 0.0
|
||||
assert summary["totals"]["savings_percent"] == 0.0
|
||||
|
||||
|
||||
def test_agent_usage_clamps_negative_token_values() -> None:
|
||||
summary = _build_agent_usage_summary(
|
||||
[
|
||||
{
|
||||
"provider": None,
|
||||
"model": None,
|
||||
"tags": {},
|
||||
"input_tokens_original": -100,
|
||||
"input_tokens_optimized": -50,
|
||||
"output_tokens": -5,
|
||||
"tokens_saved": -25,
|
||||
}
|
||||
],
|
||||
requests_by_provider={},
|
||||
requests_by_model={},
|
||||
global_before_tokens=0,
|
||||
global_after_tokens=0,
|
||||
global_tokens_saved=0,
|
||||
global_output_tokens=0,
|
||||
)
|
||||
|
||||
row = summary["agents"][0]
|
||||
|
||||
assert row["agent"] == "unknown"
|
||||
assert row["requests"] == 1
|
||||
assert row["before_tokens"] == 0
|
||||
assert row["after_tokens"] == 0
|
||||
assert row["tokens_saved"] == 0
|
||||
assert row["output_tokens"] == 0
|
||||
assert row["has_exact_tokens"] is False
|
||||
203
tests/test_gemini_function_response_waste.py
Normal file
203
tests/test_gemini_function_response_waste.py
Normal file
|
|
@ -0,0 +1,203 @@
|
|||
"""Gemini functionResponse waste-signal visibility (issue #819).
|
||||
|
||||
Gemini ``functionResponse`` parts are preserved verbatim on the wire (never
|
||||
compressed), but their payloads previously never reached ``parse_messages``,
|
||||
so tool output — where most waste lives — contributed nothing to waste
|
||||
detection on the Gemini paths.
|
||||
|
||||
The fix is telemetry-only:
|
||||
|
||||
1. ``_gemini_contents_to_messages(..., include_function_responses=True)``
|
||||
additionally emits each functionResponse payload as a ``role="tool"``
|
||||
message.
|
||||
2. ``TransformPipeline.apply(..., waste_messages=...)`` parses that richer
|
||||
list for waste signals instead of the transform input. The transform path
|
||||
and token accounting are untouched.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("fastapi")
|
||||
pytest.importorskip("httpx")
|
||||
|
||||
from headroom import OpenAIProvider, Tokenizer
|
||||
from headroom.config import HeadroomConfig
|
||||
from headroom.parser import parse_messages
|
||||
from headroom.proxy.server import HeadroomProxy, ProxyConfig
|
||||
from headroom.transforms.pipeline import TransformPipeline
|
||||
|
||||
_provider = OpenAIProvider()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def proxy() -> HeadroomProxy:
|
||||
config = ProxyConfig(
|
||||
optimize=False,
|
||||
cache_enabled=False,
|
||||
rate_limit_enabled=False,
|
||||
cost_tracking_enabled=False,
|
||||
)
|
||||
return HeadroomProxy(config)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def tokenizer() -> Tokenizer:
|
||||
return Tokenizer(_provider.get_token_counter("gpt-4o"), "gpt-4o")
|
||||
|
||||
|
||||
def _big_payload(rows: int = 200) -> dict:
|
||||
return {
|
||||
"result": [
|
||||
{"id": i, "name": f"item_{i}", "status": "ok", "score": i * 3.14} for i in range(rows)
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
def _function_response_content(payload: object, name: str = "fetch_data") -> dict:
|
||||
return {
|
||||
"role": "user",
|
||||
"parts": [{"functionResponse": {"name": name, "response": payload}}],
|
||||
}
|
||||
|
||||
|
||||
class TestFunctionResponseConversion:
|
||||
def test_default_conversion_emits_no_tool_messages(self, proxy):
|
||||
contents = [
|
||||
{"role": "user", "parts": [{"text": "fetch the data"}]},
|
||||
_function_response_content(_big_payload()),
|
||||
]
|
||||
messages, preserved = proxy._gemini_contents_to_messages(contents)
|
||||
assert [m["role"] for m in messages] == ["user"]
|
||||
assert preserved == {1}
|
||||
|
||||
def test_flag_emits_tool_message_for_dict_response(self, proxy):
|
||||
payload = _big_payload()
|
||||
contents = [
|
||||
{"role": "user", "parts": [{"text": "fetch the data"}]},
|
||||
_function_response_content(payload),
|
||||
]
|
||||
messages, preserved = proxy._gemini_contents_to_messages(
|
||||
contents, include_function_responses=True
|
||||
)
|
||||
assert [m["role"] for m in messages] == ["user", "tool"]
|
||||
assert json.loads(messages[1]["content"]) == payload
|
||||
# preserved_indices semantics unchanged: the entry is still restored
|
||||
# verbatim on the wire regardless of the telemetry conversion.
|
||||
assert preserved == {1}
|
||||
|
||||
def test_flag_passes_string_response_through(self, proxy):
|
||||
contents = [_function_response_content("plain text tool output")]
|
||||
messages, _ = proxy._gemini_contents_to_messages(contents, include_function_responses=True)
|
||||
assert messages == [{"role": "tool", "content": "plain text tool output"}]
|
||||
|
||||
def test_flag_skips_missing_response(self, proxy):
|
||||
contents = [
|
||||
{"role": "user", "parts": [{"functionResponse": {"name": "noop"}}]},
|
||||
{"role": "user", "parts": [{"functionResponse": {"name": "none", "response": None}}]},
|
||||
]
|
||||
messages, _ = proxy._gemini_contents_to_messages(contents, include_function_responses=True)
|
||||
assert messages == []
|
||||
|
||||
def test_flag_emits_text_before_tool_within_entry(self, proxy):
|
||||
contents = [
|
||||
{
|
||||
"role": "user",
|
||||
"parts": [
|
||||
{"text": "tool said:"},
|
||||
{"functionResponse": {"name": "f", "response": "output"}},
|
||||
],
|
||||
}
|
||||
]
|
||||
messages, _ = proxy._gemini_contents_to_messages(contents, include_function_responses=True)
|
||||
assert [m["role"] for m in messages] == ["user", "tool"]
|
||||
assert messages[0]["content"] == "tool said:"
|
||||
assert messages[1]["content"] == "output"
|
||||
|
||||
def test_unserializable_response_falls_back_to_str(self, proxy):
|
||||
circular: dict = {"name": "loop"}
|
||||
circular["self"] = circular
|
||||
text = proxy._function_response_text({"response": circular})
|
||||
assert "loop" in text
|
||||
|
||||
|
||||
class TestFunctionResponseWasteParsing:
|
||||
def test_function_response_payload_reaches_waste_signals(self, proxy, tokenizer):
|
||||
contents = [
|
||||
{"role": "user", "parts": [{"text": "fetch the data"}]},
|
||||
_function_response_content(_big_payload()),
|
||||
]
|
||||
messages, _ = proxy._gemini_contents_to_messages(contents, include_function_responses=True)
|
||||
blocks, _, waste = parse_messages(messages, tokenizer)
|
||||
assert any(b.kind == "tool_result" for b in blocks)
|
||||
assert waste.json_bloat_tokens > 0
|
||||
|
||||
def test_repeated_function_response_counts_as_reread(self, proxy, tokenizer):
|
||||
payload = _big_payload()
|
||||
filler = [{"role": "user", "parts": [{"text": f"working on step {i}"}]} for i in range(5)]
|
||||
contents = [
|
||||
_function_response_content(payload),
|
||||
*filler,
|
||||
_function_response_content(payload),
|
||||
]
|
||||
messages, _ = proxy._gemini_contents_to_messages(contents, include_function_responses=True)
|
||||
_, _, waste = parse_messages(messages, tokenizer)
|
||||
assert waste.reread_tokens > 0
|
||||
|
||||
|
||||
class TestPipelineWasteMessages:
|
||||
@staticmethod
|
||||
def _base_messages() -> list[dict]:
|
||||
# Compressible enough that the pipeline clears the >100 saved-token
|
||||
# gate that guards waste-signal detection.
|
||||
return [
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": "Inspect the data set."},
|
||||
{"role": "tool", "content": json.dumps(_big_payload(400)["result"])},
|
||||
]
|
||||
|
||||
def test_waste_messages_override_waste_source(self, tokenizer):
|
||||
messages = self._base_messages()
|
||||
extra_tool = {"role": "tool", "content": json.dumps(_big_payload(300))}
|
||||
|
||||
baseline = TransformPipeline(HeadroomConfig()).apply(
|
||||
[dict(m) for m in messages], model="gpt-4o", model_limit=128000
|
||||
)
|
||||
enriched = TransformPipeline(HeadroomConfig()).apply(
|
||||
[dict(m) for m in messages],
|
||||
model="gpt-4o",
|
||||
model_limit=128000,
|
||||
waste_messages=[*messages, extra_tool],
|
||||
)
|
||||
|
||||
assert baseline.waste_signals is not None
|
||||
assert enriched.waste_signals is not None
|
||||
assert enriched.waste_signals.json_bloat_tokens > baseline.waste_signals.json_bloat_tokens
|
||||
|
||||
def test_waste_messages_do_not_affect_transform_output(self, tokenizer):
|
||||
messages = self._base_messages()
|
||||
extra_tool = {"role": "tool", "content": json.dumps(_big_payload(300))}
|
||||
|
||||
baseline = TransformPipeline(HeadroomConfig()).apply(
|
||||
[dict(m) for m in messages], model="gpt-4o", model_limit=128000
|
||||
)
|
||||
enriched = TransformPipeline(HeadroomConfig()).apply(
|
||||
[dict(m) for m in messages],
|
||||
model="gpt-4o",
|
||||
model_limit=128000,
|
||||
waste_messages=[*messages, extra_tool],
|
||||
)
|
||||
|
||||
assert enriched.messages == baseline.messages
|
||||
assert enriched.tokens_before == baseline.tokens_before
|
||||
assert enriched.tokens_after == baseline.tokens_after
|
||||
|
||||
def test_no_waste_messages_falls_back_to_transform_input(self, tokenizer):
|
||||
result = TransformPipeline(HeadroomConfig()).apply(
|
||||
[dict(m) for m in self._base_messages()], model="gpt-4o", model_limit=128000
|
||||
)
|
||||
assert result.waste_signals is not None
|
||||
assert result.waste_signals.json_bloat_tokens > 0
|
||||
|
|
@ -80,6 +80,19 @@ def test_apply_and_revert_codex_provider_scope(monkeypatch, tmp_path: Path) -> N
|
|||
assert reverted.strip() == 'model = "gpt-4o"'
|
||||
|
||||
|
||||
def test_apply_codex_provider_scope_emits_flag_for_chatgpt_auth(
|
||||
monkeypatch, tmp_path: Path
|
||||
) -> None:
|
||||
config_path = tmp_path / "config.toml"
|
||||
(tmp_path / "auth.json").write_text('{"auth_mode": "chatgpt"}')
|
||||
monkeypatch.setattr("headroom.providers.codex.install.codex_config_path", lambda: config_path)
|
||||
manifest = _manifest(tmp_path)
|
||||
|
||||
apply_codex_provider_scope(manifest)
|
||||
|
||||
assert "requires_openai_auth = true" in config_path.read_text()
|
||||
|
||||
|
||||
def test_codex_build_install_env_returns_proxy_base_url() -> None:
|
||||
env = build_codex_install_env(port=5566, backend="ignored")
|
||||
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
import sys
|
||||
|
|
@ -10,6 +11,7 @@ from fastapi import Request
|
|||
|
||||
from headroom.proxy.handlers.openai import (
|
||||
OpenAIHandlerMixin,
|
||||
_openai_responses_unit_cache_key,
|
||||
_resolve_codex_routing_headers,
|
||||
)
|
||||
|
||||
|
|
@ -100,6 +102,38 @@ def test_resolve_codex_routing_ignores_invalid_jwt_payloads():
|
|||
assert headers["authorization"] == f"Bearer {token}"
|
||||
|
||||
|
||||
def test_openai_responses_unit_cache_key_includes_target_ratio() -> None:
|
||||
unit = SimpleNamespace(
|
||||
text="large tool output",
|
||||
provider="openai",
|
||||
endpoint="responses",
|
||||
role="tool",
|
||||
item_type="function_call_output",
|
||||
cache_zone="live",
|
||||
mutable=True,
|
||||
min_bytes=100,
|
||||
context=None,
|
||||
question=None,
|
||||
bias=None,
|
||||
metadata={},
|
||||
)
|
||||
|
||||
default_key = _openai_responses_unit_cache_key(unit, model="gpt-5.4")
|
||||
aggressive_key = _openai_responses_unit_cache_key(
|
||||
unit,
|
||||
model="gpt-5.4",
|
||||
target_ratio=0.10,
|
||||
)
|
||||
balanced_key = _openai_responses_unit_cache_key(
|
||||
unit,
|
||||
model="gpt-5.4",
|
||||
target_ratio=0.50,
|
||||
)
|
||||
|
||||
assert aggressive_key != default_key
|
||||
assert aggressive_key != balanced_key
|
||||
|
||||
|
||||
class _DummyMetrics:
|
||||
async def record_request(self, **kwargs): # noqa: ANN003
|
||||
return None
|
||||
|
|
@ -253,6 +287,37 @@ def test_handle_openai_responses_routes_chatgpt_auth_to_backend_api(monkeypatch)
|
|||
assert response.status_code == 200
|
||||
|
||||
|
||||
def test_handle_openai_responses_chatgpt_codex_timeout_fails_open(monkeypatch):
|
||||
token = _jwt(
|
||||
{
|
||||
"https://api.openai.com/auth": {
|
||||
"chatgpt_account_id": "acct-from-jwt",
|
||||
}
|
||||
}
|
||||
)
|
||||
request = _build_request(
|
||||
{"model": "gpt-5.4", "input": "large context"},
|
||||
{"Authorization": f"Bearer {token}"},
|
||||
)
|
||||
handler = _DummyOpenAIHandler()
|
||||
handler.config.optimize = True
|
||||
|
||||
async def timeout_compression(*args, **kwargs): # noqa: ANN002, ANN003
|
||||
raise asyncio.TimeoutError()
|
||||
|
||||
handler._compress_openai_responses_payload_in_executor = timeout_compression
|
||||
monkeypatch.setattr("headroom.tokenizers.get_tokenizer", lambda model: _DummyTokenizer())
|
||||
|
||||
response = anyio.run(handler.handle_openai_responses, request)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert handler.captured_request is not None
|
||||
method, url, headers, body = handler.captured_request
|
||||
assert method == "POST"
|
||||
assert url == "https://chatgpt.com/backend-api/codex/responses"
|
||||
assert body["input"] == "large context"
|
||||
|
||||
|
||||
def test_handle_openai_responses_routes_api_key_auth_direct_to_openai(monkeypatch):
|
||||
request = _build_request(
|
||||
{"model": "gpt-4o-mini", "input": "hello"},
|
||||
|
|
|
|||
|
|
@ -1051,3 +1051,229 @@ class TestStrandsToolResultBlocks:
|
|||
|
||||
assert anthropic_waste.total() > 0
|
||||
assert anthropic_waste.total() == openai_waste.total()
|
||||
|
||||
|
||||
# --- TestCallArgMatchReread ---
|
||||
|
||||
|
||||
class TestCallArgMatchReread:
|
||||
"""Tests for re-issued-call (arg-match) reread detection in parse_messages."""
|
||||
|
||||
LARGE_CONTENT = "def handler(event):\n return process(event)\n" * 10 # > 200 chars
|
||||
CHANGED_CONTENT = LARGE_CONTENT + "# mtime 1718000000\n"
|
||||
|
||||
def _expected_tokens(self, text):
|
||||
"""Mirror mock_tokenizer + message overhead used for tool_result blocks."""
|
||||
return len(text) // 4 + 1 + 4
|
||||
|
||||
@staticmethod
|
||||
def _filler(n):
|
||||
"""Interleaved turns that push a repeat beyond the polling gap."""
|
||||
return [
|
||||
{"role": "assistant" if i % 2 == 0 else "user", "content": f"step {i} of the task"}
|
||||
for i in range(n)
|
||||
]
|
||||
|
||||
@staticmethod
|
||||
def _openai_call(call_id, name, arguments):
|
||||
return {
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [{"id": call_id, "function": {"name": name, "arguments": arguments}}],
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _openai_result(call_id, content):
|
||||
return {"role": "tool", "tool_call_id": call_id, "content": content}
|
||||
|
||||
def test_canonical_call_key_normalizes_serialization(self):
|
||||
"""Reordered JSON-string args, dict args, and spaced JSON hash equal."""
|
||||
from headroom.parser import _canonical_call_key
|
||||
|
||||
k1 = _canonical_call_key("read_file", '{"path": "a.py", "lines": 100}')
|
||||
k2 = _canonical_call_key("read_file", '{"lines":100,"path":"a.py"}')
|
||||
k3 = _canonical_call_key("read_file", {"path": "a.py", "lines": 100})
|
||||
assert k1 == k2 == k3
|
||||
assert _canonical_call_key("read_file", '{"path": "b.py", "lines": 100}') != k1
|
||||
assert _canonical_call_key("grep", '{"path": "a.py", "lines": 100}') != k1
|
||||
|
||||
def test_reissued_call_changed_result_counts(self, mock_tokenizer):
|
||||
"""Identical call re-issued far apart counts even when result bytes differ."""
|
||||
messages = (
|
||||
[
|
||||
self._openai_call("c1", "read_file", '{"path": "a.py", "lines": 100}'),
|
||||
self._openai_result("c1", self.LARGE_CONTENT),
|
||||
]
|
||||
+ self._filler(4)
|
||||
+ [
|
||||
self._openai_call("c2", "read_file", '{"lines":100,"path":"a.py"}'),
|
||||
self._openai_result("c2", self.CHANGED_CONTENT),
|
||||
]
|
||||
)
|
||||
_, _, waste = parse_messages(messages, mock_tokenizer)
|
||||
assert waste.reread_tokens == self._expected_tokens(self.CHANGED_CONTENT)
|
||||
|
||||
def test_identical_result_not_double_counted(self, mock_tokenizer):
|
||||
"""Byte-identical repeat is counted once (content-hash pass wins)."""
|
||||
messages = (
|
||||
[
|
||||
self._openai_call("c1", "read_file", '{"path": "a.py"}'),
|
||||
self._openai_result("c1", self.LARGE_CONTENT),
|
||||
]
|
||||
+ self._filler(4)
|
||||
+ [
|
||||
self._openai_call("c2", "read_file", '{"path": "a.py"}'),
|
||||
self._openai_result("c2", self.LARGE_CONTENT),
|
||||
]
|
||||
)
|
||||
_, _, waste = parse_messages(messages, mock_tokenizer)
|
||||
assert waste.reread_tokens == self._expected_tokens(self.LARGE_CONTENT)
|
||||
|
||||
def test_adjacent_reissue_is_polling(self, mock_tokenizer):
|
||||
"""Back-to-back identical calls (poll loop) are not re-reads."""
|
||||
messages = [
|
||||
self._openai_call("c1", "check_ci", '{"run": 7}'),
|
||||
self._openai_result("c1", self.LARGE_CONTENT),
|
||||
self._openai_call("c2", "check_ci", '{"run": 7}'),
|
||||
self._openai_result("c2", self.CHANGED_CONTENT),
|
||||
]
|
||||
_, _, waste = parse_messages(messages, mock_tokenizer)
|
||||
assert waste.reread_tokens == 0
|
||||
|
||||
def test_different_args_not_matched(self, mock_tokenizer):
|
||||
"""Same tool with different arguments is not a re-issued call."""
|
||||
messages = (
|
||||
[
|
||||
self._openai_call("c1", "read_file", '{"path": "a.py"}'),
|
||||
self._openai_result("c1", self.LARGE_CONTENT),
|
||||
]
|
||||
+ self._filler(4)
|
||||
+ [
|
||||
self._openai_call("c2", "read_file", '{"path": "b.py"}'),
|
||||
self._openai_result("c2", self.CHANGED_CONTENT),
|
||||
]
|
||||
)
|
||||
_, _, waste = parse_messages(messages, mock_tokenizer)
|
||||
assert waste.reread_tokens == 0
|
||||
|
||||
def test_small_result_ignored(self, mock_tokenizer):
|
||||
"""Repeat of a call whose result is trivially small is skipped."""
|
||||
messages = (
|
||||
[
|
||||
self._openai_call("c1", "run_tests", "{}"),
|
||||
self._openai_result("c1", "ok"),
|
||||
]
|
||||
+ self._filler(4)
|
||||
+ [
|
||||
self._openai_call("c2", "run_tests", "{}"),
|
||||
self._openai_result("c2", "ok again"),
|
||||
]
|
||||
)
|
||||
_, _, waste = parse_messages(messages, mock_tokenizer)
|
||||
assert waste.reread_tokens == 0
|
||||
|
||||
def test_repeat_call_without_result_not_counted(self, mock_tokenizer):
|
||||
"""A re-issued call with no recorded result contributes nothing."""
|
||||
messages = (
|
||||
[
|
||||
self._openai_call("c1", "read_file", '{"path": "a.py"}'),
|
||||
self._openai_result("c1", self.LARGE_CONTENT),
|
||||
]
|
||||
+ self._filler(4)
|
||||
+ [self._openai_call("c2", "read_file", '{"path": "a.py"}')]
|
||||
)
|
||||
_, _, waste = parse_messages(messages, mock_tokenizer)
|
||||
assert waste.reread_tokens == 0
|
||||
|
||||
def test_anthropic_tool_use_produces_tool_call_blocks(self, mock_tokenizer):
|
||||
"""Anthropic tool_use parts become tool_call blocks with call metadata."""
|
||||
messages = [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "text", "text": "Reading the file now."},
|
||||
{
|
||||
"type": "tool_use",
|
||||
"id": "t1",
|
||||
"name": "read_file",
|
||||
"input": {"path": "a.py"},
|
||||
},
|
||||
],
|
||||
}
|
||||
]
|
||||
blocks, _, _ = parse_messages(messages, mock_tokenizer)
|
||||
tool_calls = [b for b in blocks if b.kind == "tool_call"]
|
||||
assert len(tool_calls) == 1
|
||||
assert tool_calls[0].flags["function_name"] == "read_file"
|
||||
assert tool_calls[0].flags["tool_call_id"] == "t1"
|
||||
assert tool_calls[0].flags["call_key"]
|
||||
|
||||
def test_anthropic_reissued_call_changed_result_counts(self, mock_tokenizer):
|
||||
"""Full Anthropic-format flow: re-issued tool_use with drifted result."""
|
||||
|
||||
def call(uid):
|
||||
return {
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "tool_use", "id": uid, "name": "read_file", "input": {"path": "a.py"}}
|
||||
],
|
||||
}
|
||||
|
||||
def result(uid, content):
|
||||
return {
|
||||
"role": "user",
|
||||
"content": [{"type": "tool_result", "tool_use_id": uid, "content": content}],
|
||||
}
|
||||
|
||||
messages = (
|
||||
[call("t1"), result("t1", self.LARGE_CONTENT)]
|
||||
+ self._filler(4)
|
||||
+ [call("t2"), result("t2", self.CHANGED_CONTENT)]
|
||||
)
|
||||
_, _, waste = parse_messages(messages, mock_tokenizer)
|
||||
assert waste.reread_tokens == self._expected_tokens(self.CHANGED_CONTENT)
|
||||
|
||||
def test_strands_tooluse_matched(self, mock_tokenizer):
|
||||
"""Strands/Bedrock toolUse/toolResult format is matched the same way."""
|
||||
|
||||
def call(uid):
|
||||
return {
|
||||
"role": "assistant",
|
||||
"content": [{"toolUse": {"toolUseId": uid, "name": "search", "input": {"q": "x"}}}],
|
||||
}
|
||||
|
||||
def result(uid, content):
|
||||
return {
|
||||
"role": "user",
|
||||
"content": [{"toolResult": {"toolUseId": uid, "content": [{"text": content}]}}],
|
||||
}
|
||||
|
||||
messages = (
|
||||
[call("s1"), result("s1", self.LARGE_CONTENT)]
|
||||
+ self._filler(4)
|
||||
+ [call("s2"), result("s2", self.CHANGED_CONTENT)]
|
||||
)
|
||||
_, _, waste = parse_messages(messages, mock_tokenizer)
|
||||
assert waste.reread_tokens == self._expected_tokens(self.CHANGED_CONTENT)
|
||||
|
||||
def test_cross_format_call_key_parity(self, mock_tokenizer):
|
||||
"""OpenAI JSON-string args and Anthropic dict input produce the same call_key."""
|
||||
openai_msgs = [self._openai_call("c1", "read_file", '{"lines": 100, "path": "a.py"}')]
|
||||
anthropic_msgs = [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_use",
|
||||
"id": "t1",
|
||||
"name": "read_file",
|
||||
"input": {"path": "a.py", "lines": 100},
|
||||
}
|
||||
],
|
||||
}
|
||||
]
|
||||
o_blocks, _, _ = parse_messages(openai_msgs, mock_tokenizer)
|
||||
a_blocks, _, _ = parse_messages(anthropic_msgs, mock_tokenizer)
|
||||
o_key = [b for b in o_blocks if b.kind == "tool_call"][0].flags["call_key"]
|
||||
a_key = [b for b in a_blocks if b.kind == "tool_call"][0].flags["call_key"]
|
||||
assert o_key == a_key
|
||||
|
|
|
|||
|
|
@ -1,26 +1,78 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from headroom.providers.codex.install import build_provider_section
|
||||
from pathlib import Path
|
||||
|
||||
from headroom.providers.codex.install import build_provider_section, codex_uses_chatgpt_auth
|
||||
|
||||
|
||||
def test_codex_provider_section_no_requires_openai_auth() -> None:
|
||||
"""Bug 3 (#406): build_provider_section must NOT include requires_openai_auth.
|
||||
def test_codex_provider_section_omits_requires_openai_auth_by_default() -> None:
|
||||
"""#406: the flag must default off (API-key users), and only on for OAuth.
|
||||
|
||||
Setting requires_openai_auth on a custom [model_providers.headroom] block
|
||||
forces codex to demand OpenAI OAuth login for every headroom-routed request.
|
||||
Headroom is a local proxy — it must never carry this flag.
|
||||
forces codex to demand OpenAI OAuth login for every headroom-routed request,
|
||||
which breaks API-key users; so callers opt in explicitly for ChatGPT users.
|
||||
"""
|
||||
section = build_provider_section(port=8787, name="OpenAI via Headroom proxy")
|
||||
|
||||
assert 'name = "OpenAI via Headroom proxy"' in section
|
||||
assert 'base_url = "http://127.0.0.1:8787/v1"' in section
|
||||
assert "requires_openai_auth" not in section, (
|
||||
f"requires_openai_auth must be absent from the headroom provider section; got:\n{section}"
|
||||
f"requires_openai_auth must be absent by default; got:\n{section}"
|
||||
)
|
||||
assert "supports_websockets = true" in section
|
||||
assert 'env_key = "OPENAI_API_KEY"' not in section
|
||||
|
||||
|
||||
def test_codex_provider_section_emits_requires_openai_auth_when_flagged() -> None:
|
||||
section = build_provider_section(
|
||||
port=8787, name="OpenAI via Headroom proxy", requires_openai_auth=True
|
||||
)
|
||||
|
||||
assert "requires_openai_auth = true" in section
|
||||
|
||||
|
||||
def test_codex_uses_chatgpt_auth_true_for_chatgpt_mode(tmp_path: Path) -> None:
|
||||
auth = tmp_path / "auth.json"
|
||||
auth.write_text('{"auth_mode": "chatgpt"}', encoding="utf-8")
|
||||
|
||||
assert codex_uses_chatgpt_auth(auth) is True
|
||||
|
||||
|
||||
def test_codex_uses_chatgpt_auth_true_for_account_id_without_mode(tmp_path: Path) -> None:
|
||||
auth = tmp_path / "auth.json"
|
||||
auth.write_text('{"tokens": {"account_id": "acct_1"}}', encoding="utf-8")
|
||||
|
||||
assert codex_uses_chatgpt_auth(auth) is True
|
||||
|
||||
|
||||
def test_codex_uses_chatgpt_auth_false_for_api_key(tmp_path: Path) -> None:
|
||||
auth = tmp_path / "auth.json"
|
||||
auth.write_text('{"auth_mode": "apikey", "OPENAI_API_KEY": "sk-x"}', encoding="utf-8")
|
||||
|
||||
assert codex_uses_chatgpt_auth(auth) is False
|
||||
|
||||
|
||||
def test_codex_uses_chatgpt_auth_false_for_missing_or_malformed(tmp_path: Path) -> None:
|
||||
assert codex_uses_chatgpt_auth(tmp_path / "absent.json") is False
|
||||
bad = tmp_path / "auth.json"
|
||||
bad.write_text("not json", encoding="utf-8")
|
||||
assert codex_uses_chatgpt_auth(bad) is False
|
||||
|
||||
|
||||
def test_codex_uses_chatgpt_auth_false_for_non_dict_json(tmp_path: Path) -> None:
|
||||
auth = tmp_path / "auth.json"
|
||||
auth.write_text("[]", encoding="utf-8")
|
||||
|
||||
assert codex_uses_chatgpt_auth(auth) is False
|
||||
|
||||
|
||||
def test_codex_uses_chatgpt_auth_false_for_empty_object(tmp_path: Path) -> None:
|
||||
auth = tmp_path / "auth.json"
|
||||
auth.write_text("{}", encoding="utf-8")
|
||||
|
||||
assert codex_uses_chatgpt_auth(auth) is False
|
||||
|
||||
|
||||
def test_codex_provider_section_supports_custom_markers() -> None:
|
||||
section = build_provider_section(
|
||||
port=9100,
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ from headroom.proxy.handlers.openai import (
|
|||
OpenAIHandlerMixin,
|
||||
_decode_openai_bearer_payload,
|
||||
_passthrough_usage_from_json,
|
||||
_prefers_http1_passthrough,
|
||||
)
|
||||
from headroom.proxy.helpers import _headroom_bypass_enabled
|
||||
from headroom.proxy.server import HeadroomProxy
|
||||
|
|
@ -51,6 +52,31 @@ class _TimeoutHttpClient:
|
|||
raise httpx.ConnectTimeout("connect timed out")
|
||||
|
||||
|
||||
class _RecordingHttpClient:
|
||||
def __init__(self, label: str) -> None:
|
||||
self.label = label
|
||||
self.calls = 0
|
||||
|
||||
async def request(self, **kwargs): # noqa: ANN001, ANN201
|
||||
self.calls += 1
|
||||
request = httpx.Request(kwargs["method"], kwargs["url"])
|
||||
return httpx.Response(
|
||||
200,
|
||||
request=request,
|
||||
headers={"content-type": "application/json"},
|
||||
json={"client": self.label},
|
||||
)
|
||||
|
||||
|
||||
class _ChatGPTAccountRequest:
|
||||
method = "GET"
|
||||
headers = {}
|
||||
url = SimpleNamespace(path="/backend-api/me", query="")
|
||||
|
||||
async def body(self) -> bytes:
|
||||
return b""
|
||||
|
||||
|
||||
class _PassthroughRequest:
|
||||
method = "GET"
|
||||
headers = {}
|
||||
|
|
@ -249,6 +275,61 @@ def test_openai_passthrough_connect_timeout_returns_502() -> None:
|
|||
assert "Failed to connect to upstream API" in payload["error"]["message"]
|
||||
|
||||
|
||||
def test_prefers_http1_passthrough_matches_chatgpt_hosts_only() -> None:
|
||||
assert _prefers_http1_passthrough("https://chatgpt.com") is True
|
||||
assert _prefers_http1_passthrough("https://chatgpt.com/backend-api/me") is True
|
||||
assert _prefers_http1_passthrough("https://api.chatgpt.com") is True
|
||||
assert _prefers_http1_passthrough("https://CHATGPT.COM/backend-api/me") is True
|
||||
assert _prefers_http1_passthrough("https://api.openai.com") is False
|
||||
assert _prefers_http1_passthrough("https://notchatgpt.com") is False
|
||||
assert _prefers_http1_passthrough("https://chatgpt.com.evil.com") is False
|
||||
assert _prefers_http1_passthrough("") is False
|
||||
|
||||
|
||||
def test_chatgpt_passthrough_uses_http1_client() -> None:
|
||||
handler = object.__new__(OpenAIHandlerMixin)
|
||||
handler.http_client = _RecordingHttpClient("h2")
|
||||
handler.http_client_h1 = _RecordingHttpClient("h1")
|
||||
|
||||
response = asyncio.run(
|
||||
handler.handle_passthrough(_ChatGPTAccountRequest(), "https://chatgpt.com")
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert json.loads(response.body)["client"] == "h1"
|
||||
assert handler.http_client.calls == 0
|
||||
assert handler.http_client_h1.calls == 1
|
||||
|
||||
|
||||
def test_non_chatgpt_passthrough_uses_default_client() -> None:
|
||||
handler = object.__new__(OpenAIHandlerMixin)
|
||||
handler.http_client = _RecordingHttpClient("h2")
|
||||
handler.http_client_h1 = _RecordingHttpClient("h1")
|
||||
|
||||
response = asyncio.run(
|
||||
handler.handle_passthrough(_PassthroughRequest(), "https://api.openai.com")
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert json.loads(response.body)["client"] == "h2"
|
||||
assert handler.http_client.calls == 1
|
||||
assert handler.http_client_h1.calls == 0
|
||||
|
||||
|
||||
def test_chatgpt_passthrough_falls_back_when_h1_client_missing() -> None:
|
||||
handler = object.__new__(OpenAIHandlerMixin)
|
||||
handler.http_client = _RecordingHttpClient("h2")
|
||||
handler.http_client_h1 = None
|
||||
|
||||
response = asyncio.run(
|
||||
handler.handle_passthrough(_ChatGPTAccountRequest(), "https://chatgpt.com")
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert json.loads(response.body)["client"] == "h2"
|
||||
assert handler.http_client.calls == 1
|
||||
|
||||
|
||||
def test_passthrough_usage_normalizes_vertex_usage_metadata() -> None:
|
||||
usage = _passthrough_usage_from_json(
|
||||
{
|
||||
|
|
|
|||
162
tests/test_proxy_stats_recent_requests.py
Normal file
162
tests/test_proxy_stats_recent_requests.py
Normal file
|
|
@ -0,0 +1,162 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from headroom.proxy import server
|
||||
from headroom.proxy.models import ProxyConfig
|
||||
from headroom.proxy.server import create_app
|
||||
|
||||
|
||||
class FakeRequestLogger:
|
||||
def __init__(self) -> None:
|
||||
self._logs: list[dict[str, object]] = []
|
||||
|
||||
@property
|
||||
def logs(self) -> list[dict[str, object]]:
|
||||
return self._logs
|
||||
|
||||
@logs.setter
|
||||
def logs(self, value: list[dict[str, object]]) -> None:
|
||||
self._logs = value
|
||||
|
||||
def get_recent(self, limit: int) -> list[dict[str, object]]:
|
||||
return self._logs[-limit:]
|
||||
|
||||
|
||||
class FakeLogEntry(dict[str, object]):
|
||||
def __getattr__(self, name: str) -> object:
|
||||
return self.get(name)
|
||||
|
||||
|
||||
def test_stats_refreshes_recent_requests_when_cached() -> None:
|
||||
app = create_app(
|
||||
ProxyConfig(
|
||||
optimize=False,
|
||||
cache_enabled=False,
|
||||
rate_limit_enabled=False,
|
||||
cost_tracking_enabled=False,
|
||||
log_requests=False,
|
||||
ccr_inject_tool=False,
|
||||
ccr_handle_responses=False,
|
||||
ccr_context_tracking=False,
|
||||
http2=False,
|
||||
)
|
||||
)
|
||||
logger = FakeRequestLogger()
|
||||
app.state.proxy.logger = logger
|
||||
|
||||
first_log = FakeLogEntry(
|
||||
{
|
||||
"timestamp": "2026-06-11T10:00:00Z",
|
||||
"provider": "openai",
|
||||
"model": "gpt-4.1",
|
||||
"input_tokens_original": 100,
|
||||
"input_tokens_optimized": 60,
|
||||
"tokens_saved": 40,
|
||||
"savings_percent": 40.0,
|
||||
}
|
||||
)
|
||||
second_log = FakeLogEntry(
|
||||
{
|
||||
"timestamp": "2026-06-11T10:01:00Z",
|
||||
"provider": "anthropic",
|
||||
"model": "claude-sonnet",
|
||||
"input_tokens_original": 200,
|
||||
"input_tokens_optimized": 120,
|
||||
"tokens_saved": 80,
|
||||
"savings_percent": 40.0,
|
||||
}
|
||||
)
|
||||
|
||||
with TestClient(app) as client:
|
||||
logger.logs = [first_log]
|
||||
first_response = client.get("/stats?cached=1")
|
||||
assert first_response.status_code == 200
|
||||
assert first_response.json()["recent_requests"][-1]["model"] == "gpt-4.1"
|
||||
|
||||
logger.logs = [first_log, second_log]
|
||||
second_response = client.get("/stats?cached=1")
|
||||
assert second_response.status_code == 200
|
||||
second_payload = second_response.json()
|
||||
|
||||
assert second_payload["recent_requests"][-1]["model"] == "claude-sonnet"
|
||||
assert second_payload["request_logs"][-1]["model"] == "claude-sonnet"
|
||||
|
||||
|
||||
def test_agent_usage_totals_use_proxy_only_savings(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("HEADROOM_REQUIRE_RUST_CORE", "false")
|
||||
monkeypatch.setattr(
|
||||
server,
|
||||
"_get_context_tool_stats",
|
||||
lambda: {
|
||||
"tool": "rtk",
|
||||
"label": "RTK",
|
||||
"tokens_saved": 500,
|
||||
"session": {},
|
||||
"lifetime": {},
|
||||
},
|
||||
)
|
||||
app = create_app(
|
||||
ProxyConfig(
|
||||
optimize=False,
|
||||
cache_enabled=False,
|
||||
rate_limit_enabled=False,
|
||||
cost_tracking_enabled=False,
|
||||
log_requests=False,
|
||||
ccr_inject_tool=False,
|
||||
ccr_handle_responses=False,
|
||||
ccr_context_tracking=False,
|
||||
http2=False,
|
||||
)
|
||||
)
|
||||
logger = FakeRequestLogger()
|
||||
app.state.proxy.logger = logger
|
||||
|
||||
logger.logs = [
|
||||
FakeLogEntry(
|
||||
{
|
||||
"timestamp": "2026-06-11T10:00:00Z",
|
||||
"provider": "openai",
|
||||
"model": "gpt-5.2-codex",
|
||||
"tags": {"client": "codex"},
|
||||
"input_tokens_original": 1000,
|
||||
"input_tokens_optimized": 900,
|
||||
"output_tokens": 50,
|
||||
"tokens_saved": 100,
|
||||
"savings_percent": 10.0,
|
||||
}
|
||||
)
|
||||
]
|
||||
|
||||
with TestClient(app) as client:
|
||||
proxy = client.app.state.proxy
|
||||
proxy.metrics.tokens_input_total = 900
|
||||
proxy.metrics.tokens_saved_total = 100
|
||||
proxy.metrics.tokens_output_total = 50
|
||||
|
||||
response = client.get("/stats")
|
||||
|
||||
assert response.status_code == 200
|
||||
payload = response.json()
|
||||
|
||||
assert payload["tokens"]["saved"] == 600
|
||||
assert payload["agent_usage"]["totals"]["before_tokens"] == 1000
|
||||
assert payload["agent_usage"]["totals"]["tokens_saved"] == 100
|
||||
assert payload["agent_usage"]["totals"]["savings_percent"] == 10.0
|
||||
assert payload["agent_usage"]["agents"][0]["share_of_saved_percent"] == 100.0
|
||||
|
||||
|
||||
def test_stats_preserves_default_smart_crusher_compaction_state() -> None:
|
||||
config = ProxyConfig(
|
||||
optimize=False,
|
||||
cache_enabled=False,
|
||||
rate_limit_enabled=False,
|
||||
cost_tracking_enabled=False,
|
||||
)
|
||||
client = TestClient(create_app(config))
|
||||
|
||||
response = client.get("/stats")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json()["config"]["smart_crusher_with_compaction"] is None
|
||||
|
|
@ -4,6 +4,8 @@ from __future__ import annotations
|
|||
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
ROOT = Path(__file__).resolve().parent.parent
|
||||
|
||||
|
||||
|
|
@ -103,26 +105,37 @@ def test_no_openssl_sys_in_wheel_build_tree() -> None:
|
|||
import subprocess
|
||||
|
||||
for crate in ("headroom-py", "headroom-proxy", "headroom-core"):
|
||||
result = subprocess.run(
|
||||
[
|
||||
"cargo",
|
||||
"tree",
|
||||
"--target",
|
||||
"x86_64-unknown-linux-gnu",
|
||||
"-p",
|
||||
crate,
|
||||
"-i",
|
||||
"openssl-sys",
|
||||
],
|
||||
cwd=str(ROOT),
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=False,
|
||||
)
|
||||
try:
|
||||
result = subprocess.run(
|
||||
[
|
||||
"cargo",
|
||||
"tree",
|
||||
"--target",
|
||||
"x86_64-unknown-linux-gnu",
|
||||
"-p",
|
||||
crate,
|
||||
"-i",
|
||||
"openssl-sys",
|
||||
],
|
||||
cwd=str(ROOT),
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=False,
|
||||
)
|
||||
except FileNotFoundError:
|
||||
pytest.skip("cargo is unavailable in this environment")
|
||||
# `cargo tree -i <pkg>` exits 101 with "did not match any
|
||||
# packages" when the package is NOT in the tree — the GREEN
|
||||
# case. Exit 0 with a tree of consumers means it IS pulled.
|
||||
not_in_tree = result.returncode != 0 and "did not match any packages" in result.stderr
|
||||
if (
|
||||
result.returncode != 0
|
||||
and "package ID specification `openssl-sys` did not match"
|
||||
not in (result.stderr + result.stdout)
|
||||
):
|
||||
pytest.skip(
|
||||
"cargo dependency tree for the Linux wheel target is unavailable in this environment"
|
||||
)
|
||||
assert not_in_tree, (
|
||||
f"openssl-sys is back in {crate}'s build tree:\n"
|
||||
f"stdout:\n{result.stdout}\n"
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue