Repository navigation
Add UNITY_GATEWAY_CUSTOM_CODING_AGENT header override - #1179
poojabhat0 wants to merge 1 commit into
Conversation
poojabhat0
changed the title
|
/ug-review |
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
| ) | ||
|
|
||
| headers = {"x-databricks-use-coding-agent-mode": "true"} | ||
| headers.update(custom_coding_agent_headers()) |
There was a problem hiding this comment.
other agents add this header per launch. what happens if ug configure runs inside a custom-agent shell? the Desktop profile keeps that name, so all Desktop traffic gets attributed to that agent. can we leave it out for Desktop?
| }[agent] | ||
| monkeypatch.setattr(module, "ug_version", lambda: "1.2.3") | ||
| monkeypatch.setattr(module, "agent_version", lambda _binary: "4.5.6") | ||
| monkeypatch.delenv(UNITY_GATEWAY_CUSTOM_CODING_AGENT_ENV_VAR, raising=False) |
There was a problem hiding this comment.
can we clear this in conftest's autouse _isolate_ucode_state instead? if a dev or CI has the var set, existing exact-header tests for claude/codex/gemini/pi pick up the extra header and fail.
| @pytest.mark.parametrize("identity", ["COLLIDING_ENV_VAR", "!do-not-execute"]) | ||
| def test_pi_uses_environment_reference_not_a_dynamic_literal(identity, monkeypatch): |
There was a problem hiding this comment.
can we add a padded value like " my-agent "? pi expands the raw env var, so it sends the untrimmed value while every other adapter sends the stripped one, and nothing here catches that.
Suggestion:
| @pytest.mark.parametrize("identity", ["COLLIDING_ENV_VAR", "!do-not-execute"]) | |
| def test_pi_uses_environment_reference_not_a_dynamic_literal(identity, monkeypatch): | |
| @pytest.mark.parametrize("identity", ["COLLIDING_ENV_VAR", "!do-not-execute", " padded-agent "]) | |
| def test_pi_uses_environment_reference_not_a_dynamic_literal(identity, monkeypatch): |
Polly AI ReviewBlocking issues
Security vulnerabilities
Non-blocking notes
ApproachThe shared helper and adapter-level wiring are consistent with existing boundaries. Harden validation centrally, and apply optional-attribution compatibility rules consistently to normal and skipped managed-file reconciliation. SummaryDiff size (computed): 16 files; +348 / -2 lines (350 changed text lines); 0 binary files.
Shares use additions + deletions, including generated text and lockfiles. Binary files have no line count. Tests include fixtures/helpers under test paths and colocated test/spec files; categories are path-based, not coverage measurements. ChangesThe change adds environment-based attribution across the requested CLI adapters and Claude Desktop configuration, shared validation, header ownership/conflict handling, documentation, and tests. Authentication and User-Agent generation are not intentionally changed. No dependencies, lockfiles, or generated assets change. The implementation is proportional to the outcome. The concrete reduction opportunities are redundant renderer checks and adapter-level blank-input permutations, assessed below. TestsTest-by-test assessmentAssessment is based on the frozen diff and relevant existing coverage. Offline function-level reproductions confirmed the Unicode injection and conditional Codex reconciliation failures. The affected pytest suite was not run; the PR's reported pass count is not independently verified. Native requests and gateway attribution remain unverified.
Existing Claude caller-header preservation and admin-withdrawal tests, plus Codex legacy cleanup and admin-withdrawal tests, already cover related ownership transitions. They do not cover the two demonstrated failures above. Focused unit/component regressions are sufficient for those failures; duplicate live tests are unnecessary. ScopeAll changed files support the stated attribution outcome. No clearly unrelated changes or added source-comment conciseness issues were found. The findings concern validation and managed-configuration behavior within that scope, not independent feature work. Automated review by Polly · workflow run |
1 participant
Summary
Let custom agents and wrappers identify their inference traffic through
UNITY_GATEWAY_CUSTOM_CODING_AGENT:Databricks-Custom-Coding-Agentwithout changing authentication orUser-Agent. The gateway attributes traffic to the custom identity instead of the native agent;ug configure.No custom-client-surface headers, authentication changes, or support for Copilot/MCP-only agents are included.
Validation
ty check src/passed.