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:
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user