Files
omnigent-ai--omnigent/tests/test_native_policy_hook.py
2026-07-13 13:12:00 +08:00

730 lines
27 KiB
Python

"""Unit tests for the shared native-harness policy hook converters."""
from __future__ import annotations
import httpx
import pytest
from omnigent import native_policy_hook
from omnigent.native_policy_hook import (
_is_login_redirect_or_unauthorized,
evaluation_response_to_hook_output,
fail_closed_hook_output,
hook_payload_to_evaluation_request,
post_evaluate_with_retry,
)
def test_pre_tool_use_maps_to_phase_tool_call() -> None:
"""
A PreToolUse payload becomes a PHASE_TOOL_CALL EvaluationRequest.
The tool name and arguments must land in ``event.data`` so the
server's policy engine can match on them. A failure here means the
server would evaluate an empty/garbled tool call and likely ALLOW
everything.
"""
result = hook_payload_to_evaluation_request(
"PreToolUse",
{"tool_name": "Bash", "tool_input": {"command": "rm -rf /"}},
)
assert result is not None
event = result["event"]
assert event["type"] == "PHASE_TOOL_CALL"
# The command must survive into args, or the policy can't inspect it.
assert event["data"] == {"name": "Bash", "arguments": {"command": "rm -rf /"}}
def test_post_tool_use_maps_to_phase_tool_result() -> None:
"""
A PostToolUse payload becomes a PHASE_TOOL_RESULT EvaluationRequest.
The result text goes in ``event.data.result`` and the originating
tool name/args ride along in ``request_data`` so a TOOL_RESULT
policy can correlate output to the call that produced it. A failure
means output-inspection policies would see no result or no tool.
"""
result = hook_payload_to_evaluation_request(
"PostToolUse",
{
"tool_name": "Bash",
"tool_input": {"command": "cat /etc/passwd"},
"tool_output": "root:x:0:0:...",
},
)
assert result is not None
event = result["event"]
assert event["type"] == "PHASE_TOOL_RESULT"
assert event["data"]["result"] == "root:x:0:0:..."
# request_data carries the originating call so result policies can
# correlate output back to the tool + args that produced it.
assert event["request_data"] == {
"name": "Bash",
"arguments": {"command": "cat /etc/passwd"},
}
@pytest.mark.parametrize("hook_event", ["PreToolUse", "PostToolUse"])
def test_omnigent_mcp_tools_are_skipped(hook_event: str) -> None:
"""
Omnigent MCP tools return None and are never sent to /policies/evaluate.
Omnigent MCP tool calls are already policy-checked by the relay path
(ProxyMcpManager → Omnigent /mcp endpoint → _evaluate_tool_call_policy).
If this guard regressed, every MCP tool call would be evaluated
twice — once via the relay, once via this hook.
"""
result = hook_payload_to_evaluation_request(
hook_event,
{"tool_name": "mcp__omnigent__list_comments", "tool_input": {}, "tool_output": "x"},
)
# None signals the caller to skip the POST entirely.
assert result is None
@pytest.mark.parametrize(
"hook_event,expected_type",
[("PreToolUse", "PHASE_TOOL_CALL"), ("PostToolUse", "PHASE_TOOL_RESULT")],
)
def test_connector_native_mcp_tools_are_evaluated(hook_event: str, expected_type: str) -> None:
"""
Connector-native MCP tools must not be skipped by the native pre-call hook.
Tools such as ``mcp__github__*`` are injected by the connector layer and
do not round-trip through Omnigent's MCP proxy, so this hook is their
TOOL_CALL/TOOL_RESULT policy enforcement site.
"""
result = hook_payload_to_evaluation_request(
hook_event,
{
"tool_name": "mcp__github__create_issue",
"tool_input": {"title": "blocked?"},
"tool_output": "created",
},
)
assert result is not None
event = result["event"]
assert event["type"] == expected_type
if hook_event == "PreToolUse":
assert event["data"] == {
"name": "mcp__github__create_issue",
"arguments": {"title": "blocked?"},
}
else:
assert event["request_data"] == {
"name": "mcp__github__create_issue",
"arguments": {"title": "blocked?"},
}
def test_unknown_hook_event_returns_none() -> None:
"""
A non-tool hook event (e.g. SessionStart) is not policy-relevant.
Returning None makes the hook a no-op for events that carry no tool
call. A failure (returning a request) would POST garbage to the
server for every lifecycle event.
"""
assert hook_payload_to_evaluation_request("SessionStart", {"tool_name": "Bash"}) is None
@pytest.mark.parametrize(
"action,expected_decision",
[
("POLICY_ACTION_DENY", "deny"),
("POLICY_ACTION_ASK", "deny"),
],
)
def test_pre_tool_use_response_maps_action_to_permission_decision(
action: str, expected_decision: str
) -> None:
"""
A constraining proto action maps to the matching permissionDecision.
DENY→deny. ASK→deny too: ASK is resolved server-side now (URL-based
elicitation — ``POST /policies/evaluate`` holds the gate and returns
a hard ALLOW/DENY), so the hook should never see ASK; if it does, it
must fail closed with ``deny`` rather than the old ``defer`` (which
handed control to a possibly-permissive harness permission_mode,
re-opening the bypass). ALLOW is deliberately NOT here — it returns
None (see test_pre_tool_use_allow_returns_none). A wrong mapping here
would, e.g., let a DENY verdict run the tool, defeating enforcement.
"""
output = evaluation_response_to_hook_output("PreToolUse", {"result": action})
assert output is not None
hook_specific = output["hookSpecificOutput"]
assert hook_specific["hookEventName"] == "PreToolUse"
assert hook_specific["permissionDecision"] == expected_decision
def test_pre_tool_use_allow_returns_none() -> None:
"""
A PreToolUse ALLOW yields no opinion (None), not ``"allow"``.
ALLOW is the policy engine's default verdict when no policy matches a
tool call. Emitting ``permissionDecision: "allow"`` would auto-approve
the tool in the native harness, suppressing its own permission prompt
— and, for Claude Code, the ``PermissionRequest`` hook that routes
that prompt to the web UI. Returning None keeps the policy gate and
the user's own consent gate independent: the policy layer may block
(DENY) or demand approval (ASK), but must never silence the harness's
native prompt. Regression guard for "claude-native elicitations stop
showing in the web UI" once a PreToolUse policy hook was wired in.
"""
output = evaluation_response_to_hook_output("PreToolUse", {"result": "POLICY_ACTION_ALLOW"})
assert output is None
def test_pre_tool_use_deny_includes_reason() -> None:
"""
A DENY verdict surfaces the policy reason as permissionDecisionReason.
The reason is what the user/agent sees explaining the block. A
failure (missing reason) would block tools with no explanation.
"""
output = evaluation_response_to_hook_output(
"PreToolUse",
{"result": "POLICY_ACTION_DENY", "reason": "rm blocked by admin policy"},
)
assert output is not None
hook_specific = output["hookSpecificOutput"]
assert hook_specific["permissionDecision"] == "deny"
assert hook_specific["permissionDecisionReason"] == "rm blocked by admin policy"
def test_pre_tool_use_unknown_action_returns_none() -> None:
"""
An unrecognized/unspecified action yields no opinion (None).
POLICY_ACTION_UNSPECIFIED (e.g. no agent / no policies) must not be
coerced into allow or deny — returning None lets the harness apply
its own default. A failure would fabricate a verdict from no policy.
"""
output = evaluation_response_to_hook_output(
"PreToolUse", {"result": "POLICY_ACTION_UNSPECIFIED"}
)
assert output is None
def test_post_tool_use_deny_maps_to_additional_context() -> None:
"""
A PostToolUse DENY becomes an additionalContext warning, not a block.
PostToolUse fires after the tool ran, so it cannot block — the
verdict is surfaced to the model as context. A failure would either
drop the warning or wrongly attempt to block an already-run tool.
"""
output = evaluation_response_to_hook_output(
"PostToolUse",
{"result": "POLICY_ACTION_DENY", "reason": "Sensitive data in output"},
)
assert output is not None
hook_specific = output["hookSpecificOutput"]
assert hook_specific["hookEventName"] == "PostToolUse"
# The warning text must carry the reason so the model sees why.
assert hook_specific["additionalContext"] == "[Policy violation] Sensitive data in output"
def test_post_tool_use_allow_returns_none() -> None:
"""
A PostToolUse ALLOW produces no output (nothing to inject).
Only DENY warrants an additionalContext warning. A failure
(emitting output on ALLOW) would spam the model with empty context
on every successful tool result.
"""
output = evaluation_response_to_hook_output("PostToolUse", {"result": "POLICY_ACTION_ALLOW"})
assert output is None
def test_user_prompt_submit_maps_to_phase_request() -> None:
"""
A UserPromptSubmit payload becomes a PHASE_REQUEST EvaluationRequest.
The prompt text must land in ``event.data.text`` because the server's
``_build_evaluation_context`` reads REQUEST content from ``data.text``
(falling back to ``data.content``). If the prompt were dropped, the
request-phase gate would evaluate empty content and ALLOW everything.
"""
result = hook_payload_to_evaluation_request(
"UserPromptSubmit",
{"prompt": "delete the prod database"},
)
assert result is not None
event = result["event"]
assert event["type"] == "PHASE_REQUEST"
assert event["data"] == {"text": "delete the prod database"}
# A context dict must exist so the per-harness hook can stamp model/harness.
assert event["context"] == {}
def test_user_prompt_submit_missing_prompt_yields_empty_text() -> None:
"""
A UserPromptSubmit payload with no ``prompt`` still produces a request.
The text falls back to an empty string rather than ``None`` so the
server always receives a well-formed REQUEST event.
"""
result = hook_payload_to_evaluation_request("UserPromptSubmit", {})
assert result is not None
assert result["event"]["data"] == {"text": ""}
@pytest.mark.parametrize("action", ["POLICY_ACTION_DENY", "POLICY_ACTION_ASK"])
def test_user_prompt_submit_blocking_actions_emit_decision_block(action: str) -> None:
"""
DENY (and a stray ASK) block the prompt via top-level ``decision``.
UserPromptSubmit uses the top-level ``decision`` / ``reason`` contract
(NOT ``permissionDecision``) — both harnesses parse ``decision: "block"``
to drop the prompt before the model sees it. ASK is meant to be resolved
server-side (``_hold_native_ask_gate``), so if the hook ever sees it, it
must fail closed by blocking rather than letting the prompt through.
"""
output = evaluation_response_to_hook_output(
"UserPromptSubmit",
{"result": action, "reason": "no prod mutations"},
)
assert output is not None
# Top-level decision/reason, not hookSpecificOutput.permissionDecision.
assert output == {"decision": "block", "reason": "no prod mutations"}
def test_user_prompt_submit_block_defaults_reason() -> None:
"""
A block with no reason still carries a non-empty reason.
Both harnesses drop a block whose reason is empty (the block is treated
as invalid), so a missing reason must be defaulted or the gate would
silently fail open.
"""
output = evaluation_response_to_hook_output(
"UserPromptSubmit", {"result": "POLICY_ACTION_DENY"}
)
assert output == {"decision": "block", "reason": "Denied by policy"}
@pytest.mark.parametrize("action", ["POLICY_ACTION_ALLOW", "POLICY_ACTION_UNSPECIFIED"])
def test_user_prompt_submit_non_blocking_actions_return_none(action: str) -> None:
"""
ALLOW and the no-match default proceed with no output.
Returning None lets the prompt reach the model. Unlike PreToolUse there
is no separate user-consent gate on a prompt to preserve, so ALLOW need
not emit anything.
"""
output = evaluation_response_to_hook_output("UserPromptSubmit", {"result": action})
assert output is None
def test_fail_closed_pre_tool_use_denies() -> None:
"""
An unobtainable verdict on PreToolUse fails CLOSED with ``deny``.
PreToolUse is the authoritative pre-execution gate for native tools —
the sole enforcement point for connector-native ``mcp__*`` tools and
native Bash/Write/Edit — so a verdict that cannot be fetched must deny
rather than silently let the call through (issue #536).
"""
output = fail_closed_hook_output("PreToolUse")
assert output is not None
hook_specific = output["hookSpecificOutput"]
assert hook_specific["hookEventName"] == "PreToolUse"
assert hook_specific["permissionDecision"] == "deny"
# A deny is inert without a reason on the consuming harnesses, so one
# must always be present.
assert hook_specific["permissionDecisionReason"]
def test_fail_closed_user_prompt_submit_fails_closed() -> None:
"""
``UserPromptSubmit`` (``PHASE_REQUEST``) fails CLOSED on an unobtainable verdict.
The request gate is the sole pre-turn enforcement point for native
sessions — a server hiccup must not let an over-budget or otherwise-
blocked request proceed. The output must be a top-level
``decision: "block"`` with a non-empty reason (both Claude Code and
Codex drop a block with an empty reason).
"""
output = fail_closed_hook_output("UserPromptSubmit")
assert output is not None
assert output["decision"] == "block"
assert output["reason"]
def test_fail_closed_post_tool_use_fails_open() -> None:
"""
``PostToolUse`` (``PHASE_TOOL_RESULT``) fails OPEN (``None``).
The tool has already executed by this point, so denying would only
block an already-incurred side effect. Mirrors the runner-side
``FAIL_CLOSED_PHASES``.
"""
assert fail_closed_hook_output("PostToolUse") is None
def test_fail_closed_unknown_event_fails_open() -> None:
"""
An unrecognized hook event fails OPEN (``None``), not closed.
Only the exact ``PreToolUse`` event denies; any novel event name added
by a future harness must fall through to "no opinion" rather than
accidentally blocking — the conservative default for an unknown gate.
"""
assert fail_closed_hook_output("SomeNewEvent") is None
def _resp(status: int, location: str | None = None) -> httpx.Response:
"""Build a fake response for re-auth classification tests."""
headers = {"Location": location} if location else {}
return httpx.Response(status, headers=headers, request=httpx.Request("POST", "https://ap/x"))
@pytest.mark.parametrize(
("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.
(_resp(302, "https://w.example.com/some/other/page"), False),
(_resp(302, None), False),
(_resp(200), False),
(_resp(503), False),
],
)
def test_is_login_redirect_or_unauthorized_classifies_reauth_signals(
response: httpx.Response, expected: bool
) -> None:
"""
401 and an Apps OAuth-login 302 are re-auth signals; nothing else is.
The Databricks Apps front door bounces an *expired* bearer with a
``302 → /oidc/`` (or ``/.auth/``), NOT a ``401`` — so a hook that only
checked ``401`` silently failed closed once the one-shot token lapsed.
This is the classifier that lets the hook re-mint instead.
"""
assert _is_login_redirect_or_unauthorized(response) is expected
def _make_redirect_then_ok_client(
seen_headers: list[dict[str, str]],
*,
redirect: httpx.Response,
ok: httpx.Response,
) -> type:
"""Build an httpx.Client stub: redirect on attempt 1, ``ok`` thereafter."""
class _Client:
def __init__(self, *, headers: dict[str, str], timeout: object) -> None:
del timeout
self._headers = headers
def __enter__(self) -> _Client:
return self
def __exit__(self, *args: object) -> None:
del args
def post(self, url: str, *, json: dict[str, object]) -> httpx.Response:
del url, json
seen_headers.append(dict(self._headers))
return redirect if len(seen_headers) == 1 else ok
return _Client
def test_post_evaluate_with_retry_reauths_on_login_redirect(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""
A 302→/oidc/ re-mints the bearer and retries, returning the real verdict.
Regression guard for the production bug where an "old" native session
(token past the ~1h Databricks OAuth lifetime) failed CLOSED on every tool
call. The first attempt carries the lapsed token (302), the retry carries
the fresh token and gets the ALLOW verdict — exactly as the runner's
refresh-capable ``_RunnerDatabricksAuth`` does for its own callbacks.
"""
seen_headers: list[dict[str, str]] = []
redirect = httpx.Response(
302,
headers={"Location": "https://w.example.com/oidc/oauth2/v2.0/authorize"},
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=redirect, ok=ok),
)
reauth_calls: list[int] = []
def _reauth() -> dict[str, str]:
reauth_calls.append(1)
return {"Authorization": "Bearer fresh", "X-Databricks-Org-Id": "o1"}
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] # re-minted exactly once
assert seen_headers[0]["Authorization"] == "Bearer stale" # first attempt: lapsed token
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:
"""
With no ``reauth`` callable, a login-redirect yields ``None`` (legacy).
Callers without a token source (e.g. codex/kimi today) see the same
behavior as before this change: ``raise_for_status`` rejects the 302 as a
non-retryable <500, the helper returns ``None``, and the caller fails
closed. Guards against the new branch altering that.
"""
seen_headers: list[dict[str, str]] = []
redirect = httpx.Response(
302,
headers={"Location": "https://w.example.com/oidc/x"},
request=httpx.Request("POST", "https://ap/x"),
)
monkeypatch.setattr(
native_policy_hook.httpx,
"Client",
_make_redirect_then_ok_client(seen_headers, redirect=redirect, ok=redirect),
)
resp, error = post_evaluate_with_retry(
"https://ap/x", {}, {"event": {}}, 5.0, "evaluate-policy hook"
)
assert resp is None
assert error is not None
assert len(seen_headers) == 1 # one attempt; a 302 is not retried without reauth
def test_post_evaluate_with_retry_reauth_unavailable_fails_closed(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""
When re-mint yields no token, the helper returns ``None`` (caller fails closed).
Re-auth is best-effort: a ``reauth`` that returns ``None`` (no creds /
transient mint failure) must not loop — it falls through to
``raise_for_status`` (302 → non-retryable) so the caller keeps the
fail-closed safety net.
"""
seen_headers: list[dict[str, str]] = []
redirect = httpx.Response(
302,
headers={"Location": "https://w.example.com/oidc/x"},
request=httpx.Request("POST", "https://ap/x"),
)
monkeypatch.setattr(
native_policy_hook.httpx,
"Client",
_make_redirect_then_ok_client(seen_headers, redirect=redirect, ok=redirect),
)
resp, error = post_evaluate_with_retry(
"https://ap/x",
{"Authorization": "Bearer stale"},
{"event": {}},
5.0,
"evaluate-policy hook",
reauth=lambda: None,
)
assert resp is None
assert error is not None
assert len(seen_headers) == 1 # one attempt only; no retry loop
# ── shared policy-hook header plumbing (writer + reader) ─────────────
def test_policy_hook_request_headers_merges_baked_auth(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""The reader merges the executor-baked auth + routing headers.
The import-free hook subprocess can't resolve credentials in-process, so
the executor bakes them into ``_OMNIGENT_AUTH_HEADERS``; the reader must
fold them onto ``Content-Type`` for the policy POST.
"""
monkeypatch.setenv(
"_OMNIGENT_AUTH_HEADERS",
'{"Authorization": "Bearer tok", "X-Databricks-Org-Id": "org123"}',
)
assert native_policy_hook.policy_hook_request_headers() == {
"Content-Type": "application/json",
"Authorization": "Bearer tok",
"X-Databricks-Org-Id": "org123",
}
@pytest.mark.parametrize("raw", ["", "not json", "[1,2]"])
def test_policy_hook_request_headers_tolerates_missing_or_bad_env(
monkeypatch: pytest.MonkeyPatch, raw: str
) -> None:
"""Absent / malformed env → just ``Content-Type`` (local-unauth path).
A bad value must not crash the hook nor inject garbage headers — the
server simply decides without auth (a local unauthenticated server needs
none).
"""
if raw:
monkeypatch.setenv("_OMNIGENT_AUTH_HEADERS", raw)
else:
monkeypatch.delenv("_OMNIGENT_AUTH_HEADERS", raising=False)
assert native_policy_hook.policy_hook_request_headers() == {"Content-Type": "application/json"}
def test_policy_hook_wrapper_script_bakes_auth_and_routing(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""The writer bakes bearer + routing into the wrapper via the one builder.
A new harness wiring up its hook through this helper gets auth AND
workspace routing for free — the gap that left the cursor/hermes hooks
posting unauthenticated and unrouted.
"""
import omnigent.cli_auth as cli_auth
import omnigent.runner._entry as entry
monkeypatch.setattr(
entry, "_make_auth_token_factory", lambda *, server_url=None: lambda: "tok"
)
monkeypatch.setattr(cli_auth, "load_databricks_org_id", lambda _url: "org123")
script = native_policy_hook.policy_hook_wrapper_script(
"https://acme.databricks.com/api/2.0/omnigent", "conv_x", "/path/hook.py"
)
assert script.startswith("#!/bin/sh\n")
assert "_OMNIGENT_SERVER_URL=https://acme.databricks.com/api/2.0/omnigent" in script
assert "_OMNIGENT_SESSION_ID=conv_x" in script
# The baked headers carry BOTH the bearer and the routing header.
line = next(
ln for ln in script.splitlines() if ln.startswith("export _OMNIGENT_AUTH_HEADERS=")
)
assert "Bearer tok" in line
assert "X-Databricks-Org-Id" in line and "org123" in line
def test_policy_hook_wrapper_script_omits_auth_when_unauthenticated(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""No token + no recorded selector → empty auth dict (local-unauth runs).
The wrapper still exports the (empty) header dict, so the reader yields
just ``Content-Type`` — non-workspace callers are unaffected.
"""
import omnigent.cli_auth as cli_auth
import omnigent.runner._entry as entry
monkeypatch.setattr(entry, "_make_auth_token_factory", lambda *, server_url=None: None)
monkeypatch.setattr(cli_auth, "load_databricks_org_id", lambda _url: None)
script = native_policy_hook.policy_hook_wrapper_script(
"http://127.0.0.1:6767", "conv_local", "/path/hook.py"
)
line = next(
ln for ln in script.splitlines() if ln.startswith("export _OMNIGENT_AUTH_HEADERS=")
)
assert "Bearer" not in line
assert "X-Databricks-Org-Id" not in line
def test_policy_hook_reauth_remints_and_preserves_routing_header(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""The shared reauth re-mints the bearer and keeps the routing header.
When a baked one-shot token lapses, the four hooks (codex/kimi/cursor/
hermes) that wire this builder must mint a fresh bearer through the same
factory the refresh-capable runtime auth uses — without dropping the
workspace-routing header that travels alongside it.
"""
monkeypatch.setattr(
"omnigent.runner._entry._make_auth_token_factory",
lambda _server_url: lambda: "fresh-token",
)
reauth = native_policy_hook.policy_hook_reauth(
"https://acme.databricks.com/api/2.0/omnigent",
{"Authorization": "Bearer stale", "X-Databricks-Org-Id": "o9"},
)
assert reauth() == {"Authorization": "Bearer fresh-token", "X-Databricks-Org-Id": "o9"}
def test_policy_hook_reauth_returns_none_without_factory(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""No refresh mechanism (local/unauth) → ``None`` so the caller fails closed."""
monkeypatch.setattr(
"omnigent.runner._entry._make_auth_token_factory",
lambda _server_url: None,
)
reauth = native_policy_hook.policy_hook_reauth(
"http://127.0.0.1:6767", {"Content-Type": "application/json"}
)
assert reauth() is None