730 lines
27 KiB
Python
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
|