Files
chopratejas--headroom/tests/test_provider_cortex_code.py
sfc-gh-nashukla d9d0bf4b79 feat(providers): add Cortex Code (Snowflake CoCo) as a supported agent (#1190)
## Description

Adds **Cortex Code (CoCo)** — Snowflake's AI coding CLI — as a
first-class headroom provider alongside Claude Code, Codex, and Cursor.

Cortex Code routes requests to Snowflake's Cortex inference endpoint via
the OpenAI-compatible pipeline. This PR adds the provider slice,
registers it under `"cortex-code"`, and ships tests that measure real
token savings against `claude-sonnet-4-6`.

Closes #

## Type of Change

- [x] New feature (non-breaking change that adds functionality)
- [x] Documentation update

## Changes Made

- `headroom/providers/cortex_code/__init__.py` — new provider package
- `headroom/providers/cortex_code/runtime.py` — `proxy_base_url()`,
`build_launch_env()`, `default_api_url()` (reads `SNOWFLAKE_HOST` /
`SNOWFLAKE_ACCOUNT`)
- `headroom/providers/cortex_code/install.py` — `build_install_env()`
sets `OPENAI_BASE_URL`; `render_setup_lines()`
- `headroom/providers/install_registry.py` — registers `"cortex-code"`
in `_ENV_BUILDERS`
- `tests/test_provider_cortex_code.py` — 15 unit tests
- `tests/test_cortex_code_compression.py` — 5 compression benchmark
tests (no API key needed)
- `tests/e2e_cortex_savings.py` — real REST API benchmark; reads
`SF_CONN`/`SF_HOST` from env, no hardcoded identifiers
- `docs/cortex-code.md` — integration guide (quick start, library mode,
auth, limitations)
- `README.md` — Cortex Code row added to agent compatibility matrix

## Testing

- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [x] Type checking passes (`mypy headroom`)
- [x] New tests added for new functionality
- [x] Manual testing performed

### Test Output

```text
$ uv run --with pytest pytest tests/test_provider_cortex_code.py tests/test_cortex_code_compression.py -v

tests/test_provider_cortex_code.py::test_cortex_code_proxy_base_url_is_openai_compatible PASSED
tests/test_provider_cortex_code.py::test_cortex_code_proxy_base_url_uses_given_port PASSED
tests/test_provider_cortex_code.py::test_cortex_code_build_install_env_sets_openai_base_url PASSED
tests/test_provider_cortex_code.py::test_cortex_code_build_launch_env_does_not_mutate_input PASSED
tests/test_provider_cortex_code.py::test_cortex_code_build_launch_env_applies_project_prefix PASSED
tests/test_provider_cortex_code.py::test_cortex_code_build_launch_env_ignores_blank_project PASSED
tests/test_provider_cortex_code.py::test_cortex_code_render_setup_lines_contains_proxy_url PASSED
tests/test_provider_cortex_code.py::test_cortex_code_render_setup_lines_project_attribution PASSED
tests/test_provider_cortex_code.py::test_cortex_code_default_api_url_reads_snowflake_host_env PASSED
tests/test_provider_cortex_code.py::test_cortex_code_default_api_url_constructs_url_from_account_name PASSED
tests/test_provider_cortex_code.py::test_cortex_code_default_api_url_host_takes_priority_over_account PASSED
tests/test_provider_cortex_code.py::test_cortex_code_default_api_url_falls_back_when_no_env PASSED
tests/test_provider_cortex_code.py::test_cortex_code_default_api_url_preserves_https_prefix PASSED
tests/test_provider_cortex_code.py::test_cortex_code_install_registry_includes_cortex_code PASSED
tests/test_provider_cortex_code.py::test_cortex_code_install_registry_unknown_target_skipped PASSED
tests/test_cortex_code_compression.py::test_cortex_code_headroom_compression_saves_tokens PASSED
tests/test_cortex_code_compression.py::test_cortex_code_tool_results_are_compressed_not_user_turns PASSED
tests/test_cortex_code_compression.py::test_cortex_code_tables_json_compresses PASSED
tests/test_cortex_code_compression.py::test_cortex_code_rag_search_json_compresses PASSED
tests/test_cortex_code_compression.py::test_cortex_code_compression_is_lossless_on_key_content PASSED

20 passed, 1 warning in 1.91s
```

## Real Behavior Proof

- Environment: macOS, Python 3.11, headroom 0.27.0, Snowflake Cortex
(claude-sonnet-4-6)
- Exact command / steps: `SF_CONN=<connection-name> python3
tests/e2e_cortex_savings.py`
- Observed result: 62% average token reduction across 4 payload types;
usage.prompt_tokens confirmed in live API responses (full output in Test
Output above)
- Not tested: headroom wrap cortex-code proxy mode — Cortex REST API
path /api/v2/cortex/inference:complete differs from
/v1/chat/completions; library mode is the supported path (documented in
docs/cortex-code.md Limitations)

```text
  Tokens saved  :    22,077  prompt tokens  (4 calls)
  Avg per call  :     5,519  tokens  /  $0.01656
  At 1k/day     :  $16.56/day  |  $6,044/year
```

## Review Readiness

- [x] I have performed a self-review
- [x] This PR is ready for human review

## Checklist

- [x] My code follows the project's style guidelines
- [x] I have performed a self-review of my code
- [x] I have commented my code, particularly in hard-to-understand areas
- [x] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [x] New and existing unit tests pass locally with my changes
- [ ] I have updated the CHANGELOG.md if applicable

## Additional Notes

Pre-commit hooks skipped locally due to a GPG signing / ruff-format
stash conflict in the dev environment. `ruff check` passes clean on all
new files.

---------

Co-authored-by: Cortex Code <noreply@snowflake.com>
2026-06-21 22:18:47 -07:00

100 lines
3.8 KiB
Python

from __future__ import annotations
from headroom.providers.cortex_code import build_install_env, proxy_base_url, render_setup_lines
from headroom.providers.cortex_code.runtime import build_launch_env, default_api_url
def test_cortex_code_proxy_base_url_is_openai_compatible() -> None:
assert proxy_base_url(8787) == "http://127.0.0.1:8787/v1"
def test_cortex_code_proxy_base_url_uses_given_port() -> None:
assert proxy_base_url(9999) == "http://127.0.0.1:9999/v1"
def test_cortex_code_build_install_env_sets_openai_base_url() -> None:
env = build_install_env(port=8787, backend="ignored")
assert env == {"OPENAI_BASE_URL": "http://127.0.0.1:8787/v1"}
def test_cortex_code_build_launch_env_does_not_mutate_input() -> None:
source = {"EXISTING": "val"}
env, lines = build_launch_env(port=9999, environ=source)
assert source == {"EXISTING": "val"}
assert env["OPENAI_BASE_URL"] == "http://127.0.0.1:9999/v1"
assert lines == ["OPENAI_BASE_URL=http://127.0.0.1:9999/v1"]
def test_cortex_code_build_launch_env_applies_project_prefix() -> None:
env, lines = build_launch_env(port=9999, environ={}, project="myrepo")
assert env["OPENAI_BASE_URL"] == "http://127.0.0.1:9999/p/myrepo/v1"
assert lines == ["OPENAI_BASE_URL=http://127.0.0.1:9999/p/myrepo/v1"]
def test_cortex_code_build_launch_env_ignores_blank_project() -> None:
env, lines = build_launch_env(port=9999, environ={}, project=" ")
assert env["OPENAI_BASE_URL"] == "http://127.0.0.1:9999/v1"
assert lines == ["OPENAI_BASE_URL=http://127.0.0.1:9999/v1"]
def test_cortex_code_render_setup_lines_contains_proxy_url() -> None:
lines = render_setup_lines(8787)
joined = "\n".join(lines)
assert "http://127.0.0.1:8787/v1" in joined
assert "Cortex Code" in joined
def test_cortex_code_render_setup_lines_project_attribution() -> None:
lines = render_setup_lines(8787, project="my-sf-project")
joined = "\n".join(lines)
assert "my-sf-project" in joined
plain = "\n".join(render_setup_lines(8787))
assert "attributed" not in plain
def test_cortex_code_default_api_url_reads_snowflake_host_env() -> None:
url = default_api_url({"SNOWFLAKE_HOST": "myaccount.snowflakecomputing.com"})
assert url == "https://myaccount.snowflakecomputing.com"
def test_cortex_code_default_api_url_constructs_url_from_account_name() -> None:
url = default_api_url({"SNOWFLAKE_ACCOUNT": "myaccount"})
assert url == "https://myaccount.snowflakecomputing.com"
def test_cortex_code_default_api_url_host_takes_priority_over_account() -> None:
url = default_api_url(
{
"SNOWFLAKE_HOST": "host.snowflakecomputing.com",
"SNOWFLAKE_ACCOUNT": "account",
}
)
assert url == "https://host.snowflakecomputing.com"
def test_cortex_code_default_api_url_falls_back_when_no_env() -> None:
url = default_api_url({})
assert url == "https://app.snowflake.com"
def test_cortex_code_default_api_url_preserves_https_prefix() -> None:
url = default_api_url({"SNOWFLAKE_HOST": "https://already.snowflakecomputing.com"})
assert url == "https://already.snowflakecomputing.com"
def test_cortex_code_install_registry_includes_cortex_code() -> None:
from headroom.providers.install_registry import build_install_target_envs
result = build_install_target_envs(port=1234, backend="ignored", targets=["cortex-code"])
assert result["cortex-code"]["OPENAI_BASE_URL"] == "http://127.0.0.1:1234/v1"
def test_cortex_code_install_registry_unknown_target_skipped() -> None:
from headroom.providers.install_registry import build_install_target_envs
result = build_install_target_envs(
port=1234, backend="ignored", targets=["cortex-code", "unknown-tool"]
)
assert "unknown-tool" not in result
assert "cortex-code" in result