wehub-resource-sync
611 lines
28 KiB
Python
611 lines
28 KiB
Python
"""Shared conversion between native-harness hooks and Omnigent policy events.
|
|
|
|
Both Claude Code and Codex expose a command-hook system whose
|
|
``PreToolUse`` / ``PostToolUse`` payloads use the same field names
|
|
(``hook_event_name``, ``tool_name``, ``tool_input``, ``tool_output``)
|
|
and whose ``UserPromptSubmit`` payload carries the user prompt under
|
|
``prompt``. This module owns the harness-neutral translation between
|
|
that hook shape and the server's proto-compatible ``EvaluationRequest``
|
|
/ ``EvaluationResponse`` schema served by
|
|
``POST /v1/sessions/{id}/policies/evaluate``, so the per-harness hook
|
|
entrypoints (:mod:`omnigent.claude_native_hook`,
|
|
:mod:`omnigent.codex_native_hook`) share one implementation.
|
|
|
|
The output contract differs by hook event: ``PreToolUse`` enforces via
|
|
``hookSpecificOutput.permissionDecision``, while ``UserPromptSubmit``
|
|
enforces via the top-level ``decision`` / ``reason`` fields (both
|
|
harnesses parse ``decision: "block"`` to drop the prompt before the
|
|
model sees it).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import secrets
|
|
import shlex
|
|
import sys
|
|
import time
|
|
from collections.abc import Callable
|
|
|
|
import httpx
|
|
|
|
# How long to keep retrying transient 5xx / connect errors on the
|
|
# policy evaluate POST before failing closed. Keeps the pre-execution
|
|
# gate from blocking long on a sick server while still absorbing brief
|
|
# DB hiccups on a hosted deployment.
|
|
_EVALUATE_POLICY_RETRY_BUDGET_S = 30.0
|
|
_EVALUATE_POLICY_RETRY_INITIAL_BACKOFF_S = 1.0
|
|
_EVALUATE_POLICY_RETRY_MAX_BACKOFF_S = 10.0
|
|
# Fast connect budget so an unreachable server fails into the retry
|
|
# loop quickly rather than blocking on the day-long read timeout.
|
|
_EVALUATE_POLICY_CONNECT_TIMEOUT_S = 5.0
|
|
|
|
# Hook event names that gate tool execution and therefore carry policy
|
|
# meaning. ``PreToolUse`` fires before the tool runs (can block);
|
|
# ``PostToolUse`` fires after (observational — can only warn).
|
|
_PRE_TOOL_USE = "PreToolUse"
|
|
_POST_TOOL_USE = "PostToolUse"
|
|
# ``UserPromptSubmit`` fires when a new user prompt reaches the harness —
|
|
# for native sessions this is the request-phase gate (the server-level
|
|
# ``_evaluate_input_policy`` is bypassed for native message events, so
|
|
# this hook is the sole REQUEST gate and covers both web-UI-injected and
|
|
# direct-terminal prompts). It can block the prompt before the model runs.
|
|
_USER_PROMPT_SUBMIT = "UserPromptSubmit"
|
|
|
|
# Reason surfaced when a tool call is denied because its policy verdict
|
|
# could not be obtained (server unreachable / non-2xx / empty or malformed
|
|
# body). Mirrors the runner-side fail-closed default in
|
|
# ``omnigent.runner.app._evaluate_policy_via_omnigent``.
|
|
_EVAL_UNAVAILABLE_REASON = (
|
|
"Omnigent policy evaluation unavailable (could not reach or authenticate to the "
|
|
"Omnigent server); failing closed for this tool call."
|
|
)
|
|
_EVAL_UNAVAILABLE_REQUEST_REASON = (
|
|
"Omnigent policy evaluation unavailable (could not reach or authenticate to the "
|
|
"Omnigent server); failing closed for this request."
|
|
)
|
|
|
|
|
|
# Env var carrying the one-shot auth + workspace-routing headers from the
|
|
# executor (which writes the hook wrapper) to the hook subprocess. The hook
|
|
# is import-free of the runner, so the headers are passed in rather than
|
|
# resolved in-process — the same reason ``ap_auth_headers`` is baked for the
|
|
# claude/codex/kimi hooks.
|
|
_AUTH_HEADERS_ENV = "_OMNIGENT_AUTH_HEADERS"
|
|
|
|
|
|
def policy_hook_request_headers() -> dict[str, str]:
|
|
"""Build the headers for a policy-hook subprocess's POST to the server.
|
|
|
|
Always carries ``Content-Type``, and merges the one-shot auth +
|
|
workspace-routing headers the executor baked into :data:`_AUTH_HEADERS_ENV`
|
|
at launch (see :func:`policy_hook_wrapper_script`). Without them the POST
|
|
is unauthenticated and unrouted — it 401s on an authenticated server and
|
|
misroutes to the account on a unified-account workspace. Missing or
|
|
malformed env → just ``Content-Type`` (a local unauthenticated server
|
|
needs no auth).
|
|
|
|
:returns: Request headers for ``post_evaluate_with_retry``.
|
|
"""
|
|
headers = {"Content-Type": "application/json"}
|
|
raw = os.environ.get(_AUTH_HEADERS_ENV, "")
|
|
if raw:
|
|
try:
|
|
extra = json.loads(raw)
|
|
except (json.JSONDecodeError, ValueError):
|
|
extra = None
|
|
if isinstance(extra, dict):
|
|
headers.update({str(k): str(v) for k, v in extra.items()})
|
|
return headers
|
|
|
|
|
|
def policy_hook_wrapper_script(server_url: str, session_id: str, hook_script_path: str) -> str:
|
|
"""Build the ``/bin/sh`` wrapper a native policy hook is launched as.
|
|
|
|
Resolves a one-shot Omnigent-server token and bakes the auth +
|
|
workspace-routing headers (via
|
|
:func:`omnigent.cli_auth.databricks_request_headers`) into
|
|
:data:`_AUTH_HEADERS_ENV`, so the hook's POST authenticates and routes to
|
|
the workspace. The token is a secret, so callers MUST write the returned
|
|
wrapper ``0o700`` (owner-only) — it is never world-readable.
|
|
|
|
:param server_url: Omnigent server base URL the hook posts to.
|
|
:param session_id: Session / conversation id for policy evaluation.
|
|
:param hook_script_path: Absolute path to the hook's Python entrypoint.
|
|
:returns: Shell-script text for the wrapper (write it ``0o700``).
|
|
"""
|
|
from omnigent.cli_auth import databricks_request_headers
|
|
from omnigent.runner._entry import _make_auth_token_factory
|
|
|
|
factory = _make_auth_token_factory(server_url=server_url)
|
|
token = factory() if factory is not None else None
|
|
auth_headers = databricks_request_headers(server_url, bearer_token=token)
|
|
return (
|
|
"#!/bin/sh\n"
|
|
f"export _OMNIGENT_SERVER_URL={shlex.quote(server_url)}\n"
|
|
f"export _OMNIGENT_SESSION_ID={shlex.quote(session_id)}\n"
|
|
f"export {_AUTH_HEADERS_ENV}={shlex.quote(json.dumps(auth_headers))}\n"
|
|
f"exec {shlex.quote(sys.executable)} {shlex.quote(hook_script_path)}\n"
|
|
)
|
|
|
|
|
|
class PolicyHookReauth:
|
|
"""Callable that re-mints the Omnigent bearer for a policy hook subprocess.
|
|
|
|
The baked one-shot token dies with the ~1h Databricks OAuth lifetime; on a
|
|
lapsed-token signal (401 or Apps ``302→/oidc/``) ``post_evaluate_with_retry``
|
|
calls this once to mint a fresh bearer through the same factory the
|
|
refresh-capable runtime auth uses, keeping the other headers (e.g.
|
|
``X-Databricks-Org-Id``) so routing survives.
|
|
|
|
The ``failure_reason`` attribute is set to a short diagnostic string when
|
|
the re-mint fails so callers can surface it in the fail-closed message shown
|
|
to the user — stderr from hook subprocesses is discarded by the harness, so
|
|
this is the only channel that reaches the UI.
|
|
"""
|
|
|
|
failure_reason: str | None
|
|
|
|
def __init__(self, server_url: str, headers: dict[str, str]) -> None:
|
|
self._server_url = server_url
|
|
self._headers = headers
|
|
self.failure_reason = None
|
|
|
|
def __call__(self) -> dict[str, str] | None:
|
|
# Lazy import: paid only on the rare re-auth path, off the hot path.
|
|
try:
|
|
from omnigent.runner._entry import _make_auth_token_factory
|
|
except Exception as exc: # noqa: BLE001 — best-effort; fail closed if unavailable
|
|
self.failure_reason = f"auth factory unavailable: {exc}"
|
|
return None
|
|
factory = _make_auth_token_factory(self._server_url)
|
|
if factory is None:
|
|
self.failure_reason = (
|
|
"no credential resolved "
|
|
f"(no stored token and no Databricks SDK auth for {self._server_url!r})"
|
|
)
|
|
return None
|
|
try:
|
|
token = factory()
|
|
except Exception as exc: # noqa: BLE001 — transient mint failure; fail closed
|
|
self.failure_reason = f"token mint failed: {exc}"
|
|
return None
|
|
if not token:
|
|
self.failure_reason = "auth factory returned empty token"
|
|
return None
|
|
self.failure_reason = None
|
|
return {**self._headers, "Authorization": f"Bearer {token}"}
|
|
|
|
|
|
def policy_hook_reauth(server_url: str, headers: dict[str, str]) -> PolicyHookReauth:
|
|
"""Build a :class:`PolicyHookReauth` callable for *server_url*.
|
|
|
|
:param server_url: Omnigent server base URL the hook POSTs to.
|
|
:param headers: Current (lapsed) headers; the fresh bearer is merged over
|
|
a copy so routing headers survive.
|
|
:returns: A :class:`PolicyHookReauth` instance. Call it to attempt a
|
|
re-mint; check ``.failure_reason`` afterwards when it returns ``None``.
|
|
"""
|
|
return PolicyHookReauth(server_url, headers)
|
|
|
|
|
|
def _is_login_redirect_or_unauthorized(response: httpx.Response) -> bool:
|
|
"""
|
|
Return ``True`` when a response means "the bearer is no good — re-auth".
|
|
|
|
Mirrors :func:`omnigent.runner._entry._is_login_redirect_or_unauthorized`,
|
|
duplicated here so the dependency-light hook need not import the runner
|
|
package. The Databricks Apps front door bounces an *expired* bearer with a
|
|
``302`` to the OAuth login flow (``/oidc/`` or ``/.auth/``) — **not** a
|
|
``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 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 in (401, 403):
|
|
return True
|
|
if not response.is_redirect:
|
|
return False
|
|
location = response.headers.get("location", "")
|
|
return "/oidc/" in location or "/.auth/" in location
|
|
|
|
|
|
def hook_payload_to_evaluation_request(
|
|
hook_event: str,
|
|
payload: dict[str, object],
|
|
) -> dict[str, object] | None:
|
|
"""
|
|
Convert a native-harness tool-hook payload into a proto ``EvaluationRequest``.
|
|
|
|
Maps ``PreToolUse`` to a ``PHASE_TOOL_CALL`` event, ``PostToolUse``
|
|
to a ``PHASE_TOOL_RESULT`` event, and ``UserPromptSubmit`` to a
|
|
``PHASE_REQUEST`` event (the prompt text from the payload's
|
|
``prompt`` field becomes the request content). Omnigent MCP tools
|
|
(``mcp__omnigent__*``) are skipped because they are already
|
|
policy-checked by the relay path (``ProxyMcpManager`` → Omnigent
|
|
``/mcp`` endpoint → ``_evaluate_tool_call_policy``); evaluating
|
|
them here would double-count. Connector-native MCP tools
|
|
(for example ``mcp__github__*``) still need this pre-call gate.
|
|
|
|
:param hook_event: Hook event name from the payload's
|
|
``hook_event_name`` field, e.g. ``"PreToolUse"``,
|
|
``"PostToolUse"``, or ``"UserPromptSubmit"``.
|
|
:param payload: Raw hook JSON from the harness, e.g.
|
|
``{"hook_event_name": "PreToolUse", "tool_name": "Bash",
|
|
"tool_input": {"command": "rm -rf /"}}``.
|
|
:returns: An ``EvaluationRequest`` dict suitable for POSTing to
|
|
``/policies/evaluate``, or ``None`` when the event is not
|
|
policy-relevant (unknown event or an ``mcp__omnigent__*`` tool).
|
|
"""
|
|
if hook_event == _USER_PROMPT_SUBMIT:
|
|
# Request-phase gate for native sessions. The server reads REQUEST
|
|
# content from ``data.text`` (see ``_build_evaluation_context``).
|
|
prompt = payload.get("prompt", "")
|
|
return {
|
|
"event": {
|
|
"type": "PHASE_REQUEST",
|
|
"target": "",
|
|
"data": {
|
|
"text": prompt if isinstance(prompt, str) else json.dumps(prompt),
|
|
},
|
|
"context": {},
|
|
},
|
|
}
|
|
tool_name = payload.get("tool_name", "")
|
|
# Omnigent MCP tools are already policy-checked by the relay path
|
|
# (ProxyMcpManager → Omnigent /mcp endpoint → _evaluate_tool_call_policy).
|
|
# Skip only those here to avoid double evaluation; connector-native MCP
|
|
# tools such as mcp__github__* must still go through this hook.
|
|
if isinstance(tool_name, str) and tool_name.startswith("mcp__omnigent__"):
|
|
return None
|
|
tool_input = payload.get("tool_input") or {}
|
|
if hook_event == _PRE_TOOL_USE:
|
|
return {
|
|
"event": {
|
|
"type": "PHASE_TOOL_CALL",
|
|
"target": "",
|
|
"data": {
|
|
"name": tool_name,
|
|
"arguments": tool_input,
|
|
},
|
|
"context": {},
|
|
},
|
|
}
|
|
if hook_event == _POST_TOOL_USE:
|
|
tool_output = payload.get("tool_output", "")
|
|
return {
|
|
"event": {
|
|
"type": "PHASE_TOOL_RESULT",
|
|
"target": "",
|
|
"data": {
|
|
"result": tool_output,
|
|
},
|
|
"context": {},
|
|
"request_data": {
|
|
"name": tool_name,
|
|
"arguments": tool_input,
|
|
},
|
|
},
|
|
}
|
|
return None
|
|
|
|
|
|
def evaluation_response_to_hook_output(
|
|
hook_event: str,
|
|
eval_response: dict[str, object],
|
|
) -> dict[str, object] | None:
|
|
"""
|
|
Convert an ``EvaluationResponse`` into native-harness hook output JSON.
|
|
|
|
For ``PreToolUse`` the policy layer only *enforces* — it emits a
|
|
``hookSpecificOutput.permissionDecision`` solely for verdicts that
|
|
constrain the tool: ``POLICY_ACTION_DENY`` → ``"deny"`` (with
|
|
``permissionDecisionReason``). ``POLICY_ACTION_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 fails closed with ``"deny"`` rather
|
|
than the old ``"defer"`` — ``defer`` handed control back to the
|
|
harness's ``permission_mode``, which ``acceptEdits`` /
|
|
``bypassPermissions`` would auto-approve, bypassing the human.
|
|
``POLICY_ACTION_ALLOW`` — which is the engine's default verdict when
|
|
no policy matches a tool call, not just an explicit author allow —
|
|
returns ``None`` ("no opinion") so the harness's *own* permission
|
|
system still runs. Emitting ``"allow"`` here would auto-approve the
|
|
tool and suppress the harness's native permission prompt (and, for
|
|
Claude Code, the ``PermissionRequest`` hook that routes that prompt
|
|
to the web UI), collapsing two independent gates — the deployment's
|
|
policy gate and the user's own consent gate — into one. The policy
|
|
layer may block (DENY) or demand approval (ASK); it must not silence
|
|
the user's consent. For ``PostToolUse`` a ``DENY`` is surfaced as
|
|
``additionalContext`` because the tool result is already committed
|
|
— PostToolUse hooks cannot block.
|
|
|
|
For ``UserPromptSubmit`` the output uses the top-level ``decision`` /
|
|
``reason`` contract (not ``permissionDecision``): ``DENY`` → ``{"decision":
|
|
"block", "reason": ...}``, which drops the prompt before the model sees
|
|
it. ASK is resolved server-side (``_hold_native_ask_gate`` collapses it
|
|
to a hard ALLOW/DENY before the response reaches the hook), so the hook
|
|
should never see ASK; if it somehow does, it fails closed by blocking.
|
|
ALLOW (and the engine's no-match default) returns ``None`` so the prompt
|
|
proceeds. Unlike ``PreToolUse``, there is no separate user-consent gate
|
|
on a prompt, so ALLOW need not preserve one.
|
|
|
|
Both Claude Code and Codex consume these exact output shapes, so the
|
|
``hookEventName`` echoed back is the harness-supplied ``hook_event``.
|
|
|
|
:param hook_event: Hook event name, e.g. ``"PreToolUse"``,
|
|
``"PostToolUse"``, or ``"UserPromptSubmit"``.
|
|
:param eval_response: Parsed ``EvaluationResponse`` from AP, e.g.
|
|
``{"result": "POLICY_ACTION_DENY", "reason": "blocked by policy"}``.
|
|
:returns: Hook output dict for the harness to read on stdout, or
|
|
``None`` when there is no verdict to express (allow with no
|
|
rewrite on PostToolUse, or an unknown action).
|
|
"""
|
|
action = eval_response.get("result", "POLICY_ACTION_UNSPECIFIED")
|
|
reason = eval_response.get("reason")
|
|
|
|
if hook_event == _USER_PROMPT_SUBMIT:
|
|
# DENY blocks the prompt; a stray ASK fails closed (also block) since
|
|
# ASK is meant to be resolved server-side before reaching the hook.
|
|
# ALLOW / no-match → None so the prompt proceeds. A non-empty reason
|
|
# is required for the block to take effect (both harnesses drop a
|
|
# block with an empty reason), so default one in.
|
|
if action in ("POLICY_ACTION_DENY", "POLICY_ACTION_ASK"):
|
|
return {
|
|
"decision": "block",
|
|
"reason": reason or "Denied by policy",
|
|
}
|
|
return None
|
|
|
|
if hook_event == _PRE_TOOL_USE:
|
|
# ALLOW (the engine default when no policy matches) is omitted → None,
|
|
# so the harness's own permission prompt still fires; see docstring.
|
|
decision_map = {
|
|
"POLICY_ACTION_DENY": "deny",
|
|
# 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 here. If it
|
|
# somehow does, fail closed with ``deny`` rather than the old
|
|
# ``defer`` — ``defer`` returns control to the harness's
|
|
# permission_mode, which acceptEdits / bypassPermissions would
|
|
# auto-approve, re-opening the very bypass this closes.
|
|
"POLICY_ACTION_ASK": "deny",
|
|
}
|
|
decision = decision_map.get(str(action))
|
|
if decision is None:
|
|
return None
|
|
output: dict[str, object] = {
|
|
"hookEventName": _PRE_TOOL_USE,
|
|
"permissionDecision": decision,
|
|
}
|
|
if decision == "deny" and reason:
|
|
output["permissionDecisionReason"] = reason
|
|
return {"hookSpecificOutput": output}
|
|
|
|
if hook_event == _POST_TOOL_USE:
|
|
if action == "POLICY_ACTION_DENY" and reason:
|
|
return {
|
|
"hookSpecificOutput": {
|
|
"hookEventName": _POST_TOOL_USE,
|
|
"additionalContext": f"[Policy violation] {reason}",
|
|
},
|
|
}
|
|
return None
|
|
|
|
return None
|
|
|
|
|
|
def fail_closed_hook_output(
|
|
hook_event: str, detail: str | None = None
|
|
) -> dict[str, object] | None:
|
|
"""
|
|
Build the fail-closed hook output for an unobtainable policy verdict.
|
|
|
|
Called by the per-harness hooks when the ``/policies/evaluate``
|
|
round-trip cannot produce a usable verdict for an *already-governed*
|
|
session — the server is unreachable, returns a non-2xx status, or
|
|
returns an empty / malformed body. Without this the hooks emitted "no
|
|
opinion" on those paths, silently letting the gated tool run: for
|
|
native harnesses this hook is the sole enforcement point (it gates
|
|
Bash / Write / Edit / the native Skill tool / connector-native
|
|
``mcp__*`` tools), so a transient outage disabled all DENY/ASK
|
|
enforcement.
|
|
|
|
The default is phase-aware, matching
|
|
:data:`omnigent.policies.types.FAIL_CLOSED_PHASES` (the runner-side
|
|
precedent from PR #163) — but expressed in hook-event terms so the
|
|
lightweight hook subprocess need not import the policy package:
|
|
|
|
- ``PreToolUse`` (``PHASE_TOOL_CALL``) fails CLOSED → ``deny``. This is
|
|
the authoritative pre-execution gate; an unevaluable policy must not
|
|
let the call through.
|
|
- ``UserPromptSubmit`` (``PHASE_REQUEST``) fails CLOSED →
|
|
``{"decision": "block", ...}``. This is the sole pre-turn enforcement
|
|
point for native sessions; a server hiccup must not let an over-budget
|
|
or otherwise-blocked request proceed.
|
|
- ``PostToolUse`` (``PHASE_TOOL_RESULT``) fails OPEN → ``None``. By the
|
|
result phase the tool has already executed, so denying would only block
|
|
an already-incurred side effect.
|
|
|
|
:param hook_event: Hook event name, e.g. ``"PreToolUse"``.
|
|
:param detail: Optional short diagnostic string appended to the reason
|
|
shown in the UI, e.g. a reauth failure message from
|
|
:attr:`PolicyHookReauth.failure_reason`. Omit when no detail is
|
|
available.
|
|
:returns: A ``permissionDecision: "deny"`` hook output for
|
|
``PreToolUse``; a ``decision: "block"`` output for
|
|
``UserPromptSubmit``; ``None`` for every other event (fail open).
|
|
"""
|
|
tool_reason = (
|
|
f"{_EVAL_UNAVAILABLE_REASON} Detail: {detail}" if detail else _EVAL_UNAVAILABLE_REASON
|
|
)
|
|
request_reason = (
|
|
f"{_EVAL_UNAVAILABLE_REQUEST_REASON} Detail: {detail}"
|
|
if detail
|
|
else _EVAL_UNAVAILABLE_REQUEST_REASON
|
|
)
|
|
if hook_event == _PRE_TOOL_USE:
|
|
return {
|
|
"hookSpecificOutput": {
|
|
"hookEventName": _PRE_TOOL_USE,
|
|
"permissionDecision": "deny",
|
|
"permissionDecisionReason": tool_reason,
|
|
},
|
|
}
|
|
if hook_event == _USER_PROMPT_SUBMIT:
|
|
return {
|
|
"decision": "block",
|
|
"reason": request_reason,
|
|
}
|
|
return None
|
|
|
|
|
|
def post_evaluate_with_retry(
|
|
url: str,
|
|
headers: dict[str, str],
|
|
eval_request: dict[str, object],
|
|
read_timeout: float,
|
|
hook_label: str,
|
|
reauth: Callable[[], dict[str, str] | None] | None = None,
|
|
) -> tuple[httpx.Response, None] | tuple[None, str]:
|
|
"""
|
|
POST to the Omnigent policy evaluate endpoint, retrying on transient errors.
|
|
|
|
Retries on 5xx HTTP responses and connection-level errors
|
|
(:class:`httpx.ConnectError`, :class:`httpx.ConnectTimeout`) within
|
|
:data:`_EVALUATE_POLICY_RETRY_BUDGET_S`. Returns the successful response,
|
|
or ``None`` if the budget is exhausted or a non-retryable error occurs.
|
|
|
|
A stable ``_omnigent_elicitation_id`` is minted once and stamped on
|
|
every attempt. When the server parks an ASK gate and the connection
|
|
drops (5xx or :class:`httpx.ConnectError`), the retry re-POSTs the
|
|
same id so the server re-attaches to the existing elicitation rather
|
|
than minting a new one — mirroring the ``_post_hook_with_reattach``
|
|
idiom used by the ``PermissionRequest`` hook. This prevents a
|
|
second approval card from appearing when the first was already
|
|
published before the error.
|
|
|
|
4xx responses are final — a bad request won't succeed on retry. Other
|
|
mid-stream errors (e.g. :class:`httpx.ReadTimeout`) are also not retried:
|
|
a read timeout fires *after* the server received the request and may
|
|
mean the long-polling ASK gate was severed mid-wait; retrying with the
|
|
same id will re-park the existing elicitation (no duplicate card), but
|
|
the caller's fail-closed path is equivalent and simpler. The caller is
|
|
responsible for fail-closed handling on ``None``.
|
|
|
|
:param url: Absolute URL of the evaluate endpoint.
|
|
:param headers: Auth headers for the Omnigent server.
|
|
:param eval_request: ``EvaluationRequest`` JSON body to POST.
|
|
:param read_timeout: Per-attempt read timeout in seconds. Should be
|
|
large (e.g. one day) to accommodate long-polling ASK gates.
|
|
:param hook_label: Diagnostic label used in stderr messages,
|
|
e.g. ``"evaluate-policy hook"`` or ``"codex evaluate-policy hook"``.
|
|
:param reauth: Optional callable that re-mints fresh auth headers when
|
|
the server bounces the request to its OAuth login flow (the Apps
|
|
front door 302→``/oidc/``) or returns ``401`` — i.e. the one-shot
|
|
``ap_auth_headers`` token lapsed. Called at most once; returning new
|
|
headers triggers an immediate retry with them, mirroring the runner's
|
|
refresh-capable :class:`~omnigent.runner._entry._RunnerDatabricksAuth`.
|
|
``None`` (the default) keeps the legacy behavior for callers that have
|
|
no token source. Returning ``None`` from it falls through to the
|
|
normal failure handling (the caller fails closed).
|
|
:returns: ``(response, error)`` — on success, ``(response, None)``; on
|
|
failure, ``(None, short_error_string)`` describing the last error so
|
|
callers can surface it in the deny/block reason shown to the user.
|
|
"""
|
|
# Mint one stable id for the whole retry sequence. Each retry re-sends
|
|
# it so the server can re-park the SAME elicitation rather than opening
|
|
# a second approval card. The ``elicit_evaluate_`` namespace is validated
|
|
# server-side by ``_EVALUATE_HOOK_ELICITATION_ID_RE``.
|
|
elicitation_id = f"elicit_evaluate_{secrets.token_hex(16)}"
|
|
request_body = {**eval_request, "_omnigent_elicitation_id": elicitation_id}
|
|
deadline = time.monotonic() + _EVALUATE_POLICY_RETRY_BUDGET_S
|
|
backoff_s = _EVALUATE_POLICY_RETRY_INITIAL_BACKOFF_S
|
|
timeout = httpx.Timeout(read_timeout, connect=_EVALUATE_POLICY_CONNECT_TIMEOUT_S)
|
|
reauthed = False
|
|
last_error: str = "unknown error"
|
|
while True:
|
|
try:
|
|
with httpx.Client(headers=headers, timeout=timeout) as client:
|
|
resp = client.post(url, json=request_body)
|
|
if (
|
|
reauth is not None
|
|
and not reauthed
|
|
and _is_login_redirect_or_unauthorized(resp)
|
|
):
|
|
# The one-shot ``ap_auth_headers`` token lapsed (~1h
|
|
# Databricks OAuth lifetime): the Apps front door bounces
|
|
# an expired bearer with a 302→/oidc/ (or a 401). Re-mint
|
|
# and retry once with the fresh token instead of failing
|
|
# closed — exactly as ``_RunnerDatabricksAuth`` does for the
|
|
# runner's own callbacks. Without this, every tool call on a
|
|
# session older than the token lifetime fails CLOSED while
|
|
# chat (refresh-capable) keeps working.
|
|
refreshed = reauth()
|
|
if refreshed:
|
|
headers = refreshed
|
|
reauthed = True
|
|
print(
|
|
f"omnigent {hook_label}: Omnigent auth expired "
|
|
"(login redirect/401); re-minted token and retrying",
|
|
file=sys.stderr,
|
|
)
|
|
continue
|
|
resp.raise_for_status()
|
|
return resp, None
|
|
except httpx.HTTPStatusError as exc:
|
|
status = exc.response.status_code
|
|
body_preview = exc.response.text[:200] if exc.response.content else ""
|
|
last_error = f"server returned {status}" + (
|
|
f": {body_preview}" if body_preview else ""
|
|
)
|
|
if status < 500:
|
|
print(
|
|
f"omnigent {hook_label}: Omnigent returned {status}"
|
|
+ (f": {body_preview}" if body_preview else ""),
|
|
file=sys.stderr,
|
|
)
|
|
return None, last_error
|
|
print(
|
|
f"omnigent {hook_label}: Omnigent returned {status}; retrying",
|
|
file=sys.stderr,
|
|
)
|
|
except (httpx.ConnectError, httpx.ConnectTimeout) as exc:
|
|
last_error = f"connection error: {exc}"
|
|
print(
|
|
f"omnigent {hook_label}: Omnigent request failed; retrying: {exc}",
|
|
file=sys.stderr,
|
|
)
|
|
except httpx.HTTPError as exc:
|
|
# Other HTTP errors (ReadTimeout while a long ASK poll is in flight,
|
|
# etc.) are not retried — retrying a severed ASK would open a new
|
|
# elicitation and prompt the human twice.
|
|
last_error = f"request error: {exc}"
|
|
print(
|
|
f"omnigent {hook_label}: Omnigent request failed: {exc}",
|
|
file=sys.stderr,
|
|
)
|
|
return None, last_error
|
|
if time.monotonic() + backoff_s >= deadline:
|
|
print(
|
|
f"omnigent {hook_label}: retry budget exhausted",
|
|
file=sys.stderr,
|
|
)
|
|
return None, f"retry budget exhausted (last error: {last_error})"
|
|
# Two-step backoff; not worth a retry library in this dependency-light hook.
|
|
time.sleep(backoff_s)
|
|
backoff_s = min(backoff_s * 2, _EVALUATE_POLICY_RETRY_MAX_BACKOFF_S)
|