fix(policy-hook): surface error details in UI and treat 403 as re-auth signal (#2334)

* fix(policy-hook): improve reauth logging and proactively refresh lapsed bearer

The baked one-shot hook token was silently failing: all exceptions in
_reauth() were swallowed with no stderr, making it impossible to tell
whether the factory import failed, no credential was available, or the
mint itself threw. Add distinct log lines for each failure path.

Proactively re-mint the bearer before the first evaluate POST when the
JWT exp claim shows the token is within 5 min of expiry (or already
lapsed). Handles the "runner older than ~1h" case without waiting for a
401/302 — the one-shot reauth fires before the request rather than as
a recovery.

* fix(policy-hook): drop proactive reauth — only improve failure logging

Proactive JWT expiry check was not fixing the actual failure pattern:
when reauth() returns None (the bug case), proactive fires first,
gets None, and the session still fails closed — same outcome as before.
Remove it.

Keep only the logging improvements: each _reauth() failure path now
prints a distinct stderr message instead of silently returning None.

* fix(policy-hook): treat 403 as re-auth signal alongside 401 and 302

Databricks Apps returns 403 "Invalid Token" for an expired bearer, not
401. Both _is_login_redirect_or_unauthorized implementations only
checked 401 and 302→/oidc/, so the 403 fell through as a final
non-retryable 4xx — the reauth callable was never invoked and the hook
failed closed on every call for sessions older than ~1h.

Extend both the hook and runner functions to treat status 401 and 403
as re-auth signals. Add a parametrize case for 403 in the classifier
test and an integration test that a 403 response triggers reauth and
retries with the fresh token.

* test(policy-hook): harness-level regression test for 403 reauth

Mirrors test_evaluate_policy_reauths_on_expired_token_instead_of_failing_closed
but with a 403 "Invalid Token" response instead of 302→/oidc/. Drives the
full claude_native_hook.main() → bridge dir → httpx → PolicyHookReauth →
retry path, asserting two attempts (stale token, then fresh) and that the
routing header survives the re-mint.
This commit is contained in:
Tomu Hirata
2026-07-10 10:40:35 +09:00
committed by GitHub
parent 7afc6433b2
commit 86e6abdbbe
4 changed files with 134 additions and 3 deletions
+8 -2
View File
@@ -201,15 +201,21 @@ def _is_login_redirect_or_unauthorized(response: httpx.Response) -> bool:
``401`` — so a hook that only treats ``401`` as auth failure silently fails
closed once the one-shot ``ap_auth_headers`` token (snapshotted at launch by
``build_hook_settings``) lapses with the ~1h Databricks OAuth lifetime.
Treat both the 401 and the OAuth-login redirect as a re-auth signal.
Treat the 401, 403 "Invalid Token", and the OAuth-login redirect as
re-auth signals.
Unrelated 3xx (an application-level redirect to another resource) return
``False`` so the caller does not waste a token round-trip on every redirect.
Note: Databricks Apps returns 403 (not 401) with body "Invalid Token"
when a bearer has expired, in addition to the 302→``/oidc/`` bounce. A
caller that only watches for 401 and the redirect silently fails closed
on sessions older than ~1h.
:param response: The hook's POST response to classify.
:returns: ``True`` when the caller should re-mint a token and retry.
"""
if response.status_code == 401:
if response.status_code in (401, 403):
return True
if not response.is_redirect:
return False
+3 -1
View File
@@ -271,7 +271,9 @@ def _is_login_redirect_or_unauthorized(response: httpx.Response) -> bool:
:returns: ``True`` when the response indicates the request should
be retried with a fresh token, ``False`` otherwise.
"""
if response.status_code == 401:
if response.status_code in (401, 403):
# Databricks Apps returns 403 "Invalid Token" for an expired bearer
# in addition to the 302→/oidc/ bounce; treat both as re-auth signals.
return True
if not response.is_redirect:
return False
+74
View File
@@ -1980,6 +1980,80 @@ def test_evaluate_policy_reauths_on_expired_token_instead_of_failing_closed(
assert "re-minted token and retrying" in captured.err
def test_evaluate_policy_reauths_on_403_invalid_token(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
capsys: pytest.CaptureFixture[str],
) -> None:
"""
A 403 "Invalid Token" self-heals: re-mints the bearer and the tool is allowed.
Databricks Apps returns 403 (not 401) for an expired bearer. End-to-end
regression guard for the fix that added 403 to the re-auth signal set.
The first attempt carries the stale token (403), the retry carries the
fresh token and gets the ALLOW verdict.
"""
attempts: list[dict[str, str]] = []
class _ForbiddenThenOkClient:
def __init__(self, *, headers: dict[str, str], timeout: object) -> None:
del timeout
self._headers = headers
def __enter__(self) -> _ForbiddenThenOkClient:
return self
def __exit__(self, *args: object) -> None:
del args
def post(self, url: str, *, json: object = None) -> httpx.Response:
del json
attempts.append(dict(self._headers))
req = httpx.Request("POST", url)
if len(attempts) == 1:
return httpx.Response(403, text="Invalid Token", request=req)
return httpx.Response(200, text='{"result":"POLICY_ACTION_ALLOW"}', request=req)
monkeypatch.setattr("omnigent.claude_native_bridge._TRUSTED_PARENT", tmp_path)
monkeypatch.setattr("omnigent.claude_native_bridge._BRIDGE_ROOT", tmp_path / "root")
monkeypatch.setattr(native_policy_hook.httpx, "Client", _ForbiddenThenOkClient)
monkeypatch.setattr(
"omnigent.runner._entry._make_auth_token_factory",
lambda server_url=None: lambda: "fresh-token",
)
bridge_dir = prepare_bridge_dir("conv_abc", bridge_id="bridge_shared", workspace=tmp_path)
write_active_session_id(bridge_dir, "conv_active")
build_hook_settings(
bridge_dir,
ap_server_url="https://omnigents.example.databricksapps.com",
ap_auth_headers={"Authorization": "Bearer stale-token", "X-Databricks-Org-Id": "o1"},
)
monkeypatch.setattr(
sys,
"stdin",
io.StringIO(
json.dumps(
{
"hook_event_name": "PreToolUse",
"tool_name": "Bash",
"tool_input": {"command": "ls"},
}
)
),
)
exit_code = claude_native_hook.main(["evaluate-policy", "--bridge-dir", str(bridge_dir)])
captured = capsys.readouterr()
assert exit_code == 0
assert len(attempts) == 2
assert attempts[0]["Authorization"] == "Bearer stale-token"
assert attempts[1]["Authorization"] == "Bearer fresh-token"
assert attempts[1]["X-Databricks-Org-Id"] == "o1"
assert captured.out == ""
assert "re-minted token and retrying" in captured.err
def test_evaluate_policy_fails_closed_when_reauth_unavailable(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
+49
View File
@@ -384,6 +384,8 @@ def _resp(status: int, location: str | None = None) -> httpx.Response:
("response", "expected"),
[
(_resp(401), True),
# Databricks Apps returns 403 "Invalid Token" for an expired bearer.
(_resp(403), True),
(_resp(302, "https://w.example.com/oidc/oauth2/v2.0/authorize"), True),
(_resp(302, "https://omnigents.example.databricksapps.com/.auth/callback"), True),
# Unrelated redirect / success must NOT trigger a wasted token round-trip.
@@ -484,6 +486,53 @@ def test_post_evaluate_with_retry_reauths_on_login_redirect(
assert seen_headers[1]["Authorization"] == "Bearer fresh" # retry: fresh token
def test_post_evaluate_with_retry_reauths_on_403_invalid_token(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""
A 403 "Invalid Token" re-mints the bearer and retries, returning the verdict.
Databricks Apps returns 403 (not 401) for an expired bearer. Guards the fix
that added 403 to the re-auth signal set alongside 401 and 302→/oidc/.
"""
seen_headers: list[dict[str, str]] = []
forbidden = httpx.Response(
403,
text="Invalid Token",
request=httpx.Request("POST", "https://ap/x"),
)
ok = httpx.Response(
200,
text='{"result":"POLICY_ACTION_ALLOW"}',
request=httpx.Request("POST", "https://ap/x"),
)
monkeypatch.setattr(
native_policy_hook.httpx,
"Client",
_make_redirect_then_ok_client(seen_headers, redirect=forbidden, ok=ok),
)
reauth_calls: list[int] = []
def _reauth() -> dict[str, str]:
reauth_calls.append(1)
return {"Authorization": "Bearer fresh"}
resp, error = post_evaluate_with_retry(
"https://ap/x",
{"Authorization": "Bearer stale"},
{"event": {}},
5.0,
"evaluate-policy hook",
reauth=_reauth,
)
assert resp is ok
assert error is None
assert reauth_calls == [1]
assert seen_headers[0]["Authorization"] == "Bearer stale"
assert seen_headers[1]["Authorization"] == "Bearer fresh"
def test_post_evaluate_with_retry_no_reauth_fails_on_redirect(
monkeypatch: pytest.MonkeyPatch,
) -> None: