Files
Corey Zumar bc3dc80200 perf(claude-native): Improve performance of claude native terminal typing, text streaming, etc. (#4582)
* perf(claude-native): stop taxing every hook spawn with the eager package init

Claude Code blocks its TUI on command hooks — once per streamed text
chunk (MessageDisplay), per statusline refresh, and per tool call — and
every 'python -m omnigent.<hook>' subprocess re-ran omnigent/__init__,
which eagerly imported the datamodel/executor/model-catalog graph. The
deliberately stdlib-only hot-path hooks paid ~250 ms per spawn for
imports they never use, capping visible streaming at ~4 chunks/s.

The package init now re-exports lazily (PEP 562): the FIPS md5 patch
and legacy-env mirror stay eager, every public name resolves on first
attribute access (optional executors keep their import-failure->None
contract), and submodule attribute access still works. Hot-path hook
spawns drop to ~30 ms (~interpreter cost).

A native_hook_spawn benchmark journey spawns the MessageDisplay hook
exactly as Claude Code does and rides the release/nightly regression
comparison; fresh-interpreter import-graph guards in the display-hook
test suite pin what each hook entrypoint may import so the regression
cannot silently return.

Signed-off-by: dbczumar <corey.zumar@databricks.com>

* perf(claude-native): keep the hook's hot path off the bridge's heavy imports

The observer hook — Claude blocks on it at every prompt submit, tool
call, Stop, and task event — imported claude_native_bridge, whose
module-level tools/spec/pydantic imports cost ~450 ms of interpreter
startup, plus httpx and the policy machinery besides. Enter and every
tool call paid roughly a second of subprocess overhead per event even
after the package init went lazy.

The bridge now defers its tools graph to the one launch-path function
that builds MCP tools (_build_tools) and its bundle-skills parse to
the launch args builder; the hook imports httpx and the policy
machinery inside the subcommands that actually speak HTTP. Module
import cost: bridge 450 -> ~70 ms, hook 360 -> ~70 ms, and the hook's
fresh-interpreter import graph now contains no third-party modules at
all — the import guard pins the allowance at exactly that.

Tests that reached httpx or create_os_environment through the hook's
or bridge's module attributes now patch the owning modules directly.

Signed-off-by: dbczumar <corey.zumar@databricks.com>

* perf(claude-native): cache the ungoverned policy verdict at the relay

Sessions with no policies at all still paid a full server round trip
(~0.5-1.3s measured against a Databricks App) on every policy hook
event — twice per tool call plus every prompt submit — with the server
answering the same fast-path ALLOW each time. Typing during agentic
turns stuttered in the gaps; vanilla Claude pays nothing there.

The evaluate endpoint now stamps 'governed': false on its existing
no-policies fast path (any_policies_apply's False is session-scoped —
its only phase-scoped rule forces True), and the native-harness
loopback relay caches that verdict for 30s, answering hook events
instantly. A governed response of any kind drops the cache, a
sys_add_policy call through the relay's own /tool path clears it
before the policy lands, and expiry re-validates upstream — so
enforcement for governed sessions is untouched and the attach delay
for out-of-band policy edits is bounded at the TTL.

Signed-off-by: dbczumar <corey.zumar@databricks.com>

* perf(claude-native): keep blocking Claude hooks off Python and off the WAN

Claude blocks its TUI on every command hook, and three of them still
spawned a Python interpreter per event (~30ms floor, ~77ms under EDR):
MessageDisplay once per streamed chunk, statusLine per refresh, and
evaluate-policy twice per tool call — the last one also paying a
0.5-1.3s WAN round trip whenever its 30s ungoverned-cache window
lapsed.

- MessageDisplay: a /bin/sh one-liner appends the payload (newline-
  stripped, so any valid JSON lands single-line) straight to
  message_deltas.jsonl; the deltas reader already parses by key and
  skips malformed lines.
- statusLine: the shim captures raw stdin to context_raw.json (atomic
  rename) and chains the user's own status command; the forwarder
  normalizes it into context.json on its poll loop
  (sync_raw_status_context), so the Python normalizer leaves the
  blocking path. The module entrypoint stays for older bridge dirs.
- evaluate-policy: hooks try a curl against the relay's new
  /hook/claude/evaluate-policy endpoint (advertised via a
  shell-sourceable tool_relay.env); the long-lived runner process owns
  payload→EvaluationRequest mapping, retries, the ungoverned cache,
  and verdict→hook-output shaping. When the relay is absent or
  unreachable the same stdin replays into the Python hook, which keeps
  the direct-server path and the phase-aware fail-closed contract —
  exactly the pre-curl behavior.
- The relay starts at session create (runner app) instead of at the
  first web-dispatched turn, so prompts typed directly in the TUI get
  the curl fast path too; it comes up in the background, and hooks
  that beat it use the Python fallback.

Typing during a live 25-tool-call turn against a Databricks App
measured 56.0ms median / 57.2ms p90 / 0 samples over 200ms, from
118ms median / 264ms p90 / 8 freezes before this branch.

Also pins the relay-close ownership test's trusted-parent monkeypatch
to tempfile.gettempdir() — the literal /tmp never contains the macOS
fixture root, so the test only passed on Linux.

Signed-off-by: dbczumar <corey.zumar@databricks.com>

* perf(onboarding): cache harness CLI version and login probes

Every readiness refresh on every host daemon execs vendor CLIs
(--version / auth status) whose answers change only when the binary is
swapped or a login flips; with a few dozen idle hosts that compounds
into a constant machine-wide subprocess storm (~116 spawns/min
observed) that competes with interactive terminals.

--version output is a pure function of the binary bytes, so successful
parses cache permanently against the binary's (path, mtime_ns, size)
signature; failures keep re-probing. Login verdicts can flip without a
binary change, so only positives cache, with a 120s TTL — negatives
always re-probe so the setup wizard sees a fresh login immediately, and
harness_logout invalidates its key so a successful logout is confirmed
live.

Signed-off-by: dbczumar <corey.zumar@databricks.com>

* revert(claude-native): drop the ungoverned-verdict relay cache

The cache required stamping 'governed': false on the evaluate
response so the relay could tell which ALLOWs were safe to reuse —
new response-field surface carried only by this optimization, which
we don't need right now. Remove the stamp and the relay cache
wholesale: every policy hook event consults the server again, the
relay's /policies/evaluate proxy is a plain pass-through, and the
evaluate response is byte-identical to its pre-branch shape. The
sh-shim/curl hook path (no interpreter spawns) is unchanged.

Signed-off-by: dbczumar <corey.zumar@databricks.com>

---------

Signed-off-by: dbczumar <corey.zumar@databricks.com>
2026-08-13 13:11:46 -07:00

1393 lines
59 KiB
Python

"""Claude Code hook recorder for the native Omnigent wrapper."""
from __future__ import annotations
import argparse
import json
import math
import os
import secrets
import sys
import time
from collections.abc import Callable
from pathlib import Path
from typing import TYPE_CHECKING
from omnigent.claude_native_bridge import (
BRIDGE_ID_LABEL_KEY,
read_active_session_id,
read_bridge_id,
read_claude_session_id,
read_claude_status_model,
read_permission_hook_config,
read_seen_claude_session_ids,
record_hook_event,
transcript_has_forked_from_marker,
transcript_has_recent_local_command,
url_component,
write_active_session_id,
)
# The observer path (the default, most frequent invocation — Claude blocks
# on it per Stop/UserPromptSubmit/TaskCreated/...) must not pay the
# httpx/policy import cost; those are imported inside the subcommands and
# helpers that actually speak HTTP.
if TYPE_CHECKING:
import httpx
# Client-side budget for the permission-request long-poll to AP. Held
# at one day so the hook subprocess waits ~indefinitely for a verdict
# from the web UI (or for Claude to close the connection when the user
# answers in the terminal). Kept in lockstep with the server-side
# ``_CLAUDE_NATIVE_PERMISSION_HOOK_TIMEOUT_S`` and Claude Code's own
# command-hook ``timeout`` so no single layer caps the wait early.
#
# NOTE: this bounds a SINGLE long-poll (a slow human), NOT the retry loop.
# Re-POST attempts after a *failure* are bounded separately by
# ``_PERMISSION_MAX_CONSECUTIVE_FAILURES`` — see :func:`_post_hook_with_reattach`.
# Before #1782 the retry deadline was also one day, so a persistently
# sick/unreachable server made the hook re-POST (each re-driving the turn and
# respawning harness/tool subprocesses) every ≤30s for 24h — the spin-loop half
# of the zombie pileup.
_PERMISSION_TIMEOUT_S = 86400.0
# First retry must land inside the server's re-park grace (proxies
# sever idle long-polls); later retries back off.
_PERMISSION_RETRY_INITIAL_BACKOFF_S = 1.0
_PERMISSION_RETRY_MAX_BACKOFF_S = 30.0
# Fail unreachable-server connects fast into the backoff loop instead
# of inheriting the day-long read budget.
_PERMISSION_CONNECT_TIMEOUT_S = 30.0
def _env_int(name: str, default: int) -> int:
"""Read an int env override, ignoring a malformed value.
A non-integer value must not crash the hook subprocess at import time —
that would defeat the terminal-fallback design elsewhere in this file. Bad
input logs a diagnostic and falls back to *default*.
:param name: Environment variable name.
:param default: Value used when unset or unparseable.
:returns: The parsed int, or *default*.
"""
raw = os.environ.get(name)
if raw is None:
return default
try:
return int(raw)
except ValueError:
print(
f"omnigent hook: ignoring non-integer {name}={raw!r}; using {default}",
file=sys.stderr,
)
return default
def _env_float(name: str, default: float) -> float:
"""Read a float env override, ignoring a malformed value.
Float sibling of :func:`_env_int`; a bad value logs and falls back to
*default* rather than crashing the hook at import time.
:param name: Environment variable name.
:param default: Value used when unset or unparseable.
:returns: The parsed float, or *default*.
"""
raw = os.environ.get(name)
if raw is None:
return default
try:
value = float(raw)
except ValueError:
value = math.nan
# Reject non-finite too: ``float("inf"/"nan")`` parses without ValueError,
# but an ``inf`` floor would classify every sever as a held poll (disabling
# flap detection) and ``nan`` makes every ``held_s < floor`` comparison
# False — both silent footguns, so treat them as malformed.
if not math.isfinite(value):
print(
f"omnigent hook: ignoring non-finite {name}={raw!r}; using {default}",
file=sys.stderr,
)
return default
return value
# Cap on CONSECUTIVE HARD failures before the reattach loop gives up and lets
# the caller fail-ask. A "hard failure" is the server being sick or unreachable
# (5xx, or a connection that never established) — i.e. the #1782 spin. It is
# distinguished from a proxy severing a *held* poll (a connection that WAS
# established and then dropped mid-wait), which is the re-park mechanism working
# as intended and does NOT count — see :func:`_post_hook_with_reattach`. So a
# legitimately-parked approval behind a severing proxy is never capped, while a
# down server can no longer re-POST for a day. Overridable for operators who
# want more slack against a flaky upstream.
_PERMISSION_MAX_CONSECUTIVE_FAILURES = max(1, _env_int("OMNIGENT_HOOK_MAX_RETRIES", 8))
def _never_connected_errors() -> tuple[type[Exception], ...]:
"""httpx errors meaning the request never reached a live server.
No response was ever begun — unambiguous hard failures (the server is
down / unreachable), not a held poll. Everything else under
``httpx.HTTPError`` that is not a 4xx/5xx status (RemoteProtocolError,
ReadError, ReadTimeout, …) means the connection was established and
then severed mid-poll. A function, not a module constant, so the
hook's hot observer path never imports httpx.
"""
import httpx
return (httpx.ConnectError, httpx.ConnectTimeout, httpx.PoolTimeout, httpx.ProxyError)
# An established connection that drops in under this many seconds is treated as
# a flapping/crash-looping server (a hard failure), NOT a genuinely-parked poll
# a proxy severed. Comfortably below any real idle-proxy timeout (typically
# 30-60s+) but above an instant accept-then-reset crash drop, so it tells a
# legitimate long-poll sever from a tight establish-drop spin. Operators behind
# an unusually aggressive proxy/LB whose idle timeout is under 10s can lower
# this via OMNIGENT_HOOK_HELD_POLL_FLOOR_S so their legitimate slow-human severs
# stay classified as held polls (not flaps) and are never capped; the default
# suits typical proxies. Floored at 0 (a negative would make every sever a
# held poll, disabling flap detection).
_PERMISSION_HELD_POLL_FLOOR_S = max(0.0, _env_float("OMNIGENT_HOOK_HELD_POLL_FLOOR_S", 10.0))
# Fail-fast budget for the synchronous ``/clear`` and ``/fork`` session
# rotations that run inside the SessionStart hook to gate Claude's
# welcome banner. Unlike the permission long-poll these are quick
# request/reply calls, so they must NOT inherit the day-long permission
# budget — an unresponsive Omnigent server would otherwise hang the banner.
# On timeout the rotation returns ``None`` and the background forwarder
# performs it from the recorded hook event instead.
_SESSION_ROTATION_TIMEOUT_S = 70.0
# Evaluate-policy hooks normally return immediately, but a TOOL_CALL
# ASK now parks server-side (URL-based elicitation) until a human
# resolves it via the approve URL — so the client must wait as long as
# the permission long-poll. Held at one day; the server caps the real
# wait via the deciding policy's ``ask_timeout``.
_EVALUATE_POLICY_TIMEOUT_S = 86400.0
_FORK_COMMAND_NAMES = frozenset({"/branch", "/fork"})
_FORK_TRANSCRIPT_WAIT_S = 1.0
_FORK_TRANSCRIPT_POLL_S = 0.05
def main(argv: list[str] | None = None) -> int:
"""
Record one Claude Code hook payload from stdin.
:param argv: Optional argv override excluding program name.
``None`` reads :data:`sys.argv`.
:returns: Process exit code. Returns ``0`` for malformed input
after writing a diagnostic to stderr so Claude Code itself
is not blocked by an observer failure.
"""
raw_argv = sys.argv[1:] if argv is None else argv
if raw_argv and raw_argv[0] == "permission-request":
return _main_permission_request(raw_argv[1:])
if raw_argv and raw_argv[0] == "ask-user-question":
return _main_ask_user_question(raw_argv[1:])
if raw_argv and raw_argv[0] == "evaluate-policy":
return _main_evaluate_policy(raw_argv[1:])
if raw_argv and raw_argv[0] == "route-turn":
return _main_route_turn(raw_argv[1:])
# Backwards compat: older bridge dirs may still reference the
# pre-tool-use subcommand before the terminal is restarted.
if raw_argv and raw_argv[0] == "pre-tool-use":
return _main_evaluate_policy(raw_argv[1:])
args = _parse_args(raw_argv)
raw = sys.stdin.read()
try:
payload = json.loads(raw or "{}")
except json.JSONDecodeError as exc:
print(f"omnigent claude hook: malformed JSON: {exc}", file=sys.stderr)
return 0
if not isinstance(payload, dict):
print("omnigent claude hook: expected JSON object", file=sys.stderr)
return 0
bridge_dir = Path(args.bridge_dir)
_annotate_resume_session_context(bridge_dir, payload)
if payload.get("hook_event_name") == "SessionStart" and payload.get("source") == "clear":
rotated_session_id = _rotate_session_on_clear(bridge_dir)
if rotated_session_id:
payload["omnigent_clear_rotated_to"] = rotated_session_id
elif _is_claude_branch_session_start(payload):
payload["omnigent_fork_detected"] = True
rotated_session_id = _rotate_session_on_fork(bridge_dir)
if rotated_session_id:
payload["omnigent_fork_rotated_to"] = rotated_session_id
try:
record_hook_event(bridge_dir, payload)
except Exception as exc: # noqa: BLE001 - hook must not break Claude Code.
print(f"omnigent claude hook: failed to record hook: {exc}", file=sys.stderr)
conversation_url = _conversation_url_for_active_session(bridge_dir, args.conversation_url)
if conversation_url and payload.get("hook_event_name") == "SessionStart":
print(
json.dumps({"systemMessage": (f"Open this session in Omnigent: {conversation_url}")})
)
return 0
def _annotate_resume_session_context(bridge_dir: Path, payload: dict[str, object]) -> None:
"""
Attach pre-recording Claude resume context to a hook payload.
The hook updates bridge state only after this payload is recorded.
Capturing the prior session id and seen-status here lets the
background forwarder distinguish a new Claude branch from a later
ordinary resume into an already-seen branch.
:param bridge_dir: Native Claude bridge directory.
:param payload: Hook payload read from Claude Code stdin, e.g.
``{"hook_event_name": "SessionStart", "source": "resume"}``.
:returns: None.
"""
if payload.get("hook_event_name") != "SessionStart" or payload.get("source") != "resume":
return
new_claude_session_id = payload.get("session_id")
if not isinstance(new_claude_session_id, str) or not new_claude_session_id:
return
seen = read_seen_claude_session_ids(bridge_dir)
current_claude_session_id = read_claude_session_id(bridge_dir)
if (
isinstance(current_claude_session_id, str)
and current_claude_session_id
and current_claude_session_id != new_claude_session_id
):
payload["omnigent_previous_claude_session_id"] = current_claude_session_id
payload["omnigent_claude_session_was_seen"] = new_claude_session_id in seen
def _is_claude_branch_session_start(payload: dict[str, object]) -> bool:
"""
Return whether a ``SessionStart`` payload represents Claude ``/fork``.
Claude Code reports both ordinary resumes and branch/fork switches
as ``SessionStart`` with ``source="resume"``. The branch path is
identified by a transcript fork marker plus a new Claude session
id that differs from the current bridge session and has not been
seen by this wrapper.
:param payload: Hook payload read from Claude Code stdin, e.g.
``{"hook_event_name": "SessionStart", "source": "resume"}``.
:returns: ``True`` when the event should fork the active Omnigent session.
"""
if payload.get("hook_event_name") != "SessionStart" or payload.get("source") != "resume":
return False
new_claude_session_id = payload.get("session_id")
if not isinstance(new_claude_session_id, str) or not new_claude_session_id:
return False
current_claude_session_id = payload.get("omnigent_previous_claude_session_id")
if not isinstance(current_claude_session_id, str) or not current_claude_session_id:
return False
if payload.get("omnigent_claude_session_was_seen") is True:
return False
return _payload_transcript_has_recent_branch_command(
payload,
time.time(),
source_claude_session_id=current_claude_session_id,
wait_timeout_s=_FORK_TRANSCRIPT_WAIT_S,
)
def _payload_transcript_has_recent_branch_command(
payload: dict[str, object],
recorded_at: float,
*,
source_claude_session_id: str | None = None,
wait_timeout_s: float = 0.0,
) -> bool:
"""
Return whether Claude's transcript reports a fork signal.
This avoids using the conversation display title to infer
``/fork``. The primary signal is Claude's structured
``forkedFrom`` transcript metadata; local command records are kept
as an extra signal for Claude versions that emit them.
:param payload: Hook payload read from Claude Code stdin, e.g.
``{"transcript_path": "/tmp/claude.jsonl", "session_id": "..."}``.
:param recorded_at: Unix timestamp for the current hook event,
e.g. ``1779922393.222``.
:param source_claude_session_id: Expected source Claude session
uuid, e.g. ``"9abc..."``. ``None`` accepts any different
non-empty source id.
:param wait_timeout_s: Maximum seconds to wait for Claude to
flush the fork marker, e.g. ``1.0``.
:returns: ``True`` when the transcript contains a recent
``/branch`` or ``/fork`` signal for the hook session.
"""
transcript_path = payload.get("transcript_path")
claude_session_id = payload.get("session_id")
if not isinstance(transcript_path, str) or not transcript_path:
return False
if not isinstance(claude_session_id, str) or not claude_session_id:
return False
path = Path(transcript_path)
deadline = time.monotonic() + max(wait_timeout_s, 0.0)
while True:
if transcript_has_forked_from_marker(
path,
claude_session_id=claude_session_id,
source_claude_session_id=source_claude_session_id,
) or transcript_has_recent_local_command(
path,
claude_session_id=claude_session_id,
recorded_at=recorded_at,
command_names=_FORK_COMMAND_NAMES,
):
return True
if time.monotonic() >= deadline:
return False
time.sleep(_FORK_TRANSCRIPT_POLL_S)
def _rotate_session_on_clear(bridge_dir: Path) -> str | None:
"""
Rotate Omnigent sessions synchronously for a Claude ``/clear`` SessionStart.
The SessionStart hook output is what Claude renders as the welcome
banner. Rotating here lets the banner point at the new Omnigent session
before Claude prints it. Failures return ``None`` so the background
forwarder can still perform the rotation from the recorded hook event.
:param bridge_dir: Native Claude bridge directory.
:returns: New Omnigent session id, e.g. ``"conv_new"``, or ``None`` when
rotation could not be completed from the hook.
"""
old_session_id = read_active_session_id(bridge_dir)
if not old_session_id:
return None
config = read_permission_hook_config(bridge_dir)
ap_server_url = config.get("ap_server_url")
if not isinstance(ap_server_url, str) or not ap_server_url:
return None
raw_headers = config.get("ap_auth_headers")
headers = (
{str(key): str(value) for key, value in raw_headers.items()}
if isinstance(raw_headers, dict)
else {}
)
import httpx
# Route the whole rotation sequence (GET old, POST /v1/sessions or /fork,
# PATCH new, DELETE old) to the replica holding this host's tunnel: a managed
# create/fork notifies the host inline over its pod-local tunnel, so an
# off-replica request can't reach it. This hook client carries no
# _RunnerDatabricksAuth, so key the reused headers dict from the runner-env
# host_id (databricks_request_headers reads OMNIGENT_RUNNER_SLICE_KEY when no
# explicit host_id; emitted only on the workspace mount).
from omnigent.cli_auth import databricks_request_headers
headers.update(databricks_request_headers(ap_server_url))
try:
with httpx.Client(
headers=headers, timeout=httpx.Timeout(_SESSION_ROTATION_TIMEOUT_S)
) as client:
new_session_id = _create_clear_replacement_session(
client,
ap_server_url.rstrip("/"),
old_session_id,
bridge_dir,
)
except httpx.HTTPError as exc:
print(f"omnigent claude clear hook: Omnigent rotation failed: {exc}", file=sys.stderr)
return None
except RuntimeError as exc:
print(f"omnigent claude clear hook: rotation failed: {exc}", file=sys.stderr)
return None
return new_session_id
def _rotate_session_on_fork(bridge_dir: Path) -> str | None:
"""
Fork Omnigent sessions synchronously for a Claude ``/fork``/``/branch``.
Claude renders ``SessionStart`` hook output as the welcome banner
for the new branch. Forking here lets that banner point at the
forked Omnigent session before Claude prints it. Failures return
``None`` so the background forwarder can perform the fork from the
annotated hook record.
:param bridge_dir: Native Claude bridge directory.
:returns: New Omnigent session id, e.g. ``"conv_fork"``, or ``None``
when rotation could not be completed from the hook.
"""
old_session_id = read_active_session_id(bridge_dir)
if not old_session_id:
return None
config = read_permission_hook_config(bridge_dir)
ap_server_url = config.get("ap_server_url")
if not isinstance(ap_server_url, str) or not ap_server_url:
return None
raw_headers = config.get("ap_auth_headers")
headers = (
{str(key): str(value) for key, value in raw_headers.items()}
if isinstance(raw_headers, dict)
else {}
)
import httpx
# Route the whole rotation sequence (GET old, POST /v1/sessions or /fork,
# PATCH new, DELETE old) to the replica holding this host's tunnel: a managed
# create/fork notifies the host inline over its pod-local tunnel, so an
# off-replica request can't reach it. This hook client carries no
# _RunnerDatabricksAuth, so key the reused headers dict from the runner-env
# host_id (databricks_request_headers reads OMNIGENT_RUNNER_SLICE_KEY when no
# explicit host_id; emitted only on the workspace mount).
from omnigent.cli_auth import databricks_request_headers
headers.update(databricks_request_headers(ap_server_url))
try:
with httpx.Client(
headers=headers, timeout=httpx.Timeout(_SESSION_ROTATION_TIMEOUT_S)
) as client:
new_session_id = _create_fork_replacement_session(
client,
ap_server_url.rstrip("/"),
old_session_id,
bridge_dir,
)
except httpx.HTTPError as exc:
print(f"omnigent claude fork hook: Omnigent fork failed: {exc}", file=sys.stderr)
return None
except RuntimeError as exc:
print(f"omnigent claude fork hook: fork failed: {exc}", file=sys.stderr)
return None
return new_session_id
def _create_clear_replacement_session(
client: httpx.Client,
ap_server_url: str,
old_session_id: str,
bridge_dir: Path,
) -> str:
"""
Create and activate the fresh Omnigent session for ``/clear``.
:param client: Sync Omnigent HTTP client.
:param ap_server_url: Omnigent server base URL without a trailing slash,
e.g. ``"http://127.0.0.1:8787"``.
:param old_session_id: Session being rotated away from, e.g.
``"conv_old"``.
:param bridge_dir: Native Claude bridge directory.
:returns: New Omnigent session id, e.g. ``"conv_new"``.
:raises httpx.HTTPError: If Omnigent rejects session creation,
new-session binding, or terminal transfer.
:raises RuntimeError: If Omnigent returns malformed session data.
"""
old_resp = client.get(f"{ap_server_url}/v1/sessions/{url_component(old_session_id)}")
old_resp.raise_for_status()
old = old_resp.json()
if not isinstance(old, dict):
raise RuntimeError(f"session {old_session_id!r} snapshot was not an object")
agent_id = old.get("agent_id")
if not isinstance(agent_id, str) or not agent_id:
raise RuntimeError(f"session {old_session_id!r} has no agent_id")
runner_id = old.get("runner_id")
raw_labels = old.get("labels")
labels = (
{str(key): str(value) for key, value in raw_labels.items()}
if isinstance(raw_labels, dict)
else {}
)
labels.setdefault(BRIDGE_ID_LABEL_KEY, read_bridge_id(bridge_dir) or old_session_id)
create_resp = client.post(
f"{ap_server_url}/v1/sessions",
json={
"agent_id": agent_id,
"labels": labels,
},
)
create_resp.raise_for_status()
created = create_resp.json()
if not isinstance(created, dict):
raise RuntimeError("clear replacement session response was not an object")
new_session_id = created.get("id")
if not isinstance(new_session_id, str) or not new_session_id:
raise RuntimeError("clear replacement session response did not include id")
if isinstance(runner_id, str) and runner_id:
bind_resp = client.patch(
f"{ap_server_url}/v1/sessions/{url_component(new_session_id)}",
json={"runner_id": runner_id},
)
bind_resp.raise_for_status()
from omnigent.entities.session_resources import terminal_resource_id
terminal_id = terminal_resource_id("claude", "main")
transfer_resp = client.post(
(
f"{ap_server_url}/v1/sessions/{url_component(old_session_id)}"
f"/resources/terminals/{url_component(terminal_id)}/transfer"
),
json={"target_session_id": new_session_id},
)
transfer_resp.raise_for_status()
write_active_session_id(bridge_dir, new_session_id)
clear_resp = client.patch(
f"{ap_server_url}/v1/sessions/{url_component(old_session_id)}",
json={
"runner_id": "",
# Re-key the superseded session onto a DISTINCT "-cleared" bridge id
# so its later resume gets its own isolated dir instead of the new
# session's live one (which would double-mirror the transcript and
# trip the executor guard). Mirrors the async forwarder rotation;
# ``_auto_create_claude_terminal`` recognises this marker.
"labels": {BRIDGE_ID_LABEL_KEY: f"{old_session_id}-cleared"},
},
)
if clear_resp.status_code >= 400:
print(
(
"omnigent claude clear hook: failed to clear old runner binding: "
f"{clear_resp.status_code} {clear_resp.text}"
),
file=sys.stderr,
)
return new_session_id
def _create_fork_replacement_session(
client: httpx.Client,
ap_server_url: str,
old_session_id: str,
bridge_dir: Path,
) -> str:
"""
Create and activate the forked Omnigent session for Claude ``/fork``.
:param client: Sync Omnigent HTTP client.
:param ap_server_url: Omnigent server base URL without a trailing slash,
e.g. ``"http://127.0.0.1:8787"``.
:param old_session_id: Session being forked away from, e.g.
``"conv_old"``.
:param bridge_dir: Native Claude bridge directory.
:returns: New Omnigent session id, e.g. ``"conv_fork"``.
:raises httpx.HTTPError: If Omnigent rejects session fetch, fork,
new-session binding, or terminal transfer.
:raises RuntimeError: If Omnigent returns malformed session data.
"""
old_resp = client.get(f"{ap_server_url}/v1/sessions/{url_component(old_session_id)}")
old_resp.raise_for_status()
old = old_resp.json()
if not isinstance(old, dict):
raise RuntimeError(f"session {old_session_id!r} snapshot was not an object")
runner_id = old.get("runner_id")
fork_resp = client.post(
f"{ap_server_url}/v1/sessions/{url_component(old_session_id)}/fork",
json={},
)
fork_resp.raise_for_status()
forked = fork_resp.json()
if not isinstance(forked, dict):
raise RuntimeError("fork replacement session response was not an object")
new_session_id = forked.get("id")
if not isinstance(new_session_id, str) or not new_session_id:
raise RuntimeError("fork replacement session response did not include id")
if isinstance(runner_id, str) and runner_id:
bind_resp = client.patch(
f"{ap_server_url}/v1/sessions/{url_component(new_session_id)}",
json={"runner_id": runner_id},
)
bind_resp.raise_for_status()
from omnigent.entities.session_resources import terminal_resource_id
terminal_id = terminal_resource_id("claude", "main")
transfer_resp = client.post(
(
f"{ap_server_url}/v1/sessions/{url_component(old_session_id)}"
f"/resources/terminals/{url_component(terminal_id)}/transfer"
),
json={"target_session_id": new_session_id},
)
transfer_resp.raise_for_status()
write_active_session_id(bridge_dir, new_session_id)
clear_resp = client.patch(
f"{ap_server_url}/v1/sessions/{url_component(old_session_id)}",
json={"runner_id": ""},
)
if clear_resp.status_code >= 400:
print(
(
"omnigent claude fork hook: failed to clear old runner binding: "
f"{clear_resp.status_code} {clear_resp.text}"
),
file=sys.stderr,
)
return new_session_id
def _conversation_url_for_active_session(
bridge_dir: Path,
fallback_url: str | None,
) -> str | None:
"""
Build the web URL for the bridge's current active Omnigent session.
:param bridge_dir: Native Claude bridge directory.
:param fallback_url: Legacy URL supplied by old hook settings, e.g.
``"http://127.0.0.1:8787/c/conv_old"``.
:returns: Web URL for the active session, e.g.
``"http://127.0.0.1:8787/c/conv_new"``, or ``None`` when no
URL can be constructed.
"""
config = read_permission_hook_config(bridge_dir)
ap_server_url = config.get("ap_server_url")
session_id = read_active_session_id(bridge_dir)
if isinstance(ap_server_url, str) and ap_server_url and session_id:
# ``ap_server_url`` is the API base; route through the shared
# builder so workspace-hosted servers land on the ``/omnigent``
# SPA mount (with ``?o=<org>``) rather than the JSON API mount.
from omnigent.conversation_browser import conversation_url
return conversation_url(ap_server_url, session_id)
return fallback_url
def _post_hook_with_reattach(
url: str,
headers: dict[str, str],
payload: dict[str, object],
hook_label: str,
reauth: Callable[[], dict[str, str] | None] | None = None,
) -> httpx.Response | None:
"""
POST one permission-style hook payload, surviving severed long-polls.
Proxies cut the day-long poll and the server can restart mid-wait;
a single failed POST used to fail-ask into a terminal prompt no one
watches on a headless sub-agent. Re-POSTs with one stable
``_omnigent_elicitation_id`` so the server re-parks the SAME
elicitation and can hand back a gap verdict via its pre-resolved
tombstone. Retries transport errors and 5xx; a 4xx is final.
Retries are bounded by ``_PERMISSION_MAX_CONSECUTIVE_FAILURES`` consecutive
*hard* failures, classified by kind — NOT by wall-clock. This is the #1782
fix: before it the retry deadline was ``_PERMISSION_TIMEOUT_S`` (a day), so
a persistently down server re-POSTed — re-driving the turn and respawning
harness/tool subprocesses — every ≤30s for 24h.
Failure classification:
* **Hard failure → count toward the cap.** A 5xx, or a connection that
never established (:func:`_never_connected_errors`), or an established
connection that dropped in under :data:`_PERMISSION_HELD_POLL_FLOOR_S`
(a flapping/crash-looping server). This is the spin.
* **Held-poll sever → reset the counter.** An established connection that
dropped mid-poll after being held ≥ the floor. That is a proxy severing
an idle long-poll — the re-park mechanism working as intended — so it
must NOT count, or a legitimately-parked human approval behind a
severing proxy would be capped after a handful of severs. Classifying by
*how the request failed* (established-then-severed vs never-connected),
not by elapsed wall-clock, is what makes "a slow human is never capped"
hold even when the proxy severs every 30-60s.
A hard 4xx is final (bad request won't succeed on retry).
Residual limitation (by design): a *sick* backend behind a proxy/LB that
accepts the connection and then silently severs it after its upstream
timeout (>= the floor) is indistinguishable at the transport layer from a
proxy severing a genuinely-parked human poll — both surface as an
established-then-severed error after N seconds, and the server holds the
POST silently with no client-visible "parked" ack. So this case resets the
counter and is NOT caught by the consecutive-hard-failure cap; it is bounded
only by the absolute :data:`_PERMISSION_TIMEOUT_S` ceiling below (the same
day-long human-answer window). That is the tightest safe client-side bound:
capping it sooner would necessarily cap a real slow human on the same
topology. The blast radius is limited — this loop only re-POSTs over HTTP
from one hook process (it does not itself respawn harness/tool
subprocesses), and the host-side orphan reaper (#1782 Bug A) reclaims any
subprocesses a re-driven turn does spawn — so the worst case is one hook
slow-retrying for up to a day, not the original zombie pileup.
The absolute :data:`_PERMISSION_TIMEOUT_S` ceiling bounds the total wait on
every path, matching the day-long human-answer window.
:param url: Absolute hook endpoint URL, e.g.
``"http://127.0.0.1:8787/v1/sessions/conv_x/hooks/permission-request"``.
:param headers: Outbound auth headers for the Omnigent server.
:param payload: Hook payload to POST. Not mutated; the re-attach id
rides on a copy.
:param hook_label: Diagnostic prefix for stderr lines, e.g.
``"permission"`` or ``"ask-user-question"``.
:param reauth: Optional callable that re-mints fresh auth headers when the
server bounces the POST to its OAuth login flow (Apps 302→``/oidc/``)
or returns ``401`` — i.e. the one-shot ``ap_auth_headers`` token lapsed.
Called at most once; new headers trigger an immediate retry with them.
``None`` keeps the legacy behavior.
:returns: The successful (2xx) response, or ``None`` when rejected
or out of budget — callers fail-ask as before.
"""
body = {
**payload,
"_omnigent_elicitation_id": f"elicit_claude_{secrets.token_hex(16)}",
}
backoff_s = _PERMISSION_RETRY_INITIAL_BACKOFF_S
import httpx
from omnigent.native_policy_hook import _is_login_redirect_or_unauthorized
timeout = httpx.Timeout(_PERMISSION_TIMEOUT_S, connect=_PERMISSION_CONNECT_TIMEOUT_S)
# Absolute backstop: even a run of held-poll severs (which don't count
# toward the hard-failure cap) can't loop past the day-long human-answer
# window — matches the read budget and the server's ask_timeout.
deadline = time.monotonic() + _PERMISSION_TIMEOUT_S
reauthed = False
consecutive_hard_failures = 0
while True:
attempt_started = time.monotonic()
try:
with httpx.Client(headers=headers, timeout=timeout) as client:
resp = client.post(url, json=body)
if (
reauth is not None
and not reauthed
and _is_login_redirect_or_unauthorized(resp)
):
# One-shot ``ap_auth_headers`` token lapsed (~1h OAuth
# lifetime): re-mint and retry once rather than fail-asking
# into a terminal prompt no one watches. Mirrors the
# evaluate-policy hook and ``_RunnerDatabricksAuth``.
refreshed = reauth()
if refreshed:
headers = refreshed
reauthed = True
print(
f"omnigent {hook_label} hook: Omnigent auth expired "
"(login redirect/401); re-minted token and retrying",
file=sys.stderr,
)
continue
resp.raise_for_status()
return resp
except httpx.HTTPStatusError as exc:
if exc.response.status_code < 500:
print(
f"omnigent {hook_label} hook: Omnigent request rejected: {exc}",
file=sys.stderr,
)
return None
# 5xx: server responded but is sick — a hard failure (the spin).
is_hard_failure = True
print(
f"omnigent {hook_label} hook: Omnigent request failed; retrying: {exc}",
file=sys.stderr,
)
except httpx.HTTPError as exc:
# Classify by HOW it failed, not by elapsed time (a proxy severs a
# legitimately-held poll in seconds-to-minutes, so wall-clock can't
# tell it from a down server — #1782 Polly review).
never_connected = isinstance(exc, _never_connected_errors())
held_s = time.monotonic() - attempt_started
# Hard failure iff the server was never reached, OR an established
# connection dropped so fast it's a flap rather than a parked poll.
is_hard_failure = never_connected or held_s < _PERMISSION_HELD_POLL_FLOOR_S
kind = (
"unreachable"
if never_connected
else ("flapping" if is_hard_failure else "held-poll severed")
)
print(
f"omnigent {hook_label} hook: Omnigent request failed ({kind}); retrying: {exc}",
file=sys.stderr,
)
if is_hard_failure:
consecutive_hard_failures += 1
else:
# A proxy severed a genuinely-held poll — the re-park mechanism
# working as intended. Reset so a slow human is never capped.
consecutive_hard_failures = 0
if consecutive_hard_failures >= _PERMISSION_MAX_CONSECUTIVE_FAILURES:
print(
f"omnigent {hook_label} hook: giving up after "
f"{consecutive_hard_failures} consecutive hard failures "
"(server unreachable/sick) — failing ask",
file=sys.stderr,
)
return None
# Two-line backoff; not worth a retry lib in this dependency-light hook.
time.sleep(backoff_s)
backoff_s = min(backoff_s * 2, _PERMISSION_RETRY_MAX_BACKOFF_S)
if time.monotonic() >= deadline:
print(
f"omnigent {hook_label} hook: retry budget exhausted "
f"({_PERMISSION_TIMEOUT_S:.0f}s) — failing ask",
file=sys.stderr,
)
return None
def _main_permission_request(argv: list[str]) -> int:
"""
Forward one Claude ``PermissionRequest`` hook to the active Omnigent session.
:param argv: CLI argv after the ``permission-request`` subcommand,
e.g. ``["--bridge-dir", "/tmp/x", "--omnigent-server-url",
"http://127.0.0.1:8787"]``.
:returns: Process exit code. Returns ``0`` on transport failures so
Claude Code falls back to its terminal prompt.
"""
from omnigent.native_policy_hook import policy_hook_reauth
args = _parse_permission_args(argv)
raw = sys.stdin.read()
try:
payload = json.loads(raw or "{}")
except json.JSONDecodeError as exc:
print(f"omnigent claude permission hook: malformed JSON: {exc}", file=sys.stderr)
return 0
if not isinstance(payload, dict):
print("omnigent claude permission hook: expected JSON object", file=sys.stderr)
return 0
bridge_dir = Path(args.bridge_dir)
session_id = read_active_session_id(bridge_dir)
if not session_id:
print("omnigent claude permission hook: active session missing", file=sys.stderr)
return 0
config = read_permission_hook_config(bridge_dir)
ap_server_url = args.omnigent_server_url or config.get("ap_server_url")
if not isinstance(ap_server_url, str) or not ap_server_url:
print("omnigent claude permission hook: Omnigent server URL missing", file=sys.stderr)
return 0
headers = _parse_headers(args.omnigent_auth_headers_json)
if not headers:
raw_headers = config.get("ap_auth_headers")
if isinstance(raw_headers, dict):
headers = {str(key): str(value) for key, value in raw_headers.items()}
# A permission request raises a web elicitation that parks in the pod-local
# registry on the replica holding this session's runner tunnel; an unkeyed
# POST lands elsewhere and the approval is silently lost. Key from the
# runner-env host_id (reads OMNIGENT_RUNNER_SLICE_KEY; workspace mount only).
from omnigent.cli_auth import databricks_request_headers
headers.update(databricks_request_headers(ap_server_url))
url = (
f"{ap_server_url.rstrip('/')}/v1/sessions/"
f"{url_component(session_id)}/hooks/permission-request"
)
resp = _post_hook_with_reattach(
url,
headers,
payload,
"claude permission",
reauth=policy_hook_reauth(ap_server_url, headers),
)
if resp is None:
return 0
if resp.content:
sys.stdout.write(resp.text)
return 0
def _main_ask_user_question(argv: list[str]) -> int:
"""
Handle a ``PreToolUse`` hook for Claude's built-in ``AskUserQuestion`` tool.
In ``bypassPermissions`` mode the ``PermissionRequest`` hook never fires,
so ``AskUserQuestion`` questions are silently swallowed — the web UI never
sees them. This handler intercepts the tool via ``PreToolUse`` (which fires
in all permission modes), POSTs the same payload to the Omnigent server's
``/hooks/permission-request`` endpoint, waits for the user's web-UI answer,
and converts the ``PermissionRequest``-format response into a
``PreToolUse``-format output (``updatedInput``) so Claude receives the
answers and skips its own TUI picker.
When ``permission_mode`` is anything other than ``"bypassPermissions"``
(including absent/unknown), this is a no-op: the ``PermissionRequest`` hook
will handle the call, and returning empty output here prevents a duplicate
elicitation from appearing.
:param argv: CLI argv after the ``ask-user-question`` subcommand,
e.g. ``["--bridge-dir", "/tmp/x"]``.
:returns: Process exit code. Returns ``0`` on any failure so Claude Code
falls back to its terminal TUI prompt rather than blocking.
"""
from omnigent.native_policy_hook import policy_hook_reauth
args = _parse_permission_args(argv)
raw = sys.stdin.read()
try:
payload = json.loads(raw or "{}")
except json.JSONDecodeError as exc:
print(f"omnigent ask-user-question hook: malformed JSON: {exc}", file=sys.stderr)
return 0
if not isinstance(payload, dict):
print("omnigent ask-user-question hook: expected JSON object", file=sys.stderr)
return 0
# Only intercept in bypassPermissions mode. In any other mode the
# PermissionRequest hook fires independently and owns the elicitation;
# returning empty output here is "no opinion" so Claude's own flow
# continues unimpeded and we avoid surfacing the form twice.
if payload.get("permission_mode") != "bypassPermissions":
return 0
bridge_dir = Path(args.bridge_dir)
session_id = read_active_session_id(bridge_dir)
if not session_id:
print("omnigent ask-user-question hook: active session missing", file=sys.stderr)
return 0
config = read_permission_hook_config(bridge_dir)
ap_server_url = args.omnigent_server_url or config.get("ap_server_url")
if not isinstance(ap_server_url, str) or not ap_server_url:
print("omnigent ask-user-question hook: Omnigent server URL missing", file=sys.stderr)
return 0
headers = _parse_headers(args.omnigent_auth_headers_json)
if not headers:
raw_headers = config.get("ap_auth_headers")
if isinstance(raw_headers, dict):
headers = {str(key): str(value) for key, value in raw_headers.items()}
# Same as the permission-request hook: the elicitation parks in the pod-local
# registry on the session's tunnel replica, so key the POST from the
# runner-env host_id or an off-replica landing silently drops the prompt.
from omnigent.cli_auth import databricks_request_headers
headers.update(databricks_request_headers(ap_server_url))
url = (
f"{ap_server_url.rstrip('/')}/v1/sessions/"
f"{url_component(session_id)}/hooks/permission-request"
)
resp = _post_hook_with_reattach(
url,
headers,
payload,
"ask-user-question",
reauth=policy_hook_reauth(ap_server_url, headers),
)
if resp is None or not resp.content:
return 0
# The Omnigent server returns a PermissionRequest-shaped response:
# {"hookSpecificOutput": {"hookEventName": "PermissionRequest",
# "decision": {"behavior": "allow",
# "updatedInput": {...}}}}
# Convert to PreToolUse format so Claude applies updatedInput:
# {"hookSpecificOutput": {"hookEventName": "PreToolUse",
# "permissionDecision": "allow",
# "updatedInput": {...}}}
try:
body = resp.json()
except ValueError as exc:
print(f"omnigent ask-user-question hook: invalid JSON from AP: {exc}", file=sys.stderr)
return 0
decision = (
body.get("hookSpecificOutput", {}).get("decision", {}) if isinstance(body, dict) else {}
)
behavior = decision.get("behavior") if isinstance(decision, dict) else None
if behavior not in ("allow", "deny"):
# Unexpected shape — return empty output so Claude falls back to TUI.
return 0
pre_tool_use_output: dict[str, object] = {
"hookSpecificOutput": {
"hookEventName": "PreToolUse",
"permissionDecision": behavior,
},
}
updated_input = decision.get("updatedInput") if isinstance(decision, dict) else None
if isinstance(updated_input, dict):
pre_tool_use_output["hookSpecificOutput"]["updatedInput"] = updated_input # type: ignore[index]
sys.stdout.write(json.dumps(pre_tool_use_output))
return 0
def _main_evaluate_policy(argv: list[str]) -> int:
"""
Evaluate a Claude Code ``PreToolUse`` / ``PostToolUse`` /
``UserPromptSubmit`` hook against Omnigent policies.
Reads the hook JSON payload from stdin, converts it into the
proto-compatible ``EvaluationRequest`` schema (``PHASE_TOOL_CALL``
for PreToolUse, ``PHASE_TOOL_RESULT`` for PostToolUse,
``PHASE_REQUEST`` for UserPromptSubmit), POSTs to
``/v1/sessions/{id}/policies/evaluate``, and converts the
``EvaluationResponse`` back into Claude Code's hook output format.
For ``PreToolUse``, only the constraining ``POLICY_ACTION_DENY``
verdict maps to a ``permissionDecision`` (``"deny"``); ASK is
resolved server-side (the endpoint parks via ``_hold_native_ask_gate``
and returns a hard ALLOW/DENY). ``POLICY_ACTION_ALLOW`` (the engine's
default when no policy matches) emits no output — "no opinion" — so
Claude's own permission prompt still fires and the
``PermissionRequest`` hook can route it to the web UI. See
:func:`omnigent.native_policy_hook.evaluation_response_to_hook_output`.
For ``UserPromptSubmit``, this is the request-phase gate for native
sessions (the server-level ``_evaluate_input_policy`` skips native
message events). A DENY emits top-level ``decision: "block"``, which
drops the prompt before the model sees it; ASK is resolved
server-side; ALLOW proceeds with no output.
For ``PostToolUse``, policy denials are surfaced as
``additionalContext`` (Claude sees the warning but the tool result
is already committed — PostToolUse hooks are observational).
Failure handling is phase-aware (mirroring the runner-side default
from PR #163). Once the session is known to be governed (an active
session id and a configured ``ap_server_url``) and the round-trip to
``/policies/evaluate`` cannot yield a usable verdict — the server is
unreachable, returns non-2xx, or returns an empty / malformed body —
a ``PreToolUse`` (``PHASE_TOOL_CALL``) call fails CLOSED with a
``deny`` (this hook is the sole enforcement point for native tools, so
a transient outage must not silently let a gated call through), while
``UserPromptSubmit`` and ``PostToolUse`` fail OPEN. Pre-evaluation
conditions that mean the session simply is not governed — no active
session, no ``ap_server_url``, an unparseable hook payload, or an
``mcp__omnigent__*`` tool already gated on the relay path — still
return exit 0 with no output ("no opinion") so non-Omnigent tool
calls are never blocked.
:param argv: CLI argv after the ``evaluate-policy`` subcommand,
e.g. ``["--bridge-dir", "/tmp/x"]``.
:returns: Process exit code. Always ``0`` — blocking verdicts
are expressed via the JSON output, not exit codes.
"""
from omnigent.native_policy_hook import (
evaluation_response_to_hook_output,
fail_closed_hook_output,
hook_payload_to_evaluation_request,
policy_hook_reauth,
post_evaluate_with_retry,
read_relay_policy_config,
relay_policy_evaluate_url,
)
args = _parse_evaluate_policy_args(argv)
raw = sys.stdin.read()
try:
payload = json.loads(raw or "{}")
except json.JSONDecodeError as exc:
print(f"omnigent evaluate-policy hook: malformed JSON: {exc}", file=sys.stderr)
return 0
if not isinstance(payload, dict):
print("omnigent evaluate-policy hook: expected JSON object", file=sys.stderr)
return 0
bridge_dir = Path(args.bridge_dir)
session_id = read_active_session_id(bridge_dir)
if not session_id:
return 0
hook_event = payload.get("hook_event_name", "")
eval_request = hook_payload_to_evaluation_request(hook_event, payload)
if eval_request is None:
# Unrecognized hook event — no policy to evaluate.
return 0
# Stamp the live model from this session's statusLine capture.
context = eval_request["event"]["context"]
context["harness"] = "claude-native"
status_model = read_claude_status_model(bridge_dir)
if status_model:
context["model"] = status_model
def _fail_closed(detail: str | None = None) -> int:
out = fail_closed_hook_output(hook_event, detail)
if out is not None:
sys.stdout.write(json.dumps(out))
return 0
# Prefer the relay; fall back to direct server call when relay not yet up.
relay = read_relay_policy_config(bridge_dir)
if relay:
relay_url, relay_token, _sid = relay
url = relay_policy_evaluate_url(relay_url)
headers: dict[str, str] = {
"Content-Type": "application/json",
"Authorization": f"Bearer {relay_token}",
}
reauth = None
else:
config = read_permission_hook_config(bridge_dir)
ap_server_url = config.get("ap_server_url")
if not isinstance(ap_server_url, str) or not ap_server_url:
return 0
headers = {}
raw_headers = config.get("ap_auth_headers")
if isinstance(raw_headers, dict):
headers = {str(k): str(v) for k, v in raw_headers.items()}
# This posts to the session's policy registry on the replica holding its
# tunnel; key from the runner-env host_id so it isn't misrouted. (The
# relay branch above targets a relay token URL, so it needs no key.)
from omnigent.cli_auth import databricks_request_headers
headers.update(databricks_request_headers(ap_server_url))
session_component = url_component(session_id)
url = f"{ap_server_url.rstrip('/')}/v1/sessions/{session_component}/policies/evaluate"
reauth = policy_hook_reauth(ap_server_url, headers)
resp, api_error = post_evaluate_with_retry(
url,
headers,
eval_request,
_EVALUATE_POLICY_TIMEOUT_S,
"evaluate-policy hook",
reauth=reauth,
)
if resp is None:
return _fail_closed(api_error or (reauth.failure_reason if reauth else None))
if not resp.content:
print("omnigent evaluate-policy hook: empty Omnigent response", file=sys.stderr)
return _fail_closed()
try:
eval_response = resp.json()
except json.JSONDecodeError:
print("omnigent evaluate-policy hook: malformed Omnigent response", file=sys.stderr)
return _fail_closed()
hook_output = evaluation_response_to_hook_output(hook_event, eval_response)
if hook_output is not None:
sys.stdout.write(json.dumps(hook_output))
return 0
def _parse_args(argv: list[str]) -> argparse.Namespace:
"""
Parse hook recorder arguments.
:param argv: CLI argv excluding program name, e.g.
``["--bridge-dir", "/tmp/x"]``.
:returns: Parsed argparse namespace.
"""
parser = argparse.ArgumentParser(prog="python -m omnigent.claude_native_hook")
parser.add_argument("--bridge-dir", required=True)
parser.add_argument("--conversation-url")
return parser.parse_args(argv)
def _parse_evaluate_policy_args(argv: list[str]) -> argparse.Namespace:
"""
Parse evaluate-policy hook arguments.
:param argv: CLI argv excluding program name and subcommand.
:returns: Parsed argparse namespace.
"""
parser = argparse.ArgumentParser(prog="python -m omnigent.claude_native_hook evaluate-policy")
parser.add_argument("--bridge-dir", required=True)
return parser.parse_args(argv)
def _parse_permission_args(argv: list[str]) -> argparse.Namespace:
"""
Parse permission hook forwarding arguments.
:param argv: CLI argv excluding program name and subcommand.
:returns: Parsed argparse namespace.
"""
parser = argparse.ArgumentParser(
prog="python -m omnigent.claude_native_hook permission-request"
)
parser.add_argument("--bridge-dir", required=True)
parser.add_argument("--omnigent-server-url")
parser.add_argument("--omnigent-auth-headers-json")
return parser.parse_args(argv)
def _parse_headers(raw: str | None) -> dict[str, str]:
"""
Parse serialized Omnigent auth headers for the permission hook.
:param raw: JSON object string, e.g.
``"{\"Authorization\": \"Bearer token\"}"``. ``None`` means no
headers.
:returns: Header dict with string keys and values.
"""
if not raw:
return {}
try:
parsed = json.loads(raw)
except json.JSONDecodeError as exc:
print(f"omnigent claude permission hook: bad headers JSON: {exc}", file=sys.stderr)
return {}
if not isinstance(parsed, dict):
return {}
return {str(key): str(value) for key, value in parsed.items()}
def _main_route_turn(argv: list[str]) -> int:
"""
Route the model this session runs on, from its first real prompt.
The in-harness half of first-message routing (see
:mod:`omnigent.runner.turn_routing`), registered as an extra
``UserPromptSubmit`` command alongside the forwarder's status hook and
the policy gate. On every prompt submit, in order:
1. Fast skip on ``<bridge_dir>/turn_routing_done`` **when it names this
session** — no output, no network. The authoritative gate is the
endpoint's routing-decision check; this file only saves the round
trip, and a ``/clear`` rotation hands the same bridge dir to a new
conversation whose first message must still be able to route.
2. POST ``{session_id, prompt, harness, model}`` to the advertised
loopback ``route-turn`` endpoint. Claude's hook payload carries no
model, so ``model`` is the live one from ``context.json`` (the
statusLine snapshot) — never a config file, which reports the
launch model.
3. On a routed verdict: write the marker and BLOCK the prompt. The
hook does **not** touch the model itself — the pane is frozen
waiting on this very subprocess, so keystrokes sent from here would
queue behind the block. The runner replays the prompt through the
normal turn path, which applies the routed model under the pane's
inject lock and then delivers the text.
Fails open everywhere: an absent advertisement, an unreachable
endpoint or an unroutable verdict all exit ``0`` with no output, and
the prompt runs untouched on the current model.
:param argv: CLI argv after the ``route-turn`` subcommand, e.g.
``["--bridge-dir", "/tmp/x", "--harness", "claude-native"]``.
:returns: Process exit code. Always ``0`` — the block is expressed via
the JSON on stdout, never via the exit code.
"""
from omnigent.runner.turn_routing import (
ADVERTISEMENT_FILE,
HOOK_REQUEST_TIMEOUT_S,
ROUTE_PATH_TEMPLATE,
turn_routing_marker_present,
)
parser = argparse.ArgumentParser(prog="python -m omnigent.claude_native_hook route-turn")
parser.add_argument("--bridge-dir", required=True)
parser.add_argument("--harness", default="claude-native")
args = parser.parse_args(argv)
bridge_dir = Path(args.bridge_dir)
try:
payload = json.loads(sys.stdin.read() or "{}")
except json.JSONDecodeError:
return 0
if not isinstance(payload, dict):
return 0
prompt = payload.get("prompt")
if not isinstance(prompt, str) or not prompt.strip():
return 0
from omnigent.inner.hook_scripts.subagent_router import read_router_endpoint
endpoint = read_router_endpoint(bridge_dir, filename=ADVERTISEMENT_FILE)
if endpoint is None:
return 0
# The bridge's ACTIVE session wins over the advertisement's, which is
# written once at launch and goes stale the moment ``/clear`` re-keys this
# pane onto a new conversation. Same source the permission hook reads for
# the same reason — approvals and routing both have to follow rotations.
# Reading the stale id instead made the new conversation ask (and skip) as
# the superseded one.
session_id = read_active_session_id(bridge_dir) or endpoint.session_id
if not session_id:
return 0
# The marker is checked here, after the session id is known, because it is
# scoped to a session: a ``/clear`` rotation hands this same bridge dir to
# a NEW conversation, whose first message must still be able to route.
# Still zero network on the fast path.
if turn_routing_marker_present(bridge_dir, session_id):
return 0
body = {
"harness": args.harness,
"prompt": prompt,
# Claude's payload has no turn id the runner could match a replay
# against, and a blocked prompt starts no turn at all.
"turn_id": None,
"model": read_claude_status_model(bridge_dir),
}
url = endpoint.url + ROUTE_PATH_TEMPLATE.format(session_id=url_component(session_id))
decision = _route_turn_post(url, endpoint.token, body, HOOK_REQUEST_TIMEOUT_S)
if decision is None:
return 0
model = decision.get("model")
if decision.get("action") != "route" or not isinstance(model, str) or not model:
if decision.get("terminal"):
# Nothing will route this session again, so stop asking. Covers the
# no-op verdict too (the pick equals the live model): terminal and
# unblocking, so the prompt runs where it already was.
_write_turn_routing_marker(bridge_dir, session_id, decision)
return 0
# The marker is what tells the runner "this prompt was dropped, you owe
# it a replay", so a marker we could not write means we must not block.
if not _write_turn_routing_marker(bridge_dir, session_id, decision):
return 0
sys.stdout.write(
json.dumps(
{
"decision": "block",
"reason": f"Smart Routing selected {model}; rerunning your message on it.",
}
)
)
sys.stdout.flush()
return 0
def _write_turn_routing_marker(
bridge_dir: Path, session_id: str, decision: dict[str, object]
) -> bool:
"""
Write the session-scoped turn-routing marker file.
:param bridge_dir: Native Claude bridge directory.
:param session_id: Session the verdict belongs to — the conversation a
later ``/clear`` rotation creates must not fast-skip on it.
:param decision: The verdict, for its ``decision_id``.
:returns: ``True`` when the marker is on disk.
"""
from omnigent.runner.turn_routing import write_turn_routing_marker
decision_id = decision.get("decision_id")
if write_turn_routing_marker(
bridge_dir,
session_id=session_id,
decision_id=decision_id if isinstance(decision_id, str) else None,
):
return True
print(
"omnigent claude route-turn hook: could not write the turn marker",
file=sys.stderr,
)
return False
def _route_turn_post(
url: str,
token: str,
body: dict[str, object],
timeout: float,
) -> dict[str, object] | None:
"""
POST one JSON body to the loopback ``route-turn`` endpoint.
Uses :mod:`urllib` rather than the module's ``httpx`` import so the
call stays available to a ``python -I`` hook whose interpreter may not
resolve site packages the same way the CLI's does.
:param url: Fully-qualified loopback URL.
:param token: Bearer token from the advertisement.
:param body: Request body.
:param timeout: Socket timeout in seconds.
:returns: The decoded response object, or ``None`` on any transport or
decode failure (callers treat that as "allow unrouted").
"""
import urllib.error
import urllib.request
req = urllib.request.Request(
url,
data=json.dumps(body).encode("utf-8"),
headers={"Content-Type": "application/json", "Authorization": f"Bearer {token}"},
method="POST",
)
try:
with urllib.request.urlopen(req, timeout=timeout) as resp:
decoded = json.loads(resp.read().decode("utf-8"))
except (urllib.error.URLError, OSError, ValueError, TimeoutError):
return None
return decoded if isinstance(decoded, dict) else None
if __name__ == "__main__":
raise SystemExit(main())