diff --git a/omnigent/native_policy_hook.py b/omnigent/native_policy_hook.py index cdb1b77c..b0533a28 100644 --- a/omnigent/native_policy_hook.py +++ b/omnigent/native_policy_hook.py @@ -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 diff --git a/omnigent/runner/_entry.py b/omnigent/runner/_entry.py index f4d479fe..529b86b1 100644 --- a/omnigent/runner/_entry.py +++ b/omnigent/runner/_entry.py @@ -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 diff --git a/tests/test_claude_native_hook.py b/tests/test_claude_native_hook.py index 5f15e3ef..2f80bad5 100644 --- a/tests/test_claude_native_hook.py +++ b/tests/test_claude_native_hook.py @@ -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, diff --git a/tests/test_native_policy_hook.py b/tests/test_native_policy_hook.py index 4d5b55ff..76eaef8b 100644 --- a/tests/test_native_policy_hook.py +++ b/tests/test_native_policy_hook.py @@ -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: