From 605f783d25719999728930496b4815aef024af53 Mon Sep 17 00:00:00 2001 From: CoDA PR triage Date: Fri, 7 Aug 2026 13:34:14 +1000 Subject: [PATCH 01/23] feat(omnigent): inject managed host identity from server Signed-off-by: CoDA PR triage --- app.py | 45 +++++++++++++++++++++++++- app.yaml | 3 ++ omnigents_host.py | 59 +++++++++++++++++++++++++++++++--- tests/test_auth_enforcement.py | 24 ++++++++++++++ 4 files changed, 126 insertions(+), 5 deletions(-) diff --git a/app.py b/app.py index 2cbc250f..6ab42389 100644 --- a/app.py +++ b/app.py @@ -1874,16 +1874,59 @@ def omnigent_host_status(): return jsonify(get_status()) +def _omnigent_server_request_authorized() -> bool: + """Authorize the configured Omnigent server service principal. + + Databricks Apps validates the forwarded bearer before it reaches Flask; + this check narrows the M2M endpoint to the configured server SP. + """ + expected = os.environ.get("OMNIGENT_SERVER_SP_CLIENT_ID", "").strip() + if not expected: + return False + token = ( + request.headers.get("X-Forwarded-Access-Token", "").strip() + or request.headers.get("Authorization", "").removeprefix("Bearer ").strip() + ) + try: + import base64 + import json + + payload = token.split(".")[1] + payload += "=" * (-len(payload) % 4) + claims = json.loads(base64.urlsafe_b64decode(payload)) + except (IndexError, ValueError, TypeError, json.JSONDecodeError): + return False + principals = { + str(claims.get(key, "")).strip() + for key in ("sub", "client_id", "azp", "appid") + } + return any(hmac.compare_digest(principal, expected) for principal in principals) + + @app.route("/api/omnigent-host/connect", methods=["POST"]) def omnigent_host_connect(): """Start a runtime Omnigent host tunnel for a supplied server URL.""" + if not _omnigent_server_request_authorized(): + return jsonify({"error": "Forbidden"}), 403 data = request.get_json(silent=True) or {} server_url = (data.get("server_url") or "").strip() if not server_url: return jsonify({"error": "server_url required"}), 400 from omnigents_host import connect_host - ok, status = connect_host(server_url, _omnigent_sp_creds) + + host_config = data.get("host_config") + if host_config is not None and not isinstance(host_config, dict): + return jsonify({"error": "host_config must be an object"}), 400 + ok, status = connect_host( + server_url, + _omnigent_sp_creds, + host_token=(data.get("host_token") or None), + host_id=(data.get("host_id") or None), + host_name=(data.get("host_name") or None), + host_config=host_config, + lease_id=(data.get("lease_id") or None), + ) if not ok: code = 409 if status.get("last_error") == "host already running" else 400 return jsonify(status), code diff --git a/app.yaml b/app.yaml index d4eb755a..18a008a8 100644 --- a/app.yaml +++ b/app.yaml @@ -2,6 +2,9 @@ command: - gunicorn - app:app env: + # M2M caller allowed to use /api/omnigent-host/connect. + - name: OMNIGENT_SERVER_SP_CLIENT_ID + value: "b7c82866-04b5-4d10-9667-95190f52456f" - name: HOME value: /app/python/source_code - name: ANTHROPIC_MODEL diff --git a/omnigents_host.py b/omnigents_host.py index aaa67b2d..6065942c 100644 --- a/omnigents_host.py +++ b/omnigents_host.py @@ -945,8 +945,20 @@ def _install_broker_cli_wrapper() -> None: logger.info("Installed Omnigent token-broker CLI wrapper at %s", wrapper) -def _run_host_once(server_url: str, stop_event: threading.Event | None = None) -> int: - """Run ``omnigents host`` in the foreground until it exits. Returns rc.""" +def _run_host_once( + server_url: str, + stop_event: threading.Event | None = None, + *, + host_token: str | None = None, + host_id: str | None = None, + host_name: str | None = None, + host_config: dict[str, object] | None = None, + lease_id: str | None = None, +) -> int: + """Run ``omnigents host`` in the foreground until it exits. + + Identity values are optional so legacy boot-time starts remain unchanged. + """ global _proc home = os.environ.get("HOME", "/app/python/source_code") @@ -973,6 +985,16 @@ def _run_host_once(server_url: str, stop_event: threading.Event | None = None) - broker_bin = os.path.join(home, ".coda-broker-bin") path_parts = [broker_bin, local_bin, env.get("PATH", "")] env["PATH"] = ":".join(part for part in path_parts if part) + if host_token: + env["OMNIGENT_HOST_TOKEN"] = host_token + if host_id: + env["OMNIGENT_HOST_ID"] = host_id + if host_name: + env["OMNIGENT_HOST_NAME"] = host_name + if host_config is not None: + env["OMNIGENT_HOST_CONFIG"] = json.dumps(host_config, separators=(",", ":")) + if lease_id: + env["OMNIGENT_HOST_LEASE_ID"] = lease_id stable_identity = _stable_host_identity() if stable_identity is not None: env.setdefault("OMNIGENT_HOST_ID", stable_identity[0]) @@ -1021,6 +1043,12 @@ def _supervise( server_url: str, sp_creds: dict[str, str], stop_event: threading.Event, + *, + host_token: str | None = None, + host_id: str | None = None, + host_name: str | None = None, + host_config: dict[str, object] | None = None, + lease_id: str | None = None, ) -> None: """Install, write the profile, then run the host with bounded backoff. @@ -1097,7 +1125,15 @@ def _supervise( backoff = _RESTART_BACKOFF_SECONDS while not stop_event.is_set(): try: - rc = _run_host_once(server_url, stop_event=stop_event) + rc = _run_host_once( + server_url, + stop_event=stop_event, + host_token=host_token, + host_id=host_id, + host_name=host_name, + host_config=host_config, + lease_id=lease_id, + ) if stop_event.is_set(): break logger.warning("omnigents host exited rc=%s; restarting in %ss", rc, backoff) @@ -1112,6 +1148,12 @@ def _supervise( def connect_host( server_url: str, sp_creds: dict[str, str] | None, + *, + host_token: str | None = None, + host_id: str | None = None, + host_name: str | None = None, + host_config: dict[str, object] | None = None, + lease_id: str | None = None, ) -> tuple[bool, dict[str, object]]: """Start a supervised ``omnigent host`` for a runtime-supplied server URL.""" global _sp_creds, _stop_event, _thread @@ -1150,7 +1192,16 @@ def connect_host( "last_error": None, }) _thread = threading.Thread( - target=_supervise, + target=lambda server_url, creds, stop_event: _supervise( + server_url, + creds, + stop_event, + host_token=host_token, + host_id=host_id, + host_name=host_name, + host_config=host_config, + lease_id=lease_id, + ), args=(server_url, _sp_creds, _stop_event), daemon=True, name="omnigent-host", diff --git a/tests/test_auth_enforcement.py b/tests/test_auth_enforcement.py index dd031359..5396c046 100644 --- a/tests/test_auth_enforcement.py +++ b/tests/test_auth_enforcement.py @@ -29,6 +29,30 @@ def _make_client(app_module): # --------------------------------------------------------------------------- +def test_connect_endpoint_requires_allowlisted_server_sp(monkeypatch): + """The M2M host-connect route accepts only the configured server SP.""" + import base64 + import json + + app_module = _get_app_module() + payload = base64.urlsafe_b64encode( + json.dumps({"sub": "server-sp"}).encode() + ).decode().rstrip("=") + token = f"header.{payload}.signature" + monkeypatch.setenv("OMNIGENT_SERVER_SP_CLIENT_ID", "server-sp") + + with app_module.app.test_request_context( + headers={"X-Forwarded-Access-Token": token} + ): + assert app_module._omnigent_server_request_authorized() is True + + monkeypatch.setenv("OMNIGENT_SERVER_SP_CLIENT_ID", "other-sp") + with app_module.app.test_request_context( + headers={"X-Forwarded-Access-Token": token} + ): + assert app_module._omnigent_server_request_authorized() is False + + # 1. Session endpoints MUST enforce owner check # --------------------------------------------------------------------------- From 11dd1fd59b4cb6017db4a3265a9adb036ba28c90 Mon Sep 17 00:00:00 2001 From: CoDA PR triage Date: Fri, 7 Aug 2026 14:37:08 +1000 Subject: [PATCH 02/23] feat(omnigent): fence and reap managed user leases Signed-off-by: CoDA PR triage --- app.py | 68 +++++++++++++++++--- omnigents_host.py | 116 +++++++++++++++++++++++++++++++++++ tests/test_omnigents_host.py | 38 ++++++++++++ 3 files changed, 215 insertions(+), 7 deletions(-) diff --git a/app.py b/app.py index 6ab42389..fb8dc525 100644 --- a/app.py +++ b/app.py @@ -1688,7 +1688,7 @@ def authorize_request(): # has SSO cookies — no functional regression. if request.path in ( "/health", "/api/configure-pat", "/api/inject-pat", - ) or request.path.startswith("/socket.io"): + ) or request.path.startswith(("/socket.io", "/api/omnigent-host/")): return None authorized, user = check_authorization() @@ -1903,6 +1903,39 @@ def _omnigent_server_request_authorized() -> bool: return any(hmac.compare_digest(principal, expected) for principal in principals) +@app.route("/api/omnigent-host/lease", methods=["POST"]) +def omnigent_host_lease(): + """Acquire or adopt the single user-scoped managed lease.""" + if not _omnigent_server_request_authorized(): + return jsonify({"error": "Forbidden"}), 403 + data = request.get_json(silent=True) or {} + owner = str(data.get("owner") or "").strip() + lease_id = str(data.get("lease_id") or "").strip() + if not owner or not lease_id: + return jsonify({"error": "owner and lease_id required"}), 400 + from omnigents_host import acquire_lease + + ok, lease = acquire_lease(owner, lease_id) + return jsonify(lease), (200 if ok else 409) + + +@app.route("/api/omnigent-host/workspaces", methods=["POST"]) +def omnigent_host_workspace(): + """Allocate a distinct session directory under the active lease.""" + if not _omnigent_server_request_authorized(): + return jsonify({"error": "Forbidden"}), 403 + data = request.get_json(silent=True) or {} + from omnigents_host import allocate_workspace + + try: + workspace = allocate_workspace( + str(data.get("lease_id") or ""), str(data.get("session_id") or "") + ) + except ValueError as exc: + return jsonify({"error": str(exc)}), 409 + return jsonify({"workspace": workspace}) + + @app.route("/api/omnigent-host/connect", methods=["POST"]) def omnigent_host_connect(): """Start a runtime Omnigent host tunnel for a supplied server URL.""" @@ -1913,8 +1946,11 @@ def omnigent_host_connect(): if not server_url: return jsonify({"error": "server_url required"}), 400 - from omnigents_host import connect_host + from omnigents_host import active_lease, connect_host + lease = active_lease() + if lease is None or lease.get("lease_id") != data.get("lease_id"): + return jsonify({"error": "stale or missing lease"}), 409 host_config = data.get("host_config") if host_config is not None and not isinstance(host_config, dict): return jsonify({"error": "host_config must be an object"}), 400 @@ -1930,14 +1966,31 @@ def omnigent_host_connect(): if not ok: code = 409 if status.get("last_error") == "host already running" else 400 return jsonify(status), code - return jsonify(status) + status["workspace"] = os.environ.get("HOME", "/app/python/source_code") + return jsonify(status), 202 @app.route("/api/omnigent-host/disconnect", methods=["POST"]) def omnigent_host_disconnect(): - """Stop the active runtime Omnigent host tunnel, if any.""" - from omnigents_host import disconnect_host - return jsonify(disconnect_host()) + """Release and scrub only the matching managed lease generation.""" + if not _omnigent_server_request_authorized(): + return jsonify({"error": "Forbidden"}), 403 + data = request.get_json(silent=True) or {} + lease_id = str(data.get("lease_id") or "") + from omnigents_host import active_lease, disconnect_host, release_lease + + lease = active_lease() + if lease is None or lease.get("lease_id") != lease_id: + return jsonify({"released": False, "stale": True}) + status = disconnect_host() + if data.get("scrub"): + shutil.rmtree( + os.path.join(os.environ.get("HOME", "/app/python/source_code"), "coda-sessions"), + ignore_errors=True, + ) + release_lease(lease_id) + status["released"] = True + return jsonify(status) @app.route("/api/omnigent-host/share", methods=["POST"]) @@ -2496,8 +2549,9 @@ def initialize_app(local_dev=False): # Capture the app SP's M2M OAuth creds BEFORE the strip below — the # Omnigents host tunnel needs an OAuth token (the Apps proxy rejects PATs). # No-op / returns None when disabled or creds absent. See omnigents_host.py. - from omnigents_host import capture_sp_credentials, start_host + from omnigents_host import capture_sp_credentials, start_host, start_lease_reaper _omnigent_sp_creds = capture_sp_credentials() + start_lease_reaper() # Resolve owner: Apps API (app.creator via SP) > PAT (current_user.me) app_owner = get_token_owner() diff --git a/omnigents_host.py b/omnigents_host.py index 6065942c..08fd72f9 100644 --- a/omnigents_host.py +++ b/omnigents_host.py @@ -76,6 +76,11 @@ _log_tail: list[str] = [] _LOG_TAIL_LIMIT = 80 _runner_tailer_started = False +_lease: dict[str, object] | None = None +_CODA_MAX_LEASE_S = 12 * 3600 +_CODA_IDLE_RELEASE_S = 10 * 60 +_no_runner_since: float | None = None +_lease_reaper_started = False def _stable_host_identity() -> tuple[str, str] | None: @@ -111,9 +116,117 @@ def _append_log(line: str) -> None: _status["log_tail"] = list(_log_tail) +def acquire_lease(owner: str, lease_id: str) -> tuple[bool, dict[str, object]]: + """Acquire or adopt the single user lease with an expiry fence.""" + global _lease + now = time.time() + with _lock: + if _lease is not None and float(_lease.get("expires_at", 0)) <= now: + _lease = None + if _lease is not None: + if _lease.get("owner") != owner: + return False, dict(_lease) + return True, dict(_lease) + _lease = { + "owner": owner, + "lease_id": lease_id, + "acquired_at": now, + "expires_at": now + _CODA_MAX_LEASE_S, + } + return True, dict(_lease) + + +def active_lease() -> dict[str, object] | None: + """Return the current unexpired lease, if any.""" + global _lease + with _lock: + if _lease is not None and float(_lease.get("expires_at", 0)) <= time.time(): + _lease = None + return dict(_lease) if _lease is not None else None + + +def release_lease(lease_id: str) -> bool: + """Release only the matching lease generation.""" + global _lease + with _lock: + if _lease is None or _lease.get("lease_id") != lease_id: + return False + _lease = None + return True + + +def allocate_workspace(lease_id: str, session_id: str) -> str: + """Create a session-specific directory under the active lease.""" + lease = active_lease() + if lease is None or lease.get("lease_id") != lease_id: + raise ValueError("stale lease") + allowed = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_-" + if not session_id or any(ch not in allowed for ch in session_id): + raise ValueError("invalid session id") + root = os.path.join(os.environ.get("HOME", "/app/python/source_code"), "coda-sessions") + workspace = os.path.join(root, session_id) + os.makedirs(workspace, mode=0o700, exist_ok=True) + return workspace + + +def _live_runner_count() -> int: + """Count live descendants of the supervised host process.""" + proc = _proc + if proc is None or proc.poll() is not None: + return 0 + try: + import psutil + + return sum(child.is_running() for child in psutil.Process(proc.pid).children(recursive=True)) + except Exception: + return 0 + + +def release_idle_lease(*, now: float | None = None, runner_count: int | None = None) -> bool: + """Release a lease after ten minutes with no live runner subprocesses.""" + global _no_runner_since + if active_lease() is None: + _no_runner_since = None + return False + count = _live_runner_count() if runner_count is None else runner_count + current = time.time() if now is None else now + if count > 0: + _no_runner_since = None + return False + if _no_runner_since is None: + _no_runner_since = current + return False + if current - _no_runner_since < _CODA_IDLE_RELEASE_S: + return False + lease = active_lease() + if lease is None: + return False + disconnect_host() + released = release_lease(str(lease["lease_id"])) + _no_runner_since = None + return released + + +def start_lease_reaper() -> None: + """Start the daemon safety valve that releases abandoned idle leases.""" + global _lease_reaper_started + with _lock: + if _lease_reaper_started: + return + _lease_reaper_started = True + + def _run() -> None: + while True: + time.sleep(30) + release_idle_lease() + + threading.Thread(target=_run, daemon=True, name="coda-lease-reaper").start() + + def reset_for_tests() -> None: """Reset module state between tests.""" global _proc, _sp_creds, _stop_event, _thread, _runner_tailer_started + global _lease, _no_runner_since, _lease_reaper_started if _proc is not None and _proc.poll() is None: _proc.terminate() @@ -123,6 +236,9 @@ def reset_for_tests() -> None: _stop_event = None _thread = None _runner_tailer_started = False + _lease = None + _no_runner_since = None + _lease_reaper_started = False _log_tail.clear() _status.clear() _status.update({ diff --git a/tests/test_omnigents_host.py b/tests/test_omnigents_host.py index 1dab27a5..d9bdbd35 100644 --- a/tests/test_omnigents_host.py +++ b/tests/test_omnigents_host.py @@ -8,9 +8,11 @@ from __future__ import annotations import hashlib +import os import shlex import sys +import pytest import yaml import omnigents_host as oh @@ -70,6 +72,42 @@ def test_status_initially_idle(monkeypatch): assert status["stage"] == "idle" +def test_lease_is_user_scoped_and_same_owner_adopts_existing() -> None: + oh.reset_for_tests() + ok, first = oh.acquire_lease("alice@example.com", "lease-a") + assert ok is True + ok, adopted = oh.acquire_lease("alice@example.com", "lease-b") + assert ok is True + assert adopted["lease_id"] == first["lease_id"] == "lease-a" + ok, _ = oh.acquire_lease("bob@example.com", "lease-c") + assert ok is False + assert oh.release_lease("stale") is False + assert oh.release_lease("lease-a") is True + + +def test_allocate_workspace_is_fenced_and_distinct(monkeypatch, tmp_path) -> None: + oh.reset_for_tests() + monkeypatch.setenv("HOME", str(tmp_path)) + oh.acquire_lease("alice@example.com", "lease-a") + one = oh.allocate_workspace("lease-a", "session_one") + two = oh.allocate_workspace("lease-a", "session_two") + assert one != two + assert os.path.isdir(one) + assert os.path.isdir(two) + with pytest.raises(ValueError, match="stale lease"): + oh.allocate_workspace("lease-old", "session_three") + + +def test_idle_lease_releases_after_no_runner_window(monkeypatch) -> None: + oh.reset_for_tests() + oh.acquire_lease("alice@example.com", "lease-a") + monkeypatch.setattr(oh, "disconnect_host", lambda: {}) + assert oh.release_idle_lease(now=100.0, runner_count=0) is False + assert oh.release_idle_lease(now=699.0, runner_count=0) is False + assert oh.release_idle_lease(now=700.0, runner_count=0) is True + assert oh.active_lease() is None + + def test_connect_requires_server_url(): oh.reset_for_tests() ok, status = oh.connect_host( From 376323bb4efd236a4437e396449bacfc19ac504f Mon Sep 17 00:00:00 2001 From: CoDA PR triage Date: Fri, 7 Aug 2026 15:11:35 +1000 Subject: [PATCH 03/23] fix(omnigent): reserve boot host for managed lease Signed-off-by: CoDA PR triage --- app.yaml | 2 ++ omnigents_host.py | 3 +++ tests/test_omnigents_host.py | 9 +++++++++ 3 files changed, 14 insertions(+) diff --git a/app.yaml b/app.yaml index 18a008a8..d7e3a246 100644 --- a/app.yaml +++ b/app.yaml @@ -139,6 +139,8 @@ env: # Enable CODA_DISABLE_OWNER_CHECK only in a dedicated, approved workshop overlay. # ─── Omnigent host integration ──────────────────────────────────────────── + - name: CODA_OMNIGENT_MODE + value: "managed" # Register this app as a persistent Omnigent host on boot: # initialize_app() -> start_host() dials the server as the app SP, so the # deployed app self-registers as an always-on host on every restart/redeploy. diff --git a/omnigents_host.py b/omnigents_host.py index 08fd72f9..11a923bb 100644 --- a/omnigents_host.py +++ b/omnigents_host.py @@ -1353,6 +1353,9 @@ def start_host(sp_creds: dict[str, str] | None) -> None: Runtime control should call :func:`connect_host` directly. This remains so older app.yaml deployments with ``OMNIGENTS_SERVER_URL`` still behave. """ + if os.environ.get("CODA_OMNIGENT_MODE", "external").strip().lower() == "managed": + _set(stage="idle") + return if not omnigents_host_enabled(): _set(stage="idle") return diff --git a/tests/test_omnigents_host.py b/tests/test_omnigents_host.py index d9bdbd35..8dd71aa7 100644 --- a/tests/test_omnigents_host.py +++ b/tests/test_omnigents_host.py @@ -108,6 +108,15 @@ def test_idle_lease_releases_after_no_runner_window(monkeypatch) -> None: assert oh.active_lease() is None +def test_managed_mode_skips_legacy_boot_registration(monkeypatch) -> None: + oh.reset_for_tests() + monkeypatch.setenv("CODA_OMNIGENT_MODE", "managed") + monkeypatch.setenv("OMNIGENTS_SERVER_URL", "https://omnigent.example.com") + monkeypatch.setattr(oh, "connect_host", lambda *_args, **_kwargs: (_ for _ in ()).throw(AssertionError)) + oh.start_host({"client_id": "id"}) + assert oh.get_status()["stage"] == "idle" + + def test_connect_requires_server_url(): oh.reset_for_tests() ok, status = oh.connect_host( From f63a83d2ebdda8ef4591dfe3ea29470c917c5bc3 Mon Sep 17 00:00:00 2001 From: CoDA PR triage Date: Fri, 7 Aug 2026 15:33:35 +1000 Subject: [PATCH 04/23] fix(omnigent): protect managed host status endpoint Signed-off-by: CoDA PR triage --- app.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/app.py b/app.py index fb8dc525..9f17f6f9 100644 --- a/app.py +++ b/app.py @@ -1869,7 +1869,9 @@ def omnigents_status(): @app.route("/api/omnigent-host/status") def omnigent_host_status(): - """Report runtime Omnigent host state.""" + """Report runtime Omnigent host state to the configured server SP.""" + if not _omnigent_server_request_authorized(): + return jsonify({"error": "Forbidden"}), 403 from omnigents_host import get_status return jsonify(get_status()) From ba53d561cc8e69373e4d2b4233fc761529ca502f Mon Sep 17 00:00:00 2001 From: CoDA PR triage Date: Fri, 7 Aug 2026 16:40:06 +1000 Subject: [PATCH 05/23] feat(omnigent): expose bounded runner diagnostics Signed-off-by: CoDA PR triage --- app.py | 13 +++++++++++++ omnigents_host.py | 16 ++++++++++++++++ 2 files changed, 29 insertions(+) diff --git a/app.py b/app.py index 9f17f6f9..0ff3b63b 100644 --- a/app.py +++ b/app.py @@ -1938,6 +1938,19 @@ def omnigent_host_workspace(): return jsonify({"workspace": workspace}) +@app.route("/api/omnigent-host/runner-log/") +def omnigent_host_runner_log(session_id): + """Return a bounded runner log tail to the configured server SP.""" + if not _omnigent_server_request_authorized() and get_request_user() != app_owner: + return jsonify({"error": "Forbidden"}), 403 + from omnigents_host import runner_log_tail + + try: + return jsonify({"lines": runner_log_tail(session_id)}) + except ValueError as exc: + return jsonify({"error": str(exc)}), 400 + + @app.route("/api/omnigent-host/connect", methods=["POST"]) def omnigent_host_connect(): """Start a runtime Omnigent host tunnel for a supplied server URL.""" diff --git a/omnigents_host.py b/omnigents_host.py index 11a923bb..36d1f024 100644 --- a/omnigents_host.py +++ b/omnigents_host.py @@ -93,6 +93,22 @@ def _stable_host_identity() -> tuple[str, str] | None: return f"host_{digest}", app_name +def runner_log_tail(session_id: str, *, lines: int = 80) -> list[str]: + """Return a bounded runner log tail for a validated session id.""" + if len(session_id) != 32 or any(ch not in "0123456789abcdef" for ch in session_id): + raise ValueError("invalid session id") + import glob + + home = os.environ.get("HOME", "/app/python/source_code") + matches = sorted( + glob.glob(os.path.join(home, ".omnigent", "logs", "runner", f"runner-{session_id}-*.log")) + ) + if not matches: + return [] + with open(matches[-1], errors="replace") as handle: + return handle.readlines()[-lines:] + + def get_status() -> dict[str, object]: """Return a copy of the current host-integration state.""" with _lock: From 23b12142731f6d1f04550e07a7fd3d6db18d51ba Mon Sep 17 00:00:00 2001 From: CoDA PR triage Date: Fri, 7 Aug 2026 16:50:41 +1000 Subject: [PATCH 06/23] fix(omnigent): retain full runner diagnostic tail Signed-off-by: CoDA PR triage --- omnigents_host.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/omnigents_host.py b/omnigents_host.py index 36d1f024..c3d238e5 100644 --- a/omnigents_host.py +++ b/omnigents_host.py @@ -93,7 +93,7 @@ def _stable_host_identity() -> tuple[str, str] | None: return f"host_{digest}", app_name -def runner_log_tail(session_id: str, *, lines: int = 80) -> list[str]: +def runner_log_tail(session_id: str, *, lines: int = 1000) -> list[str]: """Return a bounded runner log tail for a validated session id.""" if len(session_id) != 32 or any(ch not in "0123456789abcdef" for ch in session_id): raise ValueError("invalid session id") From 2853674eea57d8726caac4280409390d825e5ae2 Mon Sep 17 00:00:00 2001 From: CoDA PR triage Date: Sat, 8 Aug 2026 01:51:31 +1000 Subject: [PATCH 07/23] fix(omnigent): keep active runner leases alive Signed-off-by: CoDA PR triage --- pyproject.toml | 1 + requirements.lock | 23 +++++++++++++++++++++++ requirements.txt | 2 ++ tests/test_omnigents_host.py | 36 ++++++++++++++++++++++++++++++++++++ 4 files changed, 62 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 48bd5e05..c2eabb51 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,6 +16,7 @@ dependencies = [ "mlflow-skinny==3.14.0", "requests", "tomli", + "psutil>=5.9", # GHSA-g6cj-pr64-35w5 / CVE-2026-69247 — cryptography >= 44.0.0, < 50.0.0 # leaks a Bleichenbacher oracle through distinguishable errors and timing # when decrypting PKCS#7 EnvelopedData. diff --git a/requirements.lock b/requirements.lock index a89f5588..2877c9ab 100644 --- a/requirements.lock +++ b/requirements.lock @@ -574,6 +574,29 @@ protobuf==6.33.6 \ # databricks-sdk # mlflow-skinny # opentelemetry-proto +psutil==7.2.2 \ + --hash=sha256:0746f5f8d406af344fd547f1c8daa5f5c33dbc293bb8d6a16d80b4bb88f59372 \ + --hash=sha256:076a2d2f923fd4821644f5ba89f059523da90dc9014e85f8e45a5774ca5bc6f9 \ + --hash=sha256:11fe5a4f613759764e79c65cf11ebdf26e33d6dd34336f8a337aa2996d71c841 \ + --hash=sha256:1a571f2330c966c62aeda00dd24620425d4b0cc86881c89861fbc04549e5dc63 \ + --hash=sha256:1a7b04c10f32cc88ab39cbf606e117fd74721c831c98a27dc04578deb0c16979 \ + --hash=sha256:1fa4ecf83bcdf6e6c8f4449aff98eefb5d0604bf88cb883d7da3d8d2d909546a \ + --hash=sha256:2edccc433cbfa046b980b0df0171cd25bcaeb3a68fe9022db0979e7aa74a826b \ + --hash=sha256:7b6d09433a10592ce39b13d7be5a54fbac1d1228ed29abc880fb23df7cb694c9 \ + --hash=sha256:8c233660f575a5a89e6d4cb65d9f938126312bca76d8fe087b947b3a1aaac9ee \ + --hash=sha256:917e891983ca3c1887b4ef36447b1e0873e70c933afc831c6b6da078ba474312 \ + --hash=sha256:ab486563df44c17f5173621c7b198955bd6b613fb87c71c161f827d3fb149a9b \ + --hash=sha256:ae0aefdd8796a7737eccea863f80f81e468a1e4cf14d926bd9b6f5f2d5f90ca9 \ + --hash=sha256:b0726cecd84f9474419d67252add4ac0cd9811b04d61123054b9fb6f57df6e9e \ + --hash=sha256:b58fabe35e80b264a4e3bb23e6b96f9e45a3df7fb7eed419ac0e5947c61e47cc \ + --hash=sha256:c7663d4e37f13e884d13994247449e9f8f574bc4655d509c3b95e9ec9e2b9dc1 \ + --hash=sha256:e452c464a02e7dc7822a05d25db4cde564444a67e58539a00f929c51eddda0cf \ + --hash=sha256:e78c8603dcd9a04c7364f1a3e670cea95d51ee865e4efb3556a3a63adef958ea \ + --hash=sha256:eb7e81434c8d223ec4a219b5fc1c47d0417b12be7ea866e24fb5ad6e84b3d988 \ + --hash=sha256:ed0cace939114f62738d808fdcecd4c869222507e266e574799e9c0faa17d486 \ + --hash=sha256:eed63d3b4d62449571547b60578c5b2c4bcccc5387148db46e0c2313dad0ee00 \ + --hash=sha256:fd04ef36b4a6d599bbdb225dd1d3f51e00105f6d48a28f006da7f9822f2606d8 + # via -r requirements.txt pyasn1==0.6.4 \ --hash=sha256:9c447d8431c947fe4c8febc4ed9e760bc29011a5b01e5c74b67025bd9fb8ce81 \ --hash=sha256:deda9277cfd454080ec40b207fb6df82206a3a2688735233cdcd8d3d565f088b diff --git a/requirements.txt b/requirements.txt index bc784c5e..7f4c33f9 100644 --- a/requirements.txt +++ b/requirements.txt @@ -129,6 +129,8 @@ protobuf==6.33.6 # databricks-sdk # mlflow-skinny # opentelemetry-proto +psutil==7.2.2 + # via coda (pyproject.toml) pyasn1==0.6.4 # via # coda (pyproject.toml) diff --git a/tests/test_omnigents_host.py b/tests/test_omnigents_host.py index 8dd71aa7..329f8150 100644 --- a/tests/test_omnigents_host.py +++ b/tests/test_omnigents_host.py @@ -108,6 +108,42 @@ def test_idle_lease_releases_after_no_runner_window(monkeypatch) -> None: assert oh.active_lease() is None +def test_live_runner_descendant_prevents_idle_lease_release(monkeypatch) -> None: + """A live supervised child keeps the lease through the idle window.""" + oh.reset_for_tests() + oh.acquire_lease("alice@example.com", "lease-a") + + class _Child: + def is_running(self): + return True + + class _Process: + def __init__(self, _pid): + pass + + def children(self, recursive): + assert recursive is True + return [_Child()] + + class _Psutil: + Process = _Process + + class _HostProcess: + pid = 123 + + def poll(self): + return None + + monkeypatch.setitem(sys.modules, "psutil", _Psutil) + monkeypatch.setattr(oh, "_proc", _HostProcess()) + monkeypatch.setattr(oh, "disconnect_host", lambda: (_ for _ in ()).throw(AssertionError())) + + assert oh._live_runner_count() == 1 + assert oh.release_idle_lease(now=100.0) is False + assert oh.release_idle_lease(now=700.0) is False + assert oh.active_lease() is not None + + def test_managed_mode_skips_legacy_boot_registration(monkeypatch) -> None: oh.reset_for_tests() monkeypatch.setenv("CODA_OMNIGENT_MODE", "managed") From 38d74d2e7a0b781ca0f2843ae9d189ee07cf3d08 Mon Sep 17 00:00:00 2001 From: CoDA PR triage Date: Sat, 8 Aug 2026 03:05:31 +1000 Subject: [PATCH 08/23] fix(omnigent): distinguish runners from host zygote Signed-off-by: CoDA PR triage --- omnigents_host.py | 83 ++++++++++++++++-- tests/test_omnigents_host.py | 164 ++++++++++++++++++++++++++++++++--- 2 files changed, 228 insertions(+), 19 deletions(-) diff --git a/omnigents_host.py b/omnigents_host.py index c3d238e5..7bf7187f 100644 --- a/omnigents_host.py +++ b/omnigents_host.py @@ -185,17 +185,85 @@ def allocate_workspace(lease_id: str, session_id: str) -> str: return workspace -def _live_runner_count() -> int: - """Count live descendants of the supervised host process.""" +def _live_runner_count() -> int | None: + """Return live runner descendants, or ``None`` when inspection is unknown. + + The host's zygote is a persistent infrastructure process, not a runner. + Runners can either be direct host children (when zygote mode is disabled) + or children forked by that zygote. Process inspection is deliberately + fail-closed because treating an inaccessible process as absent could drop + an active user's lease. + """ proc = _proc - if proc is None or proc.poll() is not None: + if proc is None: return 0 + try: + if proc.poll() is not None: + return 0 + except Exception as exc: # pragma: no cover - defensive Popen boundary + logger.warning("unable to inspect Omnigents host process; preserving lease: %s", exc) + return None + try: import psutil - return sum(child.is_running() for child in psutil.Process(proc.pid).children(recursive=True)) - except Exception: - return 0 + try: + descendants = psutil.Process(proc.pid).children(recursive=True) + except psutil.NoSuchProcess as exc: + # Popen reported running, so a psutil disappearance here is an + # inspection race rather than definitive evidence of no runners. + logger.warning( + "unable to inspect Omnigents runner processes; preserving lease: %s", + exc, + ) + return None + + dead_statuses = {psutil.STATUS_ZOMBIE} + if hasattr(psutil, "STATUS_DEAD"): + dead_statuses.add(psutil.STATUS_DEAD) + live: list[tuple[object, list[str], int]] = [] + for child in descendants: + try: + status = child.status() + if status in dead_statuses or not child.is_running(): + continue + cmdline = child.cmdline() + ppid = child.ppid() + except psutil.NoSuchProcess: + # A child can disappear between children() and inspection. + continue + live.append((child, cmdline, ppid)) + + zygote_pids = { + child.pid + for child, cmdline, ppid in live + if ppid == proc.pid and _is_zygote_cmdline(cmdline) + } + return sum( + ppid in zygote_pids + or ( + ppid == proc.pid + and _is_runner_cmdline(cmdline) + and not _is_zygote_cmdline(cmdline) + ) + for child, cmdline, ppid in live + ) + except Exception as exc: # AccessDenied and unexpected psutil failures + logger.warning("unable to inspect Omnigents runner processes; preserving lease: %s", exc) + return None + + +def _is_zygote_cmdline(cmdline: list[str]) -> bool: + """Whether a process command line is the persistent Omnigent zygote.""" + return "omnigent.runner._zygote" in cmdline + + +def _is_runner_cmdline(cmdline: list[str]) -> bool: + """Whether a direct process is an Omnigent runner rather than infrastructure.""" + return any( + token.startswith("omnigent.runner") and token != "omnigent.runner._zygote" + for token in cmdline + ) def release_idle_lease(*, now: float | None = None, runner_count: int | None = None) -> bool: @@ -205,6 +273,9 @@ def release_idle_lease(*, now: float | None = None, runner_count: int | None = N _no_runner_since = None return False count = _live_runner_count() if runner_count is None else runner_count + if count is None: + _no_runner_since = None + return False current = time.time() if now is None else now if count > 0: _no_runner_since = None diff --git a/tests/test_omnigents_host.py b/tests/test_omnigents_host.py index 329f8150..e3f5070c 100644 --- a/tests/test_omnigents_host.py +++ b/tests/test_omnigents_host.py @@ -12,6 +12,7 @@ import shlex import sys +import psutil import pytest import yaml @@ -108,25 +109,161 @@ def test_idle_lease_releases_after_no_runner_window(monkeypatch) -> None: assert oh.active_lease() is None -def test_live_runner_descendant_prevents_idle_lease_release(monkeypatch) -> None: - """A live supervised child keeps the lease through the idle window.""" +class _FakeRunnerProcess: + def __init__(self, pid, cmdline, ppid, status=psutil.STATUS_RUNNING, running=True): + self.pid = pid + self._cmdline = cmdline + self._ppid = ppid + self._status = status + self._running = running + + def status(self): + return self._status + + def is_running(self): + return self._running + + def cmdline(self): + return self._cmdline + + def ppid(self): + return self._ppid + + +def _patch_process_tree(monkeypatch, children): + class _HostProcess: + pid = 123 + + def poll(self): + return None + + class _Process: + def __init__(self, pid): + assert pid == 123 + + def children(self, recursive): + assert recursive is True + return children + + monkeypatch.setattr(psutil, "Process", _Process) + monkeypatch.setattr(oh, "_proc", _HostProcess()) + + +def test_direct_runner_prevents_idle_lease_release(monkeypatch) -> None: + oh.reset_for_tests() + oh.acquire_lease("alice@example.com", "lease-a") + _patch_process_tree( + monkeypatch, + [_FakeRunnerProcess(201, ["python", "-m", "omnigent.runner._entry"], 123)], + ) + monkeypatch.setattr(oh, "disconnect_host", lambda: (_ for _ in ()).throw(AssertionError())) + + assert oh._live_runner_count() == 1 + assert oh.release_idle_lease(now=100.0) is False + assert oh.release_idle_lease(now=700.0) is False + assert oh.active_lease() is not None + + +def test_zygote_and_forked_runner_keep_lease_alive(monkeypatch) -> None: oh.reset_for_tests() oh.acquire_lease("alice@example.com", "lease-a") + _patch_process_tree( + monkeypatch, + [ + _FakeRunnerProcess(200, ["python", "-m", "omnigent.runner._zygote"], 123), + # A fork may inherit the zygote's exact command line; parentage + # distinguishes the runner from the infrastructure process. + _FakeRunnerProcess(201, ["python", "-m", "omnigent.runner._zygote"], 200), + ], + ) + monkeypatch.setattr(oh, "disconnect_host", lambda: (_ for _ in ()).throw(AssertionError())) - class _Child: - def is_running(self): - return True + assert oh._live_runner_count() == 1 + assert oh.release_idle_lease(now=100.0) is False + assert oh.release_idle_lease(now=700.0) is False + + +def test_zygote_alone_allows_idle_lease_release(monkeypatch) -> None: + oh.reset_for_tests() + oh.acquire_lease("alice@example.com", "lease-a") + _patch_process_tree( + monkeypatch, + [_FakeRunnerProcess(200, ["python", "-m", "omnigent.runner._zygote"], 123)], + ) + monkeypatch.setattr(oh, "disconnect_host", lambda: {}) + + assert oh._live_runner_count() == 0 + assert oh.release_idle_lease(now=100.0) is False + assert oh.release_idle_lease(now=700.0) is True + assert oh.active_lease() is None + + +def test_zombie_and_dead_descendants_are_ignored(monkeypatch) -> None: + oh.reset_for_tests() + _patch_process_tree( + monkeypatch, + [ + _FakeRunnerProcess( + 201, + ["python", "-m", "omnigent.runner"], + 123, + status=psutil.STATUS_ZOMBIE, + ), + _FakeRunnerProcess( + 202, + ["python", "-m", "omnigent.runner"], + 123, + status=getattr(psutil, "STATUS_DEAD", "dead"), + running=False, + ), + ], + ) + + assert oh._live_runner_count() == 0 + + +def test_runner_inspection_failure_resets_armed_idle_timer(monkeypatch, caplog) -> None: + oh.reset_for_tests() + oh.acquire_lease("alice@example.com", "lease-a") + + # Arm the timer, then make inspection unknown after the threshold. The + # unknown result must reset the timer rather than release the lease. + assert oh.release_idle_lease(now=100.0, runner_count=0) is False + + class _DeniedRunner(_FakeRunnerProcess): + def status(self): + raise psutil.AccessDenied(self.pid) + + _patch_process_tree( + monkeypatch, + [_DeniedRunner(201, ["python", "-m", "omnigent.runner._entry"], 123)], + ) + monkeypatch.setattr(oh, "disconnect_host", lambda: {}) + + assert oh.release_idle_lease(now=700.0) is False + assert oh.active_lease() is not None + assert "preserving lease" in caplog.text + + # A later definitive zero starts a fresh idle window; it does not inherit + # the pre-failure timer and release immediately. + _patch_process_tree(monkeypatch, []) + assert oh.release_idle_lease(now=800.0) is False + assert oh.release_idle_lease(now=1399.0) is False + assert oh.release_idle_lease(now=1400.0) is True + assert oh.active_lease() is None + + +def test_root_disappearing_during_process_inspection_preserves_lease(monkeypatch, caplog) -> None: + oh.reset_for_tests() + oh.acquire_lease("alice@example.com", "lease-a") class _Process: - def __init__(self, _pid): - pass + def __init__(self, pid): + assert pid == 123 def children(self, recursive): assert recursive is True - return [_Child()] - - class _Psutil: - Process = _Process + raise psutil.NoSuchProcess(123) class _HostProcess: pid = 123 @@ -134,14 +271,15 @@ class _HostProcess: def poll(self): return None - monkeypatch.setitem(sys.modules, "psutil", _Psutil) + monkeypatch.setattr(psutil, "Process", _Process) monkeypatch.setattr(oh, "_proc", _HostProcess()) monkeypatch.setattr(oh, "disconnect_host", lambda: (_ for _ in ()).throw(AssertionError())) - assert oh._live_runner_count() == 1 + assert oh._live_runner_count() is None assert oh.release_idle_lease(now=100.0) is False assert oh.release_idle_lease(now=700.0) is False assert oh.active_lease() is not None + assert "preserving lease" in caplog.text def test_managed_mode_skips_legacy_boot_registration(monkeypatch) -> None: From 6599953bdd49682fdd9d94b48a828a676601832c Mon Sep 17 00:00:00 2001 From: CoDA PR triage Date: Sat, 8 Aug 2026 04:12:33 +1000 Subject: [PATCH 09/23] fix(omnigent): resolve managed server identity per app Signed-off-by: CoDA PR triage --- Makefile | 9 ++-- app.yaml | 2 +- attach_omnigent_resources.sh | 84 ++++++++++++++++++++++++------------ 3 files changed, 63 insertions(+), 32 deletions(-) diff --git a/Makefile b/Makefile index 735e956b..ad232d41 100644 --- a/Makefile +++ b/Makefile @@ -226,9 +226,10 @@ redeploy-git: grant-omnigent-host deploy-git ## (Re)grant Omnigent host IAM, the OMNIGENT_SERVER_URL ?= OMNIGENT_SECRET_SCOPE ?= coda-omnigent OMNIGENT_SECRET_KEY ?= omnigent-server-url +OMNIGENT_CLIENT_ID_SECRET_KEY ?= omnigent-server-client-id -attach-omnigent-resources: ## Attach the per-app omnigent-wheels (UC Volume) + omnigent-server-url (Secret) resources the generic app.yaml resolves via valueFrom - @# The generic app.yaml references two resource keys at runtime: +attach-omnigent-resources: ## Attach the workspace-specific Omnigent volume, URL, and server-SP resources + @# The generic app.yaml references three resource keys at runtime: @# OMNIGENTS_WHEEL_SPEC valueFrom: omnigent-wheels @# OMNIGENTS_SERVER_URL valueFrom: omnigent-server-url @# This target attaches those resources to the app (merging with existing @@ -247,10 +248,12 @@ attach-omnigent-resources: ## Attach the per-app omnigent-wheels (UC Volume) + o @./attach_omnigent_resources.sh \ --profile $(PROFILE) \ --coda-app $(APP_NAME) \ + --server-app $(OMNIGENT_SERVER_APP) \ --server-url $(OMNIGENT_SERVER_URL) \ --wheel-volume $(WHEEL_VOLUME) \ --secret-scope $(OMNIGENT_SECRET_SCOPE) \ - --secret-key $(OMNIGENT_SECRET_KEY) + --secret-key $(OMNIGENT_SECRET_KEY) \ + --client-id-secret-key $(OMNIGENT_CLIENT_ID_SECRET_KEY) # ── Monitoring ─────────────────────────────────────── diff --git a/app.yaml b/app.yaml index d7e3a246..57232fc4 100644 --- a/app.yaml +++ b/app.yaml @@ -4,7 +4,7 @@ command: env: # M2M caller allowed to use /api/omnigent-host/connect. - name: OMNIGENT_SERVER_SP_CLIENT_ID - value: "b7c82866-04b5-4d10-9667-95190f52456f" + valueFrom: omnigent-server-client-id - name: HOME value: /app/python/source_code - name: ANTHROPIC_MODEL diff --git a/attach_omnigent_resources.sh b/attach_omnigent_resources.sh index 1f769fd6..bd1bcdd3 100755 --- a/attach_omnigent_resources.sh +++ b/attach_omnigent_resources.sh @@ -3,23 +3,24 @@ # valueFrom, so workspace-specific values (the Omnigent server URL, the wheel # volume) never have to be committed in app.yaml. # -# The generic app.yaml references two resource keys: -# - name: OMNIGENTS_SERVER_URL valueFrom: omnigent-server-url -# - name: OMNIGENTS_WHEEL_SPEC valueFrom: omnigent-wheels +# The generic app.yaml references three resource keys: +# - name: OMNIGENTS_SERVER_URL valueFrom: omnigent-server-url +# - name: OMNIGENT_SERVER_SP_CLIENT_ID valueFrom: omnigent-server-client-id +# - name: OMNIGENTS_WHEEL_SPEC valueFrom: omnigent-wheels # -# This script attaches those two resources to the app: -# 1. omnigent-wheels — a UC Volume resource pointing at the wheel volume -# (the same .. grant_omnigent_host.sh grants -# READ_VOLUME on). Resolves at runtime to /Volumes///. -# 2. omnigent-server-url — a Secret resource holding the Omnigent server app -# URL for this workspace. Stored in a Databricks secret scope/key (created -# if missing) because app.yaml's valueFrom can only reference secrets, not -# arbitrary strings. +# This script attaches those three resources to the app: +# 1. omnigent-wheels — a UC Volume resource pointing at the wheel volume. +# 2. omnigent-server-url — a Secret resource holding the server app URL. +# 3. omnigent-server-client-id — a Secret resource holding the server app's +# service-principal client ID, used to authorize only that M2M caller. +# +# String values use Databricks secrets because app.yaml's valueFrom can only +# reference app resources, not arbitrary workspace-specific strings. # # Uses `apps create-update resources` (the targeted field-mask patch) so # ONLY the resources field is touched — `apps update --json` is a full-body # write that clears unset fields (notably git_repository on git-linked apps). -# Merges the two resources with the app's existing ones (read → merge → write) +# Merges the three resources with the app's existing ones (read → merge → write) # to avoid clobbering unrelated resources (e.g. workshop challenge-repo-token). # # Run AFTER grant_omnigent_host.sh (which grants the SP the UC traversal it @@ -32,19 +33,23 @@ # ./attach_omnigent_resources.sh \ # --profile DEFAULT \ # --coda-app coda \ +# --server-app omnigent \ # --server-url https://omnigent-..databricksapps.com \ # --wheel-volume .. \ # --secret-scope coda-omnigent \ -# --secret-key omnigent-server-url +# --secret-key omnigent-server-url \ +# --client-id-secret-key omnigent-server-client-id set -euo pipefail PROFILE="" CODA_APP="" +SERVER_APP="" SERVER_URL="" WHEEL_VOLUME="" SECRET_SCOPE="coda-omnigent" SECRET_KEY="omnigent-server-url" +CLIENT_ID_SECRET_KEY="omnigent-server-client-id" usage() { sed -n '2,29p' "$0" | sed 's/^# \{0,1\}//' @@ -53,18 +58,20 @@ usage() { while [[ $# -gt 0 ]]; do case "$1" in - --profile) PROFILE="$2"; shift 2 ;; - --coda-app) CODA_APP="$2"; shift 2 ;; - --server-url) SERVER_URL="$2"; shift 2 ;; - --wheel-volume) WHEEL_VOLUME="$2"; shift 2 ;; - --secret-scope) SECRET_SCOPE="$2"; shift 2 ;; - --secret-key) SECRET_KEY="$2"; shift 2 ;; + --profile) PROFILE="$2"; shift 2 ;; + --coda-app) CODA_APP="$2"; shift 2 ;; + --server-app) SERVER_APP="$2"; shift 2 ;; + --server-url) SERVER_URL="$2"; shift 2 ;; + --wheel-volume) WHEEL_VOLUME="$2"; shift 2 ;; + --secret-scope) SECRET_SCOPE="$2"; shift 2 ;; + --secret-key) SECRET_KEY="$2"; shift 2 ;; + --client-id-secret-key) CLIENT_ID_SECRET_KEY="$2"; shift 2 ;; -h|--help) usage 0 ;; *) echo "unknown arg: $1" >&2; usage 1 ;; esac done -for req in PROFILE CODA_APP SERVER_URL WHEEL_VOLUME; do +for req in PROFILE CODA_APP SERVER_APP SERVER_URL WHEEL_VOLUME; do if [[ -z "${!req}" ]]; then echo "ERROR: --$(echo "$req" | tr 'A-Z_' 'a-z-') is required" >&2 usage 1 @@ -74,9 +81,18 @@ done DBX=(databricks --profile "$PROFILE") echo "==> Attaching Omnigent resources to '$CODA_APP' on profile '$PROFILE'..." +SERVER_CLIENT_ID=$("${DBX[@]}" apps get "$SERVER_APP" --output json \ + | python3 -c "import sys,json; print(json.load(sys.stdin).get('service_principal_client_id',''))") +if [[ -z "$SERVER_CLIENT_ID" ]]; then + echo "ERROR: could not resolve service principal for server app '$SERVER_APP'." >&2 + exit 1 +fi + +echo " server app: $SERVER_APP" echo " server URL: $SERVER_URL" echo " wheel volume: $WHEEL_VOLUME" -echo " secret: $SECRET_SCOPE/$SECRET_KEY" +echo " URL secret: $SECRET_SCOPE/$SECRET_KEY" +echo " SP secret: $SECRET_SCOPE/$CLIENT_ID_SECRET_KEY" # ---- 1. Store the server URL in a Databricks secret ------------------------- echo "==> Storing server URL in secret $SECRET_SCOPE/$SECRET_KEY..." @@ -88,7 +104,9 @@ echo "==> Storing server URL in secret $SECRET_SCOPE/$SECRET_KEY..." || echo " scope '$SECRET_SCOPE' already exists — reusing" # Put the secret value via stdin so it never lands on argv or in shell history. printf '%s' "$SERVER_URL" | "${DBX[@]}" secrets put-secret "$SECRET_SCOPE" "$SECRET_KEY" -echo " secret stored" +printf '%s' "$SERVER_CLIENT_ID" | "${DBX[@]}" secrets put-secret \ + "$SECRET_SCOPE" "$CLIENT_ID_SECRET_KEY" +echo " secrets stored" # ---- 2. Read the app's current resources (merge, don't replace) ------------ echo "==> Reading current resources on '$CODA_APP'..." @@ -101,7 +119,7 @@ print(json.dumps(d.get('resources') or [])) ") echo " existing resources: $(printf '%s' "$CURRENT" | python3 -c "import sys,json; print(len(json.load(sys.stdin)))")" -# ---- 3. Merge the two omnigent resources and write ------------------------- +# ---- 3. Merge the three Omnigent resources and write ----------------------- # Use the Apps SDK's create_update(app, update_mask='resources', app=App(...)) # — the targeted field-mask patch — so ONLY the resources field is touched. # The `apps update --json` CLI path is a full-body write that CLEARS unset @@ -111,15 +129,16 @@ echo " existing resources: $(printf '%s' "$CURRENT" | python3 -c "import sys, # Merge with existing resources (indexed by name) so we don't clobber unrelated # ones (e.g. workshop challenge-repo-token). echo "==> Merging + attaching resources..." -DATABRICKS_CONFIG_PROFILE="$PROFILE" python3 - "$CODA_APP" "$WHEEL_VOLUME" "$SECRET_SCOPE" "$SECRET_KEY" "$CURRENT" <<'PY' +DATABRICKS_CONFIG_PROFILE="$PROFILE" python3 - "$CODA_APP" "$WHEEL_VOLUME" \ + "$SECRET_SCOPE" "$SECRET_KEY" "$CLIENT_ID_SECRET_KEY" "$CURRENT" <<'PY' import json, os, sys from databricks.sdk import WorkspaceClient from databricks.sdk.service.apps import App, AppResource, AppResourceUcSecurable, AppResourceUcSecurableUcSecurableType, AppResourceUcSecurableUcSecurablePermission, AppResourceSecret, AppResourceSecretSecretPermission coda_app = sys.argv[1] wheel_volume = sys.argv[2] -scope, key = sys.argv[3], sys.argv[4] -current = json.loads(sys.argv[5]) +scope, url_key, client_id_key = sys.argv[3], sys.argv[4], sys.argv[5] +current = json.loads(sys.argv[6]) w = WorkspaceClient(profile=os.environ['DATABRICKS_CONFIG_PROFILE']) # Index existing resources by name so we update in place, not duplicate. @@ -134,7 +153,11 @@ by_name['omnigent-wheels'] = { } by_name['omnigent-server-url'] = { 'name': 'omnigent-server-url', - 'secret': {'scope': scope, 'key': key, 'permission': 'READ'}, + 'secret': {'scope': scope, 'key': url_key, 'permission': 'READ'}, +} +by_name['omnigent-server-client-id'] = { + 'name': 'omnigent-server-client-id', + 'secret': {'scope': scope, 'key': client_id_key, 'permission': 'READ'}, } def to_resource(d): @@ -177,12 +200,17 @@ if 'omnigent-server-url' in res: out.append('omnigent-server-url=%s/%s perm=%s' % (s.get('scope'), s.get('key'), s.get('permission'))) else: out.append('omnigent-server-url=MISSING') +if 'omnigent-server-client-id' in res: + s=res['omnigent-server-client-id'].get('secret',{}) + out.append('omnigent-server-client-id=%s/%s perm=%s' % (s.get('scope'), s.get('key'), s.get('permission'))) +else: + out.append('omnigent-server-client-id=MISSING') print(' ' + ' '.join(out)) ") echo "$FINAL" if echo "$FINAL" | grep -q MISSING; then - echo "ERROR: one or both resources did not attach — see above." >&2 + echo "ERROR: one or more resources did not attach — see above." >&2 exit 1 fi echo "==> Done. Redeploy '$CODA_APP' for the valueFrom refs to resolve." From fbd3e8006548059d147201f3e89e2112a20e9c19 Mon Sep 17 00:00:00 2001 From: CoDA PR triage Date: Sat, 8 Aug 2026 06:22:26 +1000 Subject: [PATCH 10/23] test(omnigent): exercise fenced managed control API Signed-off-by: CoDA PR triage --- tests/test_omnigents_host_api.py | 35 +++++++++++++++++++++++++------- 1 file changed, 28 insertions(+), 7 deletions(-) diff --git a/tests/test_omnigents_host_api.py b/tests/test_omnigents_host_api.py index 9578fdda..6f7ac7b3 100644 --- a/tests/test_omnigents_host_api.py +++ b/tests/test_omnigents_host_api.py @@ -6,7 +6,14 @@ def _import_app(): import importlib import app - return importlib.reload(app) + module = importlib.reload(app) + import omnigents_host + + omnigents_host.reset_for_tests() + # Endpoint behavior is exercised here; M2M authorization has dedicated + # coverage in test_auth_enforcement.py. + module._omnigent_server_request_authorized = lambda: True + return module def test_omnigent_host_status_returns_state(monkeypatch): @@ -36,23 +43,28 @@ def test_omnigent_host_connect_calls_supervisor(monkeypatch): app_module._omnigent_sp_creds = {"client_id": "c", "client_secret": "s", "host": "https://h"} called = {} - def fake_connect(url, sp_creds): + def fake_connect(url, sp_creds, **kwargs): called["url"] = url called["sp_creds"] = sp_creds + called.update(kwargs) return True, {"stage": "starting", "server_url": url} monkeypatch.setattr("omnigents_host.connect_host", fake_connect) + from omnigents_host import acquire_lease + + acquire_lease("owner@example.com", "lease-a") with app_module.app.test_client() as client: with mock.patch.object(app_module, "_is_databricks_apps", return_value=False): resp = client.post( "/api/omnigent-host/connect", - json={"server_url": "https://omnigent.example.com"}, + json={"server_url": "https://omnigent.example.com", "lease_id": "lease-a"}, ) - assert resp.status_code == 200 + assert resp.status_code == 202 assert called["url"] == "https://omnigent.example.com" assert called["sp_creds"] == app_module._omnigent_sp_creds + assert called["lease_id"] == "lease-a" def test_omnigent_host_connect_conflict(monkeypatch): @@ -60,14 +72,20 @@ def test_omnigent_host_connect_conflict(monkeypatch): app_module._omnigent_sp_creds = {"client_id": "c", "client_secret": "s", "host": "https://h"} monkeypatch.setattr( "omnigents_host.connect_host", - lambda url, sp_creds: (False, {"stage": "running", "last_error": "host already running"}), + lambda url, sp_creds, **kwargs: ( + False, + {"stage": "running", "last_error": "host already running"}, + ), ) + from omnigents_host import acquire_lease + + acquire_lease("owner@example.com", "lease-a") with app_module.app.test_client() as client: with mock.patch.object(app_module, "_is_databricks_apps", return_value=False): resp = client.post( "/api/omnigent-host/connect", - json={"server_url": "https://omnigent.example.com"}, + json={"server_url": "https://omnigent.example.com", "lease_id": "lease-a"}, ) assert resp.status_code == 409 @@ -76,10 +94,13 @@ def test_omnigent_host_connect_conflict(monkeypatch): def test_omnigent_host_disconnect_calls_supervisor(monkeypatch): app_module = _import_app() monkeypatch.setattr("omnigents_host.disconnect_host", lambda: {"stage": "stopped", "running": False}) + from omnigents_host import acquire_lease + + acquire_lease("owner@example.com", "lease-a") with app_module.app.test_client() as client: with mock.patch.object(app_module, "_is_databricks_apps", return_value=False): - resp = client.post("/api/omnigent-host/disconnect") + resp = client.post("/api/omnigent-host/disconnect", json={"lease_id": "lease-a"}) assert resp.status_code == 200 assert resp.get_json()["stage"] == "stopped" From 1d4d213b68e3d89c2fedd7502413cb5453c01d7d Mon Sep 17 00:00:00 2001 From: CoDA PR triage Date: Sat, 8 Aug 2026 06:36:08 +1000 Subject: [PATCH 11/23] fix(omnigent): trust only verified Apps identity Signed-off-by: CoDA PR triage --- app.py | 12 ++++++++---- tests/test_auth_enforcement.py | 4 ++++ tests/test_omnigents_host_api.py | 23 +++++++++++++++++++++++ 3 files changed, 35 insertions(+), 4 deletions(-) diff --git a/app.py b/app.py index 0ff3b63b..ae246c7c 100644 --- a/app.py +++ b/app.py @@ -1885,10 +1885,10 @@ def _omnigent_server_request_authorized() -> bool: expected = os.environ.get("OMNIGENT_SERVER_SP_CLIENT_ID", "").strip() if not expected: return False - token = ( - request.headers.get("X-Forwarded-Access-Token", "").strip() - or request.headers.get("Authorization", "").removeprefix("Bearer ").strip() - ) + # Only trust the Apps-proxy-injected token. Accepting a caller-supplied + # Authorization header here would make unverified JWT payload decoding an + # authorization bypass if the Flask port were ever exposed directly. + token = request.headers.get("X-Forwarded-Access-Token", "").strip() try: import base64 import json @@ -1913,8 +1913,12 @@ def omnigent_host_lease(): data = request.get_json(silent=True) or {} owner = str(data.get("owner") or "").strip() lease_id = str(data.get("lease_id") or "").strip() + requested_app = str(data.get("app_name") or "").strip() + app_name = os.environ.get("DATABRICKS_APP_NAME", "").strip() if not owner or not lease_id: return jsonify({"error": "owner and lease_id required"}), 400 + if not app_name or requested_app != app_name: + return jsonify({"error": "app_name does not match this CoDA instance"}), 409 from omnigents_host import acquire_lease ok, lease = acquire_lease(owner, lease_id) diff --git a/tests/test_auth_enforcement.py b/tests/test_auth_enforcement.py index 5396c046..4f7c37e1 100644 --- a/tests/test_auth_enforcement.py +++ b/tests/test_auth_enforcement.py @@ -52,6 +52,10 @@ def test_connect_endpoint_requires_allowlisted_server_sp(monkeypatch): ): assert app_module._omnigent_server_request_authorized() is False + monkeypatch.setenv("OMNIGENT_SERVER_SP_CLIENT_ID", "server-sp") + with app_module.app.test_request_context(headers={"Authorization": f"Bearer {token}"}): + assert app_module._omnigent_server_request_authorized() is False + # 1. Session endpoints MUST enforce owner check # --------------------------------------------------------------------------- diff --git a/tests/test_omnigents_host_api.py b/tests/test_omnigents_host_api.py index 6f7ac7b3..0fd2c45d 100644 --- a/tests/test_omnigents_host_api.py +++ b/tests/test_omnigents_host_api.py @@ -28,6 +28,29 @@ def test_omnigent_host_status_returns_state(monkeypatch): assert resp.get_json()["stage"] == "idle" +def test_omnigent_host_lease_requires_matching_app_name(monkeypatch): + app_module = _import_app() + monkeypatch.setenv("DATABRICKS_APP_NAME", "coda-main") + + with app_module.app.test_client() as client: + mismatch = client.post( + "/api/omnigent-host/lease", + json={"owner": "owner@example.com", "lease_id": "lease-a", "app_name": "coda"}, + ) + matched = client.post( + "/api/omnigent-host/lease", + json={ + "owner": "owner@example.com", + "lease_id": "lease-a", + "app_name": "coda-main", + }, + ) + + assert mismatch.status_code == 409 + assert matched.status_code == 200 + assert matched.get_json()["lease_id"] == "lease-a" + + def test_omnigent_host_connect_requires_url(): app_module = _import_app() From 757160e8c7bc56462442427be1bfb07e5f50b6cd Mon Sep 17 00:00:00 2001 From: CoDA PR triage Date: Sat, 8 Aug 2026 16:22:20 +1000 Subject: [PATCH 12/23] fix(omnigent): fence lease generation cleanup Signed-off-by: CoDA PR triage --- app.py | 15 ++++++++---- omnigents_host.py | 41 +++++++++++++++++++++++++------- static/index.html | 2 +- tests/test_omnigents_host.py | 27 +++++++++++++++++++++ tests/test_omnigents_host_api.py | 26 ++++++++++++++++++++ 5 files changed, 96 insertions(+), 15 deletions(-) diff --git a/app.py b/app.py index ae246c7c..b4e7f2d2 100644 --- a/app.py +++ b/app.py @@ -1964,6 +1964,9 @@ def omnigent_host_connect(): server_url = (data.get("server_url") or "").strip() if not server_url: return jsonify({"error": "server_url required"}), 400 + configured_server_url = os.environ.get("OMNIGENTS_SERVER_URL", "").strip() + if configured_server_url and server_url.rstrip("/") != configured_server_url.rstrip("/"): + return jsonify({"error": "server_url does not match configured Omnigent server"}), 409 from omnigents_host import active_lease, connect_host @@ -1996,17 +1999,19 @@ def omnigent_host_disconnect(): return jsonify({"error": "Forbidden"}), 403 data = request.get_json(silent=True) or {} lease_id = str(data.get("lease_id") or "") - from omnigents_host import active_lease, disconnect_host, release_lease + from omnigents_host import ( + _scrub_session_workspaces, + active_lease, + disconnect_host, + release_lease, + ) lease = active_lease() if lease is None or lease.get("lease_id") != lease_id: return jsonify({"released": False, "stale": True}) status = disconnect_host() if data.get("scrub"): - shutil.rmtree( - os.path.join(os.environ.get("HOME", "/app/python/source_code"), "coda-sessions"), - ignore_errors=True, - ) + _scrub_session_workspaces() release_lease(lease_id) status["released"] = True return jsonify(status) diff --git a/omnigents_host.py b/omnigents_host.py index 7bf7187f..887ac60a 100644 --- a/omnigents_host.py +++ b/omnigents_host.py @@ -132,14 +132,30 @@ def _append_log(line: str) -> None: _status["log_tail"] = list(_log_tail) +def _lease_expired(lease: dict[str, object], *, now: float | None = None) -> bool: + """Return whether a lease crossed its unchanged hard expiry fence.""" + current = time.time() if now is None else now + return float(lease.get("expires_at", 0)) <= current + + +def _scrub_session_workspaces() -> None: + """Remove all workspace data belonging to the released lease generation.""" + shutil.rmtree( + os.path.join(os.environ.get("HOME", "/app/python/source_code"), "coda-sessions"), + ignore_errors=True, + ) + + def acquire_lease(owner: str, lease_id: str) -> tuple[bool, dict[str, object]]: """Acquire or adopt the single user lease with an expiry fence.""" global _lease now = time.time() with _lock: - if _lease is not None and float(_lease.get("expires_at", 0)) <= now: - _lease = None if _lease is not None: + # The reaper stops and scrubs expired generations outside this + # lock. Fail closed until that cleanup completes. + if _lease_expired(_lease, now=now): + return False, dict(_lease) if _lease.get("owner") != owner: return False, dict(_lease) return True, dict(_lease) @@ -154,11 +170,10 @@ def acquire_lease(owner: str, lease_id: str) -> tuple[bool, dict[str, object]]: def active_lease() -> dict[str, object] | None: """Return the current unexpired lease, if any.""" - global _lease with _lock: - if _lease is not None and float(_lease.get("expires_at", 0)) <= time.time(): - _lease = None - return dict(_lease) if _lease is not None else None + if _lease is None or _lease_expired(_lease): + return None + return dict(_lease) def release_lease(lease_id: str) -> bool: @@ -267,16 +282,24 @@ def _is_runner_cmdline(cmdline: list[str]) -> bool: def release_idle_lease(*, now: float | None = None, runner_count: int | None = None) -> bool: - """Release a lease after ten minutes with no live runner subprocesses.""" + """Release an expired generation or one idle for ten minutes.""" global _no_runner_since - if active_lease() is None: + current = time.time() if now is None else now + with _lock: + lease_snapshot = dict(_lease) if _lease is not None else None + if lease_snapshot is None: _no_runner_since = None return False + if _lease_expired(lease_snapshot, now=current): + disconnect_host() + _scrub_session_workspaces() + released = release_lease(str(lease_snapshot["lease_id"])) + _no_runner_since = None + return released count = _live_runner_count() if runner_count is None else runner_count if count is None: _no_runner_since = None return False - current = time.time() if now is None else now if count > 0: _no_runner_since = None return False diff --git a/static/index.html b/static/index.html index 59d4a84c..bd4f9fa8 100644 --- a/static/index.html +++ b/static/index.html @@ -712,7 +712,7 @@

General

let omnigentShareFired = false; async function refreshOmnigentHostStatus() { - const resp = await fetch('/api/omnigent-host/status'); + const resp = await fetch('/api/omnigents-status'); const status = await resp.json(); renderOmnigentHostStatus(status); if (status.stage === 'running' && !omnigentShareFired) { diff --git a/tests/test_omnigents_host.py b/tests/test_omnigents_host.py index e3f5070c..59b9322d 100644 --- a/tests/test_omnigents_host.py +++ b/tests/test_omnigents_host.py @@ -99,6 +99,33 @@ def test_allocate_workspace_is_fenced_and_distinct(monkeypatch, tmp_path) -> Non oh.allocate_workspace("lease-old", "session_three") +def test_expired_lease_cleans_generation_before_owner_handoff(monkeypatch, tmp_path) -> None: + oh.reset_for_tests() + now = [100.0] + monkeypatch.setattr(oh.time, "time", lambda: now[0]) + monkeypatch.setenv("HOME", str(tmp_path)) + ok, _ = oh.acquire_lease("alice@example.com", "lease-a") + assert ok is True + workspace = oh.allocate_workspace("lease-a", "session_one") + assert os.path.isdir(workspace) + + now[0] += oh._CODA_MAX_LEASE_S + 1 + assert oh.active_lease() is None + ok, stale = oh.acquire_lease("bob@example.com", "lease-b") + assert ok is False + assert stale["lease_id"] == "lease-a" + + disconnected: list[bool] = [] + monkeypatch.setattr(oh, "disconnect_host", lambda: disconnected.append(True) or {}) + assert oh.release_idle_lease(now=now[0], runner_count=1) is True + assert disconnected == [True] + assert not os.path.exists(tmp_path / "coda-sessions") + + ok, lease = oh.acquire_lease("bob@example.com", "lease-b") + assert ok is True + assert lease["lease_id"] == "lease-b" + + def test_idle_lease_releases_after_no_runner_window(monkeypatch) -> None: oh.reset_for_tests() oh.acquire_lease("alice@example.com", "lease-a") diff --git a/tests/test_omnigents_host_api.py b/tests/test_omnigents_host_api.py index 0fd2c45d..a684bc5d 100644 --- a/tests/test_omnigents_host_api.py +++ b/tests/test_omnigents_host_api.py @@ -1,3 +1,4 @@ +from pathlib import Path from unittest import mock @@ -90,6 +91,31 @@ def fake_connect(url, sp_creds, **kwargs): assert called["lease_id"] == "lease-a" +def test_omnigent_host_connect_rejects_configured_server_mismatch(monkeypatch): + app_module = _import_app() + monkeypatch.setenv("OMNIGENTS_SERVER_URL", "https://omnigent.example.com/") + from omnigents_host import acquire_lease + + acquire_lease("owner@example.com", "lease-a") + with app_module.app.test_client() as client: + resp = client.post( + "/api/omnigent-host/connect", + json={"server_url": "https://attacker.example", "lease_id": "lease-a"}, + ) + + assert resp.status_code == 409 + assert resp.get_json() == { + "error": "server_url does not match configured Omnigent server" + } + + +def test_browser_polls_owner_authenticated_status_endpoint(): + static_html = Path(__file__).parents[1] / "static" / "index.html" + source = static_html.read_text() + assert "fetch('/api/omnigents-status')" in source + assert "fetch('/api/omnigent-host/status')" not in source + + def test_omnigent_host_connect_conflict(monkeypatch): app_module = _import_app() app_module._omnigent_sp_creds = {"client_id": "c", "client_secret": "s", "host": "https://h"} From 1993871358a276a963b95e48b4a1cf6b1bf67c8a Mon Sep 17 00:00:00 2001 From: CoDA PR triage Date: Sat, 8 Aug 2026 16:29:56 +1000 Subject: [PATCH 13/23] fix(omnigent): sanitize browser host status Signed-off-by: CoDA PR triage --- app.py | 14 ++++++++++++-- omnigents_host.py | 5 ++++- tests/test_omnigents_host_api.py | 31 +++++++++++++++++++++++++++++++ 3 files changed, 47 insertions(+), 3 deletions(-) diff --git a/app.py b/app.py index b4e7f2d2..3a35b783 100644 --- a/app.py +++ b/app.py @@ -1862,9 +1862,19 @@ def capacity_status(): @app.route("/api/omnigents-status") def omnigents_status(): - """Report Omnigents host-integration state (FR-9 observability).""" + """Report browser-safe host state without runner or host log content.""" from omnigents_host import get_status - return jsonify(get_status()) + + status = get_status() + browser_fields = ( + "configured", + "running", + "installed", + "host_launched", + "server_url", + "stage", + ) + return jsonify({key: status.get(key) for key in browser_fields}) @app.route("/api/omnigent-host/status") diff --git a/omnigents_host.py b/omnigents_host.py index 887ac60a..acaca091 100644 --- a/omnigents_host.py +++ b/omnigents_host.py @@ -328,7 +328,10 @@ def start_lease_reaper() -> None: def _run() -> None: while True: time.sleep(30) - release_idle_lease() + try: + release_idle_lease() + except Exception: + logger.exception("CoDA lease reaper cleanup failed; preserving lease") threading.Thread(target=_run, daemon=True, name="coda-lease-reaper").start() diff --git a/tests/test_omnigents_host_api.py b/tests/test_omnigents_host_api.py index a684bc5d..8b9c38ab 100644 --- a/tests/test_omnigents_host_api.py +++ b/tests/test_omnigents_host_api.py @@ -17,6 +17,37 @@ def _import_app(): return module +def test_browser_status_omits_logs_and_error_details(monkeypatch): + app_module = _import_app() + monkeypatch.setattr( + "omnigents_host.get_status", + lambda: { + "configured": True, + "running": True, + "installed": True, + "host_launched": True, + "server_url": "https://omnigent.example.com", + "stage": "running", + "last_error": "Authorization: Bearer secret", + "log_tail": ["SECRET_REPOSITORY_OUTPUT"], + }, + ) + + with app_module.app.test_client() as client: + with mock.patch.object(app_module, "_is_databricks_apps", return_value=False): + resp = client.get("/api/omnigents-status") + + assert resp.status_code == 200 + assert resp.get_json() == { + "configured": True, + "host_launched": True, + "installed": True, + "running": True, + "server_url": "https://omnigent.example.com", + "stage": "running", + } + + def test_omnigent_host_status_returns_state(monkeypatch): app_module = _import_app() monkeypatch.setattr("omnigents_host.get_status", lambda: {"stage": "idle", "running": False}) From 633aef041b8651e0225f40cc6109beab00edb1ff Mon Sep 17 00:00:00 2001 From: CoDA PR triage Date: Tue, 1 Sep 2026 23:04:48 +1000 Subject: [PATCH 14/23] feat(omnigent): gate managed host control behind opt-in mode --- app.py | 43 ++++++++++++++++++++++++++------ app.yaml | 16 +++++++++--- tests/test_auth_enforcement.py | 12 +++++++++ tests/test_omnigents_host_api.py | 39 +++++++++++++++++++++++++++-- 4 files changed, 96 insertions(+), 14 deletions(-) diff --git a/app.py b/app.py index 3a35b783..c274c2c4 100644 --- a/app.py +++ b/app.py @@ -1879,19 +1879,27 @@ def omnigents_status(): @app.route("/api/omnigent-host/status") def omnigent_host_status(): - """Report runtime Omnigent host state to the configured server SP.""" - if not _omnigent_server_request_authorized(): + """Report runtime host state; require server-SP auth in managed mode.""" + if _managed_omnigent_enabled() and not _omnigent_server_request_authorized(): return jsonify({"error": "Forbidden"}), 403 from omnigents_host import get_status return jsonify(get_status()) +def _managed_omnigent_enabled() -> bool: + """Return whether server-managed host leasing is explicitly enabled.""" + return os.environ.get("CODA_OMNIGENT_MODE", "external").strip().lower() == "managed" + + def _omnigent_server_request_authorized() -> bool: """Authorize the configured Omnigent server service principal. Databricks Apps validates the forwarded bearer before it reaches Flask; - this check narrows the M2M endpoint to the configured server SP. + this check narrows the M2M endpoint to the configured server SP. Managed + control is default-off even if a stale server-SP resource remains attached. """ + if not _managed_omnigent_enabled(): + return False expected = os.environ.get("OMNIGENT_SERVER_SP_CLIENT_ID", "").strip() if not expected: return False @@ -1918,6 +1926,8 @@ def _omnigent_server_request_authorized() -> bool: @app.route("/api/omnigent-host/lease", methods=["POST"]) def omnigent_host_lease(): """Acquire or adopt the single user-scoped managed lease.""" + if not _managed_omnigent_enabled(): + return jsonify({"error": "Managed OmniGENT mode is disabled"}), 404 if not _omnigent_server_request_authorized(): return jsonify({"error": "Forbidden"}), 403 data = request.get_json(silent=True) or {} @@ -1938,6 +1948,8 @@ def omnigent_host_lease(): @app.route("/api/omnigent-host/workspaces", methods=["POST"]) def omnigent_host_workspace(): """Allocate a distinct session directory under the active lease.""" + if not _managed_omnigent_enabled(): + return jsonify({"error": "Managed OmniGENT mode is disabled"}), 404 if not _omnigent_server_request_authorized(): return jsonify({"error": "Forbidden"}), 403 data = request.get_json(silent=True) or {} @@ -1955,6 +1967,8 @@ def omnigent_host_workspace(): @app.route("/api/omnigent-host/runner-log/") def omnigent_host_runner_log(session_id): """Return a bounded runner log tail to the configured server SP.""" + if not _managed_omnigent_enabled(): + return jsonify({"error": "Managed OmniGENT mode is disabled"}), 404 if not _omnigent_server_request_authorized() and get_request_user() != app_owner: return jsonify({"error": "Forbidden"}), 403 from omnigents_host import runner_log_tail @@ -1967,13 +1981,21 @@ def omnigent_host_runner_log(session_id): @app.route("/api/omnigent-host/connect", methods=["POST"]) def omnigent_host_connect(): - """Start a runtime Omnigent host tunnel for a supplied server URL.""" - if not _omnigent_server_request_authorized(): - return jsonify({"error": "Forbidden"}), 403 + """Start a runtime host tunnel, with leases only in managed mode.""" data = request.get_json(silent=True) or {} server_url = (data.get("server_url") or "").strip() if not server_url: return jsonify({"error": "server_url required"}), 400 + if not _managed_omnigent_enabled(): + from omnigents_host import connect_host + + ok, status = connect_host(server_url, _omnigent_sp_creds) + if not ok: + code = 409 if status.get("last_error") == "host already running" else 400 + return jsonify(status), code + return jsonify(status) + if not _omnigent_server_request_authorized(): + return jsonify({"error": "Forbidden"}), 403 configured_server_url = os.environ.get("OMNIGENTS_SERVER_URL", "").strip() if configured_server_url and server_url.rstrip("/") != configured_server_url.rstrip("/"): return jsonify({"error": "server_url does not match configured Omnigent server"}), 409 @@ -2004,7 +2026,11 @@ def omnigent_host_connect(): @app.route("/api/omnigent-host/disconnect", methods=["POST"]) def omnigent_host_disconnect(): - """Release and scrub only the matching managed lease generation.""" + """Stop an external host or release the matching managed lease.""" + if not _managed_omnigent_enabled(): + from omnigents_host import disconnect_host + + return jsonify(disconnect_host()) if not _omnigent_server_request_authorized(): return jsonify({"error": "Forbidden"}), 403 data = request.get_json(silent=True) or {} @@ -2585,7 +2611,8 @@ def initialize_app(local_dev=False): # No-op / returns None when disabled or creds absent. See omnigents_host.py. from omnigents_host import capture_sp_credentials, start_host, start_lease_reaper _omnigent_sp_creds = capture_sp_credentials() - start_lease_reaper() + if _managed_omnigent_enabled(): + start_lease_reaper() # Resolve owner: Apps API (app.creator via SP) > PAT (current_user.me) app_owner = get_token_owner() diff --git a/app.yaml b/app.yaml index 57232fc4..e0496afb 100644 --- a/app.yaml +++ b/app.yaml @@ -2,9 +2,12 @@ command: - gunicorn - app:app env: - # M2M caller allowed to use /api/omnigent-host/connect. - - name: OMNIGENT_SERVER_SP_CLIENT_ID - valueFrom: omnigent-server-client-id + # Managed OmniGENT host leasing is opt-in. To enable it, set + # CODA_OMNIGENT_MODE=managed and attach the omnigent-server-client-id + # resource, then uncomment this entry. Persistent host mode remains the + # default and does not expose the managed control plane. + # - name: OMNIGENT_SERVER_SP_CLIENT_ID + # valueFrom: omnigent-server-client-id - name: HOME value: /app/python/source_code - name: ANTHROPIC_MODEL @@ -139,8 +142,13 @@ env: # Enable CODA_DISABLE_OWNER_CHECK only in a dedicated, approved workshop overlay. # ─── Omnigent host integration ──────────────────────────────────────────── + # Host mode defaults to "external": existing deployments register one + # persistent host at boot when the resources below are attached. Managed + # leasing, workspace isolation, cleanup and server control APIs are enabled + # only when an approved deployment explicitly changes this to "managed" and + # configures OMNIGENT_SERVER_SP_CLIENT_ID above. - name: CODA_OMNIGENT_MODE - value: "managed" + value: "external" # Register this app as a persistent Omnigent host on boot: # initialize_app() -> start_host() dials the server as the app SP, so the # deployed app self-registers as an always-on host on every restart/redeploy. diff --git a/tests/test_auth_enforcement.py b/tests/test_auth_enforcement.py index 4f7c37e1..387e621f 100644 --- a/tests/test_auth_enforcement.py +++ b/tests/test_auth_enforcement.py @@ -29,6 +29,17 @@ def _make_client(app_module): # --------------------------------------------------------------------------- +def test_managed_control_is_disabled_by_default(monkeypatch): + """A stale server-SP resource cannot enable managed control on its own.""" + app_module = _get_app_module() + monkeypatch.delenv("CODA_OMNIGENT_MODE", raising=False) + monkeypatch.setenv("OMNIGENT_SERVER_SP_CLIENT_ID", "server-sp") + + with app_module.app.test_request_context(): + assert app_module._managed_omnigent_enabled() is False + assert app_module._omnigent_server_request_authorized() is False + + def test_connect_endpoint_requires_allowlisted_server_sp(monkeypatch): """The M2M host-connect route accepts only the configured server SP.""" import base64 @@ -39,6 +50,7 @@ def test_connect_endpoint_requires_allowlisted_server_sp(monkeypatch): json.dumps({"sub": "server-sp"}).encode() ).decode().rstrip("=") token = f"header.{payload}.signature" + monkeypatch.setenv("CODA_OMNIGENT_MODE", "managed") monkeypatch.setenv("OMNIGENT_SERVER_SP_CLIENT_ID", "server-sp") with app_module.app.test_request_context( diff --git a/tests/test_omnigents_host_api.py b/tests/test_omnigents_host_api.py index 8b9c38ab..098328c9 100644 --- a/tests/test_omnigents_host_api.py +++ b/tests/test_omnigents_host_api.py @@ -11,8 +11,9 @@ def _import_app(): import omnigents_host omnigents_host.reset_for_tests() - # Endpoint behavior is exercised here; M2M authorization has dedicated - # coverage in test_auth_enforcement.py. + # Managed endpoint behavior is exercised here; feature gating and M2M + # authorization have dedicated coverage in test_auth_enforcement.py. + module._managed_omnigent_enabled = lambda: True module._omnigent_server_request_authorized = lambda: True return module @@ -48,6 +49,40 @@ def test_browser_status_omits_logs_and_error_details(monkeypatch): } +def test_managed_endpoints_are_unavailable_in_external_mode(monkeypatch): + app_module = _import_app() + monkeypatch.setattr(app_module, "_managed_omnigent_enabled", lambda: False) + + with app_module.app.test_client() as client: + response = client.post("/api/omnigent-host/lease", json={}) + + assert response.status_code == 404 + + +def test_external_mode_preserves_runtime_connect(monkeypatch): + app_module = _import_app() + monkeypatch.setattr(app_module, "_managed_omnigent_enabled", lambda: False) + app_module._omnigent_sp_creds = {"client_id": "c"} + called = {} + + def fake_connect(url, sp_creds): + called.update(url=url, sp_creds=sp_creds) + return True, {"stage": "starting"} + + monkeypatch.setattr("omnigents_host.connect_host", fake_connect) + with app_module.app.test_client() as client: + response = client.post( + "/api/omnigent-host/connect", + json={"server_url": "https://omnigent.example.com"}, + ) + + assert response.status_code == 200 + assert called == { + "url": "https://omnigent.example.com", + "sp_creds": app_module._omnigent_sp_creds, + } + + def test_omnigent_host_status_returns_state(monkeypatch): app_module = _import_app() monkeypatch.setattr("omnigents_host.get_status", lambda: {"stage": "idle", "running": False}) From b457afda973a2a297c0817aaf14ae6198277d056 Mon Sep 17 00:00:00 2001 From: CoDA PR triage Date: Tue, 1 Sep 2026 23:15:24 +1000 Subject: [PATCH 15/23] refactor(omnigent): support managed mode only --- README.md | 36 +++++++++++++++++--------------- app.py | 20 +++++++----------- app.yaml | 28 ++++++++----------------- docs/agent-instructions.md | 11 ++++++---- omnigents_host.py | 15 ++++++------- tests/test_omnigents_host.py | 17 +++++++++++---- tests/test_omnigents_host_api.py | 21 +++++++------------ 7 files changed, 68 insertions(+), 80 deletions(-) diff --git a/README.md b/README.md index 08d57649..285b8876 100644 --- a/README.md +++ b/README.md @@ -185,26 +185,26 @@ Tracing setup is skipped gracefully when `APP_OWNER` is not set (e.g., local dev ## Omnigent Host Integration -CoDA can register itself as a persistent **[Omnigent](https://github.com/omnigent-ai/omnigent) agent host** — an always-on target the Omnigent server can drive coding-agent sessions into. Those sessions run *inside this container* and use the same filesystem as browser terminals. They authenticate to Databricks as the CoDA app service principal, not as the interactive browser user, so their Unity Catalog authority may differ. A deployed CoDA app becomes both an interactive terminal **and** a headless host that survives restarts and redeploys. +CoDA can act as a managed **[Omnigent](https://github.com/omnigent-ai/omnigent) agent host**. The Omnigent server acquires a fenced lease before connecting the host and launching coding-agent sessions inside the container. Each session receives a separate workspace, and abandoned idle leases are reaped automatically. Runners authenticate to Databricks as the CoDA app service principal, not as the interactive browser user, so their Unity Catalog authority may differ. -**Off by default.** With `OMNIGENTS_SERVER_URL` unset, none of this runs and CoDA behaves exactly as before. This is opt-in, environment-specific wiring — the committed `app.yaml` keeps it commented out. +**Disabled by default.** Attached resources or `OMNIGENTS_SERVER_URL` alone do not activate the integration. CoDA exposes the managed control plane only when `CODA_OMNIGENT_MODE=managed`; persistent boot-host registration is not supported. ### Turning it on -Set three variables in your deployed `app.yaml` (see `app.yaml.workshop` for a ready-to-copy overlay template): +Set managed mode and configure the server identity and host artifacts through app resources (the `attach-omnigent-resources` target attaches all three): ```yaml # app.yaml env: - # The Omnigent server this app registers against on boot. + - name: CODA_OMNIGENT_MODE + value: "managed" + # Client ID of the only Omnigent server SP allowed to control this host. + - name: OMNIGENT_SERVER_SP_CLIENT_ID + valueFrom: omnigent-server-client-id - name: OMNIGENTS_SERVER_URL - value: "https://..databricksapps.com" - # UC Volume holding the omnigent host wheels (app SP needs READ_VOLUME). + valueFrom: omnigent-server-url - name: OMNIGENTS_WHEEL_SPEC - value: "/Volumes///artifacts/wheels" - # Optional: force-reinstall the host CLI on boot while rolling out a new wheel. - - name: OMNIGENTS_FORCE_REINSTALL - value: "1" + valueFrom: omnigent-wheels ``` Before deploying, grant the CoDA app service principal `CAN_USE` on the @@ -216,7 +216,7 @@ the complete prerequisite set: make grant-omnigent-host PROFILE= APP_NAME= ``` -On boot, `initialize_app()` calls `start_host()`, which — only when `OMNIGENTS_SERVER_URL` is set — installs the `omnigents host` CLI from the wheel volume and launches it as a supervised background process that dials the server over an outbound WSS tunnel. +On boot, CoDA captures its service-principal credentials and waits in `idle`. The authorised Omnigent server acquires a lease and then asks CoDA to launch the supervised outbound WSS host tunnel. A stale lease generation cannot connect, disconnect, or scrub a newer allocation. ### Two credentials, two jobs @@ -238,15 +238,17 @@ The non-obvious part of this design is that the host uses **two separate credent ### Runtime controls -Beyond boot registration, the host can be driven at runtime: +The managed endpoints require the configured Omnigent server SP and return `404` while managed mode is disabled: | Endpoint | Method | Purpose | |----------|--------|---------| -| `/api/omnigents-status` | GET | Host-integration state (FR-9 observability) | -| `/api/omnigent-host/status` | GET | Current runtime host state | -| `/api/omnigent-host/connect` | POST | Start a host tunnel for a supplied `server_url` | -| `/api/omnigent-host/disconnect` | POST | Stop the active host tunnel | -| `/api/omnigent-host/share` | POST | Share the SP-owned host with a connecting user | +| `/api/omnigents-status` | GET | Sanitised browser-visible integration state | +| `/api/omnigent-host/status` | GET | Current managed-host state | +| `/api/omnigent-host/lease` | POST | Acquire or adopt the fenced user lease | +| `/api/omnigent-host/workspaces` | POST | Allocate a session-specific workspace | +| `/api/omnigent-host/connect` | POST | Connect the leased host tunnel | +| `/api/omnigent-host/disconnect` | POST | Release and optionally scrub the matching lease | +| `/api/omnigent-host/runner-log/` | GET | Return a bounded runner diagnostic tail | ### Related diff --git a/app.py b/app.py index c274c2c4..b0b183f0 100644 --- a/app.py +++ b/app.py @@ -1879,8 +1879,10 @@ def omnigents_status(): @app.route("/api/omnigent-host/status") def omnigent_host_status(): - """Report runtime host state; require server-SP auth in managed mode.""" - if _managed_omnigent_enabled() and not _omnigent_server_request_authorized(): + """Report managed runtime host state to the configured server SP.""" + if not _managed_omnigent_enabled(): + return jsonify({"error": "Managed OmniGENT mode is disabled"}), 404 + if not _omnigent_server_request_authorized(): return jsonify({"error": "Forbidden"}), 403 from omnigents_host import get_status return jsonify(get_status()) @@ -1888,7 +1890,7 @@ def omnigent_host_status(): def _managed_omnigent_enabled() -> bool: """Return whether server-managed host leasing is explicitly enabled.""" - return os.environ.get("CODA_OMNIGENT_MODE", "external").strip().lower() == "managed" + return os.environ.get("CODA_OMNIGENT_MODE", "disabled").strip().lower() == "managed" def _omnigent_server_request_authorized() -> bool: @@ -1987,13 +1989,7 @@ def omnigent_host_connect(): if not server_url: return jsonify({"error": "server_url required"}), 400 if not _managed_omnigent_enabled(): - from omnigents_host import connect_host - - ok, status = connect_host(server_url, _omnigent_sp_creds) - if not ok: - code = 409 if status.get("last_error") == "host already running" else 400 - return jsonify(status), code - return jsonify(status) + return jsonify({"error": "Managed OmniGENT mode is disabled"}), 404 if not _omnigent_server_request_authorized(): return jsonify({"error": "Forbidden"}), 403 configured_server_url = os.environ.get("OMNIGENTS_SERVER_URL", "").strip() @@ -2028,9 +2024,7 @@ def omnigent_host_connect(): def omnigent_host_disconnect(): """Stop an external host or release the matching managed lease.""" if not _managed_omnigent_enabled(): - from omnigents_host import disconnect_host - - return jsonify(disconnect_host()) + return jsonify({"error": "Managed OmniGENT mode is disabled"}), 404 if not _omnigent_server_request_authorized(): return jsonify({"error": "Forbidden"}), 403 data = request.get_json(silent=True) or {} diff --git a/app.yaml b/app.yaml index e0496afb..4e93c9d4 100644 --- a/app.yaml +++ b/app.yaml @@ -4,8 +4,7 @@ command: env: # Managed OmniGENT host leasing is opt-in. To enable it, set # CODA_OMNIGENT_MODE=managed and attach the omnigent-server-client-id - # resource, then uncomment this entry. Persistent host mode remains the - # default and does not expose the managed control plane. + # resource, then uncomment this entry. Integration is disabled by default. # - name: OMNIGENT_SERVER_SP_CLIENT_ID # valueFrom: omnigent-server-client-id - name: HOME @@ -142,24 +141,15 @@ env: # Enable CODA_DISABLE_OWNER_CHECK only in a dedicated, approved workshop overlay. # ─── Omnigent host integration ──────────────────────────────────────────── - # Host mode defaults to "external": existing deployments register one - # persistent host at boot when the resources below are attached. Managed - # leasing, workspace isolation, cleanup and server control APIs are enabled - # only when an approved deployment explicitly changes this to "managed" and - # configures OMNIGENT_SERVER_SP_CLIENT_ID above. + # Integration is disabled unless an approved deployment explicitly changes + # this to "managed" and configures OMNIGENT_SERVER_SP_CLIENT_ID above. + # Managed mode enables fenced leasing, workspace isolation, cleanup and the + # server-controlled host APIs; it never registers a persistent boot host. - name: CODA_OMNIGENT_MODE - value: "external" - # Register this app as a persistent Omnigent host on boot: - # initialize_app() -> start_host() dials the server as the app SP, so the - # deployed app self-registers as an always-on host on every restart/redeploy. - # - # The host is ON when the two valueFrom resources below are attached to the - # app (via `make attach-omnigent-resources` or the UI); it's OFF otherwise — - # an unresolved valueFrom yields an empty string and omnigents_host_enabled() - # returns False. So the on/off switch is "are the resources attached", not - # "did you edit this committed file" — which keeps app.yaml generic and - # git-deploy friendly, and satisfies docs/agent-instructions.md §5 (these must - # not carry one workspace's values on main). + value: "disabled" + # The resources below configure managed mode but do not enable it by + # themselves. Keeping activation separate prevents stale attached resources + # from unexpectedly exposing the managed control plane. # See attach_omnigent_resources.sh. # # OMNIGENTS_SERVER_URL — the Omnigent server app URL for THIS workspace. diff --git a/docs/agent-instructions.md b/docs/agent-instructions.md index 777f8935..f015fd92 100644 --- a/docs/agent-instructions.md +++ b/docs/agent-instructions.md @@ -228,10 +228,13 @@ exist"; stderr swallowed). On workspaces without an Omnigent server, use ## 5. Omnigent host integration -`OMNIGENTS_SERVER_URL` is the on/off switch — empty/absent means host attach is -off. Before PRing branch code to `main`, ensure `OMNIGENTS_SERVER_URL`, -`OMNIGENTS_WHEEL_SPEC`, `OMNIGENTS_FORCE_REINSTALL`, and personal-workspace -values like `CLAUDE_CODE_OTEL_CATALOG_SCHEMA` are commented out or defaulted off. +`CODA_OMNIGENT_MODE=managed` is the sole on switch; absent or any other value +means host integration is disabled. Attached resources alone must never activate +it, and persistent boot-host registration is unsupported. Before PRing branch +code to `main`, ensure the mode defaults to `disabled` and that +`OMNIGENT_SERVER_SP_CLIENT_ID`, `OMNIGENTS_SERVER_URL`, `OMNIGENTS_WHEEL_SPEC`, +`OMNIGENTS_FORCE_REINSTALL`, and personal-workspace values like +`CLAUDE_CODE_OTEL_CATALOG_SCHEMA` are commented out or defaulted off. **Liveness check:** use `GET /api/omnigents-status` on the CoDA app itself (`stage=running` + `host_launched=True`). Do **not** use `/v1/hosts` as a health diff --git a/omnigents_host.py b/omnigents_host.py index acaca091..680a7dc0 100644 --- a/omnigents_host.py +++ b/omnigents_host.py @@ -1461,18 +1461,15 @@ def disconnect_host() -> dict[str, object]: def start_host(sp_creds: dict[str, str] | None) -> None: - """Legacy boot-time wrapper around :func:`connect_host`. + """Initialise managed mode without registering a persistent boot host. - Runtime control should call :func:`connect_host` directly. This remains so - older app.yaml deployments with ``OMNIGENTS_SERVER_URL`` still behave. + The OmniGENT server connects this host only after acquiring a fenced lease. + Any value other than the explicit ``managed`` mode leaves integration off. """ - if os.environ.get("CODA_OMNIGENT_MODE", "external").strip().lower() == "managed": - _set(stage="idle") + if os.environ.get("CODA_OMNIGENT_MODE", "disabled").strip().lower() != "managed": + _set(configured=False, running=False, stage="disabled", pid=None) return - if not omnigents_host_enabled(): - _set(stage="idle") - return - connect_host(os.environ["OMNIGENTS_SERVER_URL"], sp_creds) + _set(stage="idle") def _sp_bearer(sp_creds: dict[str, str]) -> str: diff --git a/tests/test_omnigents_host.py b/tests/test_omnigents_host.py index 59b9322d..5342ae98 100644 --- a/tests/test_omnigents_host.py +++ b/tests/test_omnigents_host.py @@ -43,15 +43,15 @@ def test_start_host_noop_when_disabled(monkeypatch): monkeypatch.setattr(oh.threading, "Thread", _fail("Thread")) # Must return cleanly without touching install or threads. oh.start_host(sp_creds={"client_id": "x", "client_secret": "y", "host": "h"}) - assert oh.get_status()["stage"] == "idle" + assert oh.get_status()["stage"] == "disabled" -def test_start_host_legacy_noop_without_env(monkeypatch): +def test_start_host_stays_disabled_without_explicit_mode(monkeypatch): oh.reset_for_tests() monkeypatch.delenv("OMNIGENTS_SERVER_URL", raising=False) monkeypatch.setattr(oh, "connect_host", _fail("connect_host")) oh.start_host({"client_id": "c", "client_secret": "s", "host": "https://h"}) - assert oh.get_status()["stage"] == "idle" + assert oh.get_status()["stage"] == "disabled" def test_start_host_refuses_without_sp_creds(monkeypatch): @@ -309,7 +309,7 @@ def poll(self): assert "preserving lease" in caplog.text -def test_managed_mode_skips_legacy_boot_registration(monkeypatch) -> None: +def test_managed_mode_waits_for_a_server_lease(monkeypatch) -> None: oh.reset_for_tests() monkeypatch.setenv("CODA_OMNIGENT_MODE", "managed") monkeypatch.setenv("OMNIGENTS_SERVER_URL", "https://omnigent.example.com") @@ -318,6 +318,15 @@ def test_managed_mode_skips_legacy_boot_registration(monkeypatch) -> None: assert oh.get_status()["stage"] == "idle" +def test_host_integration_is_disabled_by_default(monkeypatch) -> None: + oh.reset_for_tests() + monkeypatch.delenv("CODA_OMNIGENT_MODE", raising=False) + monkeypatch.setenv("OMNIGENTS_SERVER_URL", "https://omnigent.example.com") + monkeypatch.setattr(oh, "connect_host", lambda *_args, **_kwargs: (_ for _ in ()).throw(AssertionError)) + oh.start_host({"client_id": "id"}) + assert oh.get_status()["stage"] == "disabled" + + def test_connect_requires_server_url(): oh.reset_for_tests() ok, status = oh.connect_host( diff --git a/tests/test_omnigents_host_api.py b/tests/test_omnigents_host_api.py index 098328c9..9dc4641a 100644 --- a/tests/test_omnigents_host_api.py +++ b/tests/test_omnigents_host_api.py @@ -49,7 +49,7 @@ def test_browser_status_omits_logs_and_error_details(monkeypatch): } -def test_managed_endpoints_are_unavailable_in_external_mode(monkeypatch): +def test_managed_endpoints_are_unavailable_when_disabled(monkeypatch): app_module = _import_app() monkeypatch.setattr(app_module, "_managed_omnigent_enabled", lambda: False) @@ -59,28 +59,21 @@ def test_managed_endpoints_are_unavailable_in_external_mode(monkeypatch): assert response.status_code == 404 -def test_external_mode_preserves_runtime_connect(monkeypatch): +def test_runtime_connect_is_unavailable_when_disabled(monkeypatch): app_module = _import_app() monkeypatch.setattr(app_module, "_managed_omnigent_enabled", lambda: False) - app_module._omnigent_sp_creds = {"client_id": "c"} - called = {} - - def fake_connect(url, sp_creds): - called.update(url=url, sp_creds=sp_creds) - return True, {"stage": "starting"} + monkeypatch.setattr( + "omnigents_host.connect_host", + lambda *_args, **_kwargs: (_ for _ in ()).throw(AssertionError), + ) - monkeypatch.setattr("omnigents_host.connect_host", fake_connect) with app_module.app.test_client() as client: response = client.post( "/api/omnigent-host/connect", json={"server_url": "https://omnigent.example.com"}, ) - assert response.status_code == 200 - assert called == { - "url": "https://omnigent.example.com", - "sp_creds": app_module._omnigent_sp_creds, - } + assert response.status_code == 404 def test_omnigent_host_status_returns_state(monkeypatch): From cc2e60116a988f1bd36af8a2a0d1a34b73527089 Mon Sep 17 00:00:00 2001 From: CoDA PR triage Date: Wed, 2 Sep 2026 00:35:24 +1000 Subject: [PATCH 16/23] fix(gateway): provision app model access and remove static pins --- Makefile | 18 ++++- app.py | 23 +++--- app.yaml | 9 +-- configure_gateway_resources.py | 89 +++++++++++++++++++++++ docs/deployment.md | 6 +- setup_claude.py | 27 +++++-- setup_pi.py | 8 +- tests/test_configure_gateway_resources.py | 51 +++++++++++++ tests/test_gateway_discovery.py | 11 ++- tests/test_gateway_models.py | 13 ++-- 10 files changed, 216 insertions(+), 39 deletions(-) create mode 100644 configure_gateway_resources.py create mode 100644 tests/test_configure_gateway_resources.py diff --git a/Makefile b/Makefile index ad232d41..1133008d 100644 --- a/Makefile +++ b/Makefile @@ -79,12 +79,12 @@ help: ## Show this help # ── Workflows ──────────────────────────────────────── -deploy: create-app grant-omnigent-host sync deploy-app ## Full deploy (create app, grant Omnigent host IAM, sync, deploy) +deploy: create-app configure-gateway-resources grant-omnigent-host sync deploy-app ## Full deploy (create app, attach Gateway models, grant Omnigent IAM, sync, deploy) @echo "" @echo "Deployment complete! App URL:" @databricks apps get $(APP_NAME) --profile $(PROFILE) --output json 2>/dev/null | python3 -c "import sys,json; print(json.load(sys.stdin).get('url','(pending)'))" -redeploy: grant-omnigent-host sync deploy-app ## Redeploy: (re)grant Omnigent host IAM + sync + deploy +redeploy: configure-gateway-resources grant-omnigent-host sync deploy-app ## Redeploy: refresh Gateway models + Omnigent IAM, sync, deploy @echo "" @echo "Redeployment complete!" @@ -206,7 +206,7 @@ configure-git-credential: ## Add a Git credential to the app SP for private repo && echo " Git credential added to SP $$sp_id for '$(APP_NAME)'." \ || echo " git-credentials create failed (credential may already exist for $(GIT_PROVIDER))." -deploy-git: ## Deploy the app from the configured Git ref ($(GIT_REF_TYPE)=$(GIT_REF)) +deploy-git: configure-gateway-resources ## Attach Gateway models, then deploy from configured Git ref ($(GIT_REF_TYPE)=$(GIT_REF)) @echo "==> Deploying '$(APP_NAME)' from Git $(GIT_REF_TYPE)='$(GIT_REF)'..." @databricks apps deploy $(APP_NAME) --profile $(PROFILE) --no-wait \ --json '{"git_source":{"$(GIT_REF_TYPE)":"$(GIT_REF)"}}' @@ -217,6 +217,18 @@ redeploy-git: grant-omnigent-host deploy-git ## (Re)grant Omnigent host IAM, the @echo "" @echo "Git redeployment complete!" +# ── AI Gateway model resources ───────────────────── + +AUTO_CONFIGURE_GATEWAY ?= true + +configure-gateway-resources: ## Grant the app SP CAN_QUERY on READY Foundation Model chat endpoints + @if [ "$(AUTO_CONFIGURE_GATEWAY)" = "true" ]; then \ + echo "==> Configuring AI Gateway model resources for '$(APP_NAME)'..."; \ + python3 configure_gateway_resources.py --profile $(PROFILE) --app $(APP_NAME); \ + else \ + echo "==> AI Gateway resource configuration disabled (AUTO_CONFIGURE_GATEWAY=$(AUTO_CONFIGURE_GATEWAY))."; \ + fi + # ── Omnigent host resources ───────────────────────── # The generic app.yaml resolves workspace-specific Omnigent values at runtime # via valueFrom resource references. This target attaches those resources to diff --git a/app.py b/app.py index b0b183f0..ba75a892 100644 --- a/app.py +++ b/app.py @@ -33,7 +33,7 @@ import app_state import enterprise_config from claude_otel import apply_claude_otel_env -from utils import add_1m_context_suffix, ensure_https, get_gateway_host +from utils import ensure_https from token_helper import write_databricks_token_wrapper from pat_rotator import PATRotator from sp_token_broker import ( @@ -784,13 +784,9 @@ def _configure_all_cli_auth(token): claude_dir = os.path.join(home, ".claude") os.makedirs(claude_dir, exist_ok=True) - gateway_host = get_gateway_host() databricks_host = ensure_https(os.environ.get("DATABRICKS_HOST", "").rstrip("/")) - - if gateway_host: - anthropic_base_url = f"{gateway_host}/anthropic" - else: - anthropic_base_url = f"{databricks_host}/serving-endpoints/anthropic" + from gateway_models import pi_base_urls + anthropic_base_url = pi_base_urls(databricks_host)["claude"] # Read-merge-write to preserve env vars from other setup scripts (e.g. setup_mlflow.py) settings_path = os.path.join(claude_dir, "settings.json") @@ -801,8 +797,10 @@ def _configure_all_cli_auth(token): settings = {} settings.setdefault("env", {}) - settings["env"]["ANTHROPIC_MODEL"] = os.environ.get("ANTHROPIC_MODEL", "databricks-claude-opus-4-8") + settings["env"].pop("ANTHROPIC_MODEL", None) settings["env"]["ANTHROPIC_BASE_URL"] = anthropic_base_url + settings["env"]["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] = "1" + settings["env"]["CLAUDE_CODE_USE_GATEWAY"] = "1" # Respect the spec-C apiKeyHelper: when it owns model auth (setup_claude.py # installed the "apiKeyHelper" key), don't re-pin a static token here — the # helper fetches its own per-TTL. Otherwise write the PAT as before. @@ -810,11 +808,8 @@ def _configure_all_cli_auth(token): settings["env"].pop("ANTHROPIC_AUTH_TOKEN", None) else: settings["env"]["ANTHROPIC_AUTH_TOKEN"] = token - # [1m] suffix requests the 1M context window via the gateway (opus/sonnet - # only; see utils.add_1m_context_suffix). Keep in sync with setup_claude.py. - settings["env"]["ANTHROPIC_DEFAULT_OPUS_MODEL"] = add_1m_context_suffix("databricks-claude-opus-4-8") - settings["env"]["ANTHROPIC_DEFAULT_SONNET_MODEL"] = add_1m_context_suffix("databricks-claude-sonnet-4-6") - settings["env"]["ANTHROPIC_DEFAULT_HAIKU_MODEL"] = "databricks-claude-haiku-4-5" + # Model inventory and family defaults are owned by setup_claude.py. Token + # rotation must not replace them with stale static model ids. settings["env"]["ANTHROPIC_CUSTOM_HEADERS"] = "x-databricks-use-coding-agent-mode: true" settings["env"]["CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS"] = "1" if apply_claude_otel_env(settings, token, databricks_host): @@ -834,7 +829,7 @@ def _configure_all_cli_auth(token): # They are idempotent: detect CLI already installed, just write config files env = {**os.environ, "DATABRICKS_TOKEN": token, "CODA_VENV_PYTHON": _venv_python()} - for script in ["setup_pi.py", "setup_codex.py", "setup_opencode.py", "setup_gemini.py", "setup_hermes.py"]: + for script in ["setup_claude.py", "setup_pi.py", "setup_codex.py", "setup_opencode.py", "setup_gemini.py", "setup_hermes.py"]: try: result = subprocess.run( [_venv_python(), script], diff --git a/app.yaml b/app.yaml index 4e93c9d4..041c6d8c 100644 --- a/app.yaml +++ b/app.yaml @@ -9,12 +9,9 @@ env: # valueFrom: omnigent-server-client-id - name: HOME value: /app/python/source_code - - name: ANTHROPIC_MODEL - value: system.ai.claude-sonnet-5 - # Pi routes to the same /anthropic gateway route as Claude Code; PI_MODEL is - # the serving-endpoint name it addresses (validated against in-geo models). - - name: PI_MODEL - value: system.ai.claude-sonnet-5 + # Claude and Pi discover their model inventory from this workspace at setup + # time. Do not pin ANTHROPIC_MODEL or PI_MODEL here: a static value collapses + # Claude's picker and can select a model the app SP cannot query. - name: GEMINI_MODEL value: system.ai.gemini-3-flash # The setup script discovers system.ai models served via the Responses API diff --git a/configure_gateway_resources.py b/configure_gateway_resources.py new file mode 100644 index 00000000..c3951154 --- /dev/null +++ b/configure_gateway_resources.py @@ -0,0 +1,89 @@ +#!/usr/bin/env python3 +"""Attach available Foundation Model endpoints to a CoDA Databricks App. + +This is the deployment-time half of CoDA's ucode-compatible Gateway setup. +The deploying identity discovers READY chat endpoints; the Apps resource API +then grants the app service principal only CAN_QUERY. Existing resources are +preserved, embeddings are excluded, and repeated runs are idempotent. +""" +from __future__ import annotations + +import argparse +from collections.abc import Iterable + +from databricks.sdk import WorkspaceClient +from databricks.sdk.service.apps import App, AppResource + +_RESOURCE_PREFIX = "gateway-model-" + + +def discover_gateway_endpoint_names(endpoints: Iterable[object]) -> list[str]: + """Return READY Foundation Model chat endpoint names.""" + names: set[str] = set() + for endpoint in endpoints: + data = endpoint.as_dict() if hasattr(endpoint, "as_dict") else endpoint + if not isinstance(data, dict) or data.get("task") != "llm/v1/chat": + continue + state = data.get("state") or {} + if state.get("ready") != "READY": + continue + entities = ((data.get("config") or {}).get("served_entities") or []) + if not any( + isinstance(entity, dict) + and str(entity.get("entity_name") or "").startswith("system.ai.") + for entity in entities + ): + continue + name = data.get("name") + if isinstance(name, str) and name.startswith("databricks-"): + names.add(name) + return sorted(names) + + +def merge_resources(current: list[dict], endpoint_names: list[str]) -> list[AppResource]: + """Preserve non-Gateway resources and replace our managed endpoint set.""" + by_name = { + resource["name"]: resource + for resource in current + if isinstance(resource, dict) + and isinstance(resource.get("name"), str) + and not resource["name"].startswith(_RESOURCE_PREFIX) + } + for endpoint_name in endpoint_names: + resource_name = _RESOURCE_PREFIX + endpoint_name.removeprefix("databricks-") + by_name[resource_name] = { + "name": resource_name, + "description": "ucode-compatible AI Gateway model access", + "serving_endpoint": { + "name": endpoint_name, + "permission": "CAN_QUERY", + }, + } + return [AppResource.from_dict(resource) for resource in by_name.values()] + + +def configure(profile: str, app_name: str) -> list[str]: + w = WorkspaceClient(profile=profile) + app = w.apps.get(app_name) + endpoint_names = discover_gateway_endpoint_names(w.serving_endpoints.list()) + if not endpoint_names: + raise RuntimeError("no READY Foundation Model chat endpoints were discovered") + current = [resource.as_dict() for resource in (app.resources or [])] + resources = merge_resources(current, endpoint_names) + w.apps.create_update(app_name, "resources", app=App(name=app_name, resources=resources)) + return endpoint_names + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--profile", required=True) + parser.add_argument("--app", required=True) + args = parser.parse_args() + endpoints = configure(args.profile, args.app) + print(f"Attached {len(endpoints)} READY chat endpoints to {args.app} with CAN_QUERY:") + for endpoint in endpoints: + print(f" {endpoint}") + + +if __name__ == "__main__": + main() diff --git a/docs/deployment.md b/docs/deployment.md index 40849f0b..95d23d01 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -20,7 +20,7 @@ The app pulls the code directly from Git. To update later, just re-deploy — it > **Note:** On first startup, the app automatically removes the template's `.git` history and reinitializes a clean, remote-free git repo. This prevents accidental pushes back to the template repo from the in-browser terminal. -> **Optional (Highly Recommended):** If you use [Databricks AI Gateway](https://docs.databricks.com/aws/en/ai-gateway/), also add `DATABRICKS_GATEWAY_HOST` as a secret or environment variable. Otherwise the app falls back to direct model serving endpoints. +> **AI Gateway setup:** CLI deployment targets automatically discover READY Foundation Model chat endpoints and attach them to the app with `CAN_QUERY`. This is required because Gateway visibility and invocation are identity-scoped; a model visible to the deployer may otherwise return a masked 404 to the app service principal. Set `AUTO_CONFIGURE_GATEWAY=false` only when permissions are managed externally. ## Alternative: Deploy with CLI @@ -75,8 +75,8 @@ make configure-git APP_NAME=coda-04 PROFILE= gh auth token | make configure-git-credential APP_NAME=coda-04 PROFILE= # 3. Deploy from a ref (branch | tag | commit) -make deploy-git APP_NAME=coda-04 PROFILE= GIT_REF=main -make redeploy-git APP_NAME=coda-04 PROFILE= GIT_REF=main # + (re)grant Omnigent IAM +make deploy-git APP_NAME=coda-04 PROFILE= GIT_REF=main # refresh Gateway CAN_QUERY resources +make redeploy-git APP_NAME=coda-04 PROFILE= GIT_REF=main # + (re)grant Omnigent IAM ``` Overridable vars: `GIT_URL`, `GIT_PROVIDER` (`gitHub`, `gitLab`, …), `GIT_REF`, diff --git a/setup_claude.py b/setup_claude.py index 5b1fc663..65618087 100644 --- a/setup_claude.py +++ b/setup_claude.py @@ -83,17 +83,32 @@ def _write_apikey_helper(claude_dir: Path) -> Path: served = catalog["anthropic"] print(f"Discovered {len(served)} anthropic-dialect model services") - requested_model = os.environ.get("ANTHROPIC_MODEL", "system.ai.claude-sonnet-5") - sonnet_model = family_model("sonnet", served, fallback=requested_model) + if not served: + print( + "ERROR: the CoDA service principal discovered no Anthropic-dialect " + "Gateway models; attach serving-endpoint resources with CAN_QUERY" + ) + raise SystemExit(1) + requested_model = os.environ.get("ANTHROPIC_MODEL", "").strip() + sonnet_model = family_model("sonnet", served, fallback=served[0]) opus_model = family_model("opus", served, fallback=sonnet_model) haiku_model = family_model("haiku", served, fallback=sonnet_model) - active_model = requested_model if requested_model in served else sonnet_model - if served and active_model != requested_model: - print(f"ANTHROPIC_MODEL={requested_model} not served here, using {active_model}") + if requested_model and requested_model not in served: + print(f"Ignoring unavailable ANTHROPIC_MODEL={requested_model}") settings.setdefault("env", {}) - settings["env"]["ANTHROPIC_MODEL"] = active_model + # Match ucode's contract: native Gateway discovery plus modelOverrides owns + # the picker. ANTHROPIC_MODEL must remain absent or Claude collapses the + # picker to that one static row. + settings["env"].pop("ANTHROPIC_MODEL", None) settings["env"]["ANTHROPIC_BASE_URL"] = anthropic_base_url + settings["env"]["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] = "1" + settings["env"]["CLAUDE_CODE_USE_GATEWAY"] = "1" + settings["modelOverrides"] = { + model.removeprefix("system.ai."): model + for model in served + if model.startswith("system.ai.claude-") + } # Token source (spec C): by default install an apiKeyHelper that fetches a # fresh token per-TTL -- Claude Code re-runs it on the interval below, so diff --git a/setup_pi.py b/setup_pi.py index c1fa3928..29e1b775 100644 --- a/setup_pi.py +++ b/setup_pi.py @@ -105,7 +105,13 @@ base_urls = pi_base_urls(host) catalog = discover_model_catalog(host, token) -claude_models = catalog["anthropic"] or [pi_model] +claude_models = catalog["anthropic"] +if not claude_models: + print( + "ERROR: the CoDA service principal discovered no Anthropic-dialect " + "Gateway models; attach serving-endpoint resources with CAN_QUERY" + ) + raise SystemExit(1) active_model = preferred_model(pi_model, claude_models) print(f"Using workspace AI Gateway: {base_urls['claude']}") print( diff --git a/tests/test_configure_gateway_resources.py b/tests/test_configure_gateway_resources.py new file mode 100644 index 00000000..ced456a0 --- /dev/null +++ b/tests/test_configure_gateway_resources.py @@ -0,0 +1,51 @@ +from configure_gateway_resources import discover_gateway_endpoint_names, merge_resources + + +def _endpoint(name, *, task="llm/v1/chat", ready="READY", entity=None): + return { + "name": name, + "task": task, + "state": {"ready": ready}, + "config": { + "served_entities": [ + {"entity_name": entity or f"system.ai.{name}"} + ] + }, + } + + +def test_discovers_only_ready_foundation_model_chat_endpoints(): + endpoints = [ + _endpoint("databricks-claude-sonnet-5"), + _endpoint("databricks-gpt-oss-120b"), + _endpoint("databricks-qwen3-embedding", task="llm/v1/embeddings"), + _endpoint("databricks-not-ready", ready="NOT_READY"), + _endpoint("custom-chat", entity="catalog.schema.model"), + ] + + assert discover_gateway_endpoint_names(endpoints) == [ + "databricks-claude-sonnet-5", + "databricks-gpt-oss-120b", + ] + + +def test_merge_preserves_unrelated_resources_and_replaces_managed_set(): + current = [ + {"name": "challenge", "secret": {"scope": "s", "key": "k", "permission": "READ"}}, + { + "name": "gateway-model-stale", + "serving_endpoint": {"name": "databricks-stale", "permission": "CAN_QUERY"}, + }, + ] + + merged = [resource.as_dict() for resource in merge_resources( + current, ["databricks-claude-sonnet-5"] + )] + by_name = {resource["name"]: resource for resource in merged} + + assert "challenge" in by_name + assert "gateway-model-stale" not in by_name + assert by_name["gateway-model-claude-sonnet-5"]["serving_endpoint"] == { + "name": "databricks-claude-sonnet-5", + "permission": "CAN_QUERY", + } diff --git a/tests/test_gateway_discovery.py b/tests/test_gateway_discovery.py index 3d24a44d..40f76092 100644 --- a/tests/test_gateway_discovery.py +++ b/tests/test_gateway_discovery.py @@ -157,7 +157,7 @@ def _run_setup(self, script_name, tmp_path, env_overrides=None): "DATABRICKS_TOKEN": "dapi_test_token", "DATABRICKS_WORKSPACE_ID": "1234567890123456", "PATH": os.environ.get("PATH", ""), - "PYTHONPATH": str(SETUP_DIR), + "PYTHONPATH": os.pathsep.join((str(tmp_path), str(SETUP_DIR))), # Pre-resolve gateway so subprocess skips the network probe "_GATEWAY_RESOLVED": "", "CODA_SKIP_CLAUDE_INSTALL": "true", @@ -167,6 +167,15 @@ def _run_setup(self, script_name, tmp_path, env_overrides=None): if env_overrides: env.update(env_overrides) + # Keep this subprocess test hermetic: production now fails closed when + # discovery returns no models, so provide a deterministic catalog rather + # than relying on the old static Sonnet fallback. + (tmp_path / "sitecustomize.py").write_text( + "import gateway_models\n" + "gateway_models.discover_model_catalog = lambda *_a, **_k: {" + "'anthropic':['system.ai.claude-sonnet-5','system.ai.claude-opus-5']," + "'anthropic_specs':[], 'openai':[], 'gemini':[], 'oss':[], 'oss_specs':[]}\n" + ) # Create required dirs (tmp_path / ".claude").mkdir(exist_ok=True) diff --git a/tests/test_gateway_models.py b/tests/test_gateway_models.py index e26c834f..60a9863b 100644 --- a/tests/test_gateway_models.py +++ b/tests/test_gateway_models.py @@ -348,7 +348,8 @@ def test_app_yaml_enables_the_default_harnesses(): default now that compatible endpoints are available; Hermes remains opt-in. """ source = (Path(__file__).parents[1] / "app.yaml").read_text() - assert source.count("value: system.ai.claude-sonnet-5") == 2 + assert "name: ANTHROPIC_MODEL" not in source + assert "name: PI_MODEL" not in source assert "value: system.ai.gpt-5" in source assert "value: system.ai.gemini-3-flash" in source assert '- name: ENABLE_CLAUDE\n value: "true"' in source @@ -393,10 +394,10 @@ def test_default_model_is_sonnet_not_opus(monkeypatch): assert gm.preferred_model("system.ai.claude-opus-5", models) == "system.ai.claude-opus-5" -def test_app_yaml_defaults_the_pickers_to_sonnet(monkeypatch): +def test_app_yaml_does_not_pin_claude_or_pi_picker_models(monkeypatch): source = (Path(__file__).parents[1] / "app.yaml").read_text() - assert "value: system.ai.claude-sonnet-5" in source - assert "value: system.ai.claude-opus-5" not in source + assert "name: ANTHROPIC_MODEL" not in source + assert "name: PI_MODEL" not in source assert 'name: ENABLE_FABLE_MODELS' in source @@ -495,7 +496,9 @@ def test_setup_claude_uses_the_workspace_gateway_and_discovered_models(): assert "pick_in_geo_model" not in source assert 'pi_base_urls(databricks_host)["claude"]' in source assert "discover_model_catalog" in source - assert '"ANTHROPIC_MODEL", "system.ai.claude-sonnet-5"' in source + assert 'settings["env"].pop("ANTHROPIC_MODEL", None)' in source + assert 'settings["env"]["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] = "1"' in source + assert 'settings["env"]["CLAUDE_CODE_USE_GATEWAY"] = "1"' in source def test_family_model_picks_newest_and_falls_back(): From 01297c67fd067b1d30ef685b847120f6bed540bf Mon Sep 17 00:00:00 2001 From: CoDA PR triage Date: Wed, 2 Sep 2026 00:36:52 +1000 Subject: [PATCH 17/23] fix(deploy): bound Gateway resource names --- configure_gateway_resources.py | 12 ++++++++++-- tests/test_configure_gateway_resources.py | 20 ++++++++++++++++---- 2 files changed, 26 insertions(+), 6 deletions(-) diff --git a/configure_gateway_resources.py b/configure_gateway_resources.py index c3951154..c6b49bb3 100644 --- a/configure_gateway_resources.py +++ b/configure_gateway_resources.py @@ -9,12 +9,20 @@ from __future__ import annotations import argparse +import hashlib from collections.abc import Iterable from databricks.sdk import WorkspaceClient from databricks.sdk.service.apps import App, AppResource -_RESOURCE_PREFIX = "gateway-model-" +_RESOURCE_PREFIX = "coda-gw-" + + +def gateway_resource_name(endpoint_name: str) -> str: + """Return a deterministic Apps resource name within the 30-char limit.""" + slug = endpoint_name.removeprefix("databricks-").replace("_", "-") + digest = hashlib.sha256(endpoint_name.encode()).hexdigest()[:8] + return f"{_RESOURCE_PREFIX}{slug[:13]}-{digest}" def discover_gateway_endpoint_names(endpoints: Iterable[object]) -> list[str]: @@ -50,7 +58,7 @@ def merge_resources(current: list[dict], endpoint_names: list[str]) -> list[AppR and not resource["name"].startswith(_RESOURCE_PREFIX) } for endpoint_name in endpoint_names: - resource_name = _RESOURCE_PREFIX + endpoint_name.removeprefix("databricks-") + resource_name = gateway_resource_name(endpoint_name) by_name[resource_name] = { "name": resource_name, "description": "ucode-compatible AI Gateway model access", diff --git a/tests/test_configure_gateway_resources.py b/tests/test_configure_gateway_resources.py index ced456a0..4ad65c40 100644 --- a/tests/test_configure_gateway_resources.py +++ b/tests/test_configure_gateway_resources.py @@ -1,4 +1,8 @@ -from configure_gateway_resources import discover_gateway_endpoint_names, merge_resources +from configure_gateway_resources import ( + discover_gateway_endpoint_names, + gateway_resource_name, + merge_resources, +) def _endpoint(name, *, task="llm/v1/chat", ready="READY", entity=None): @@ -29,11 +33,18 @@ def test_discovers_only_ready_foundation_model_chat_endpoints(): ] +def test_gateway_resource_names_are_stable_and_fit_apps_limit(): + name = gateway_resource_name("databricks-claude-sonnet-5") + assert name == gateway_resource_name("databricks-claude-sonnet-5") + assert len(name) <= 30 + assert name != gateway_resource_name("databricks-claude-sonnet-4-5") + + def test_merge_preserves_unrelated_resources_and_replaces_managed_set(): current = [ {"name": "challenge", "secret": {"scope": "s", "key": "k", "permission": "READ"}}, { - "name": "gateway-model-stale", + "name": "coda-gw-stale", "serving_endpoint": {"name": "databricks-stale", "permission": "CAN_QUERY"}, }, ] @@ -44,8 +55,9 @@ def test_merge_preserves_unrelated_resources_and_replaces_managed_set(): by_name = {resource["name"]: resource for resource in merged} assert "challenge" in by_name - assert "gateway-model-stale" not in by_name - assert by_name["gateway-model-claude-sonnet-5"]["serving_endpoint"] == { + assert "coda-gw-stale" not in by_name + resource_name = gateway_resource_name("databricks-claude-sonnet-5") + assert by_name[resource_name]["serving_endpoint"] == { "name": "databricks-claude-sonnet-5", "permission": "CAN_QUERY", } From 7705d1ff88614eb44c746ab74871d91d45b70ce2 Mon Sep 17 00:00:00 2001 From: CoDA PR triage Date: Wed, 2 Sep 2026 00:43:02 +1000 Subject: [PATCH 18/23] fix(setup): expose bounded agent configuration diagnostics --- app.py | 14 ++++++++++++++ tests/test_setup_pi.py | 29 +++++++++++++++++++++-------- 2 files changed, 35 insertions(+), 8 deletions(-) diff --git a/app.py b/app.py index ba75a892..0177eba9 100644 --- a/app.py +++ b/app.py @@ -301,9 +301,23 @@ def _run_step(step_id, command): if result.returncode == 0: if step_id == "dbcli" and os.environ.get(BROKER_URL_ENV): _ensure_broker_cli_wrapper() + if step_id in {"claude", "pi"}: + safe_markers = ( + "Using workspace AI Gateway:", + "Discovered ", + "Pi configured:", + "Claude configured:", + ) + summary = [ + line.strip() + for line in result.stdout.splitlines() + if line.strip().startswith(safe_markers) + ] + logger.info("%s setup: %s", step_id, " | ".join(summary) or "complete") _update_step(step_id, status="complete", completed_at=time.time()) else: err = result.stderr.strip() or result.stdout.strip() or "Unknown error" + logger.error("%s setup failed (rc=%s): %s", step_id, result.returncode, err[-500:]) _update_step(step_id, status="error", completed_at=time.time(), error=err[:500]) except subprocess.TimeoutExpired: _update_step(step_id, status="error", completed_at=time.time(), error="Timed out after 300s") diff --git a/tests/test_setup_pi.py b/tests/test_setup_pi.py index 96dcd34e..a2322ef0 100644 --- a/tests/test_setup_pi.py +++ b/tests/test_setup_pi.py @@ -1,9 +1,9 @@ """Tests for setup_pi.py — verify the Pi models.json config is written correctly. Runs the real setup_pi.py as a subprocess against a fake HOME. A fake `pi` -binary is pre-seeded so the npm install is skipped. Model-services discovery -fails closed against the fake workspace, so only the requested system.ai model -is configured and the write path remains deterministic without inference. +binary is pre-seeded so the npm install is skipped. A sitecustomize fixture +provides a deterministic catalogue; production fails closed when discovery is +empty rather than writing an invalid fallback model. """ import json @@ -30,11 +30,18 @@ def _seed_fake_pi_binary(home: Path): def run_setup_pi(tmp_path, env_overrides=None): + (tmp_path / "sitecustomize.py").write_text( + "import gateway_models\n" + "gateway_models.discover_model_catalog = lambda *_a, **_k: {" + "'anthropic':['system.ai.claude-sonnet-5','system.ai.claude-opus-5']," + "'anthropic_specs':[], 'openai':[], 'gemini':[], 'oss':[], 'oss_specs':[]}\n" + ) env = { "HOME": str(tmp_path), "DATABRICKS_HOST": "https://workspace.example.test", "DATABRICKS_TOKEN": "dapi_test_token", "PATH": os.environ.get("PATH", ""), + "PYTHONPATH": os.pathsep.join((str(tmp_path), str(SETUP_PI.parent))), "_GATEWAY_RESOLVED": "", } if env_overrides: @@ -69,15 +76,21 @@ def test_writes_databricks_claude_provider_schema(self, tmp_path): assert provider["baseUrl"] == "https://workspace.example.test/ai-gateway/anthropic" assert ".ai-gateway." not in provider["baseUrl"] assert provider["compat"] == {"supportsEagerToolInputStreaming": False} - assert [m["id"] for m in provider["models"]] == ["system.ai.claude-opus-5"] + assert [m["id"] for m in provider["models"]] == [ + "system.ai.claude-sonnet-5", + "system.ai.claude-opus-5", + ] # Limits and thinking come from the shared Claude version policy: opus 5 # is a >= 4.6 tier, so 1M/128k with adaptive thinking. Without # forceAdaptiveThinking Pi sends `thinking: {type: "enabled"}` and the # endpoint answers 400 "thinking.type.enabled is not supported". - assert provider["models"][0]["reasoning"] is True - assert provider["models"][0]["compat"] == {"forceAdaptiveThinking": True} - assert provider["models"][0]["contextWindow"] == 1_000_000 - assert provider["models"][0]["maxTokens"] == 128_000 + opus = {model["id"]: model for model in provider["models"]}[ + "system.ai.claude-opus-5" + ] + assert opus["reasoning"] is True + assert opus["compat"] == {"forceAdaptiveThinking": True} + assert opus["contextWindow"] == 1_000_000 + assert opus["maxTokens"] == 128_000 def test_models_json_is_chmod_600(self, tmp_path): _seed_fake_pi_binary(tmp_path) From c835d039e24ad08784d6a47a49a06b22fef4b71e Mon Sep 17 00:00:00 2001 From: CoDA PR triage Date: Wed, 2 Sep 2026 00:46:25 +1000 Subject: [PATCH 19/23] fix(gateway): union app-visible endpoints with UC models --- gateway_models.py | 14 ++++++++++++++ tests/test_gateway_models.py | 25 +++++++++++++++++++++++++ 2 files changed, 39 insertions(+) diff --git a/gateway_models.py b/gateway_models.py index 34d13498..2a3d827f 100644 --- a/gateway_models.py +++ b/gateway_models.py @@ -374,6 +374,20 @@ def discover_model_catalog(workspace: str, token: str) -> dict[str, Any]: """ ids = list_model_services(workspace, token) metadata = fetch_foundation_models(workspace, token) + # Match ucode's Gateway-only union: an app SP may have CAN_QUERY through + # attached serving-endpoint resources while lacking UC browse permission. + # In that case model-services is empty, but the app-SP-visible Foundation + # Models catalogue is authoritative. Project endpoint IDs back to the + # routable system.ai aliases used in request bodies. + gateway_ids = { + "system.ai." + canonical + for canonical, entry in metadata.items() + if entry.get("gateway_v2") is True + and isinstance(entry.get("id"), str) + and entry["id"].startswith("databricks-") + and not any(marker in canonical for marker in NON_CHAT_MARKERS) + } + ids = sorted(set(ids) | gateway_ids) # Every served version of each offered family, newest first — not just the # newest per family. A picker that lists one model per family cannot switch # to an older opus the workspace still serves, which is the whole point of diff --git a/tests/test_gateway_models.py b/tests/test_gateway_models.py index 60a9863b..75fe89cf 100644 --- a/tests/test_gateway_models.py +++ b/tests/test_gateway_models.py @@ -461,6 +461,31 @@ def test_picker_lists_only_models_the_gateway_serves_for_that_dialect(monkeypatc assert catalog["gemini"] == ["system.ai.gemini-3-pro"] +def test_gateway_catalog_supplies_models_when_app_sp_cannot_browse_uc(monkeypatch): + """CAN_QUERY resources work without broad model-services browse grants.""" + monkeypatch.setattr(gm, "list_model_services", lambda *_a, **_kw: []) + monkeypatch.setattr( + gm, + "_get_json", + lambda *_a, **_kw: _fm_payload( + [ + ("databricks-claude-sonnet-5", ["anthropic/v1/messages"]), + ("databricks-claude-opus-5", ["anthropic/v1/messages"]), + ("databricks-gpt-oss-120b", ["mlflow/v1/chat/completions"]), + ("databricks-qwen3-embedding", ["mlflow/v1/embeddings"]), + ] + ), + ) + + catalog = gm.discover_model_catalog(WORKSPACE, "tok") + + assert catalog["anthropic"] == [ + "system.ai.claude-sonnet-5", + "system.ai.claude-opus-5", + ] + assert catalog["oss"] == ["system.ai.gpt-oss-120b"] + + def test_unknown_models_are_kept_so_a_picker_never_collapses(monkeypatch): """Discovery failure must not silently reduce the picker to nothing.""" ids = ["system.ai.claude-sonnet-5", "system.ai.claude-opus-5"] From a9dedd825ad067752e8e943bbe97bd802ba767db Mon Sep 17 00:00:00 2001 From: CoDA PR triage Date: Wed, 2 Sep 2026 00:51:01 +1000 Subject: [PATCH 20/23] fix(gateway): publish app-visible model inventory --- app.yaml | 8 +++-- configure_gateway_resources.py | 37 +++++++++++++++++++++-- gateway_models.py | 32 +++++++++++++++++--- tests/test_configure_gateway_resources.py | 14 ++++++++- tests/test_gateway_models.py | 14 +++++++++ 5 files changed, 96 insertions(+), 9 deletions(-) diff --git a/app.yaml b/app.yaml index 041c6d8c..9fa18eb6 100644 --- a/app.yaml +++ b/app.yaml @@ -9,8 +9,12 @@ env: # valueFrom: omnigent-server-client-id - name: HOME value: /app/python/source_code - # Claude and Pi discover their model inventory from this workspace at setup - # time. Do not pin ANTHROPIC_MODEL or PI_MODEL here: a static value collapses + # The deploy-time Gateway configurator writes the exact CAN_QUERY inventory + # into this secret resource. Runtime unions it with live discovery because an + # app SP can invoke attached endpoints without UC model-services browse access. + - name: CODA_GATEWAY_MODEL_CATALOG + valueFrom: gateway-model-catalog + # Do not pin ANTHROPIC_MODEL or PI_MODEL here: a static value collapses # Claude's picker and can select a model the app SP cannot query. - name: GEMINI_MODEL value: system.ai.gemini-3-flash diff --git a/configure_gateway_resources.py b/configure_gateway_resources.py index c6b49bb3..f1b90db6 100644 --- a/configure_gateway_resources.py +++ b/configure_gateway_resources.py @@ -10,12 +10,16 @@ import argparse import hashlib +import json from collections.abc import Iterable from databricks.sdk import WorkspaceClient +from databricks.sdk.errors import ResourceAlreadyExists from databricks.sdk.service.apps import App, AppResource _RESOURCE_PREFIX = "coda-gw-" +_CATALOG_RESOURCE = "gateway-model-catalog" +_CATALOG_SCOPE = "coda-gateway" def gateway_resource_name(endpoint_name: str) -> str: @@ -48,7 +52,14 @@ def discover_gateway_endpoint_names(endpoints: Iterable[object]) -> list[str]: return sorted(names) -def merge_resources(current: list[dict], endpoint_names: list[str]) -> list[AppResource]: +def endpoint_model_id(endpoint_name: str) -> str: + """Map a Foundation Model endpoint name to its routable request model ID.""" + return "system.ai." + endpoint_name.removeprefix("databricks-") + + +def merge_resources( + current: list[dict], endpoint_names: list[str], *, catalog_secret_key: str +) -> list[AppResource]: """Preserve non-Gateway resources and replace our managed endpoint set.""" by_name = { resource["name"]: resource @@ -67,6 +78,15 @@ def merge_resources(current: list[dict], endpoint_names: list[str]) -> list[AppR "permission": "CAN_QUERY", }, } + by_name[_CATALOG_RESOURCE] = { + "name": _CATALOG_RESOURCE, + "description": "Gateway models granted to this CoDA app", + "secret": { + "scope": _CATALOG_SCOPE, + "key": catalog_secret_key, + "permission": "READ", + }, + } return [AppResource.from_dict(resource) for resource in by_name.values()] @@ -76,8 +96,21 @@ def configure(profile: str, app_name: str) -> list[str]: endpoint_names = discover_gateway_endpoint_names(w.serving_endpoints.list()) if not endpoint_names: raise RuntimeError("no READY Foundation Model chat endpoints were discovered") + catalog_secret_key = f"{app_name}-model-catalog" + try: + w.secrets.create_scope(_CATALOG_SCOPE) + except ResourceAlreadyExists: + pass + model_ids = [endpoint_model_id(name) for name in endpoint_names] + w.secrets.put_secret( + _CATALOG_SCOPE, + catalog_secret_key, + string_value=json.dumps(model_ids, separators=(",", ":")), + ) current = [resource.as_dict() for resource in (app.resources or [])] - resources = merge_resources(current, endpoint_names) + resources = merge_resources( + current, endpoint_names, catalog_secret_key=catalog_secret_key + ) w.apps.create_update(app_name, "resources", app=App(name=app_name, resources=resources)) return endpoint_names diff --git a/gateway_models.py b/gateway_models.py index 2a3d827f..94e16e61 100644 --- a/gateway_models.py +++ b/gateway_models.py @@ -6,6 +6,7 @@ """ from __future__ import annotations +import json import os import re from typing import Any @@ -365,6 +366,26 @@ def discover_oss_specs( return [specs[key] for key in sorted(specs)] +def configured_model_ids() -> list[str]: + """Return the deployment-time inventory attached with CAN_QUERY.""" + raw = os.environ.get("CODA_GATEWAY_MODEL_CATALOG", "").strip() + if not raw: + return [] + try: + values = json.loads(raw) + except (TypeError, ValueError): + return [] + if not isinstance(values, list): + return [] + return sorted( + { + value.strip() + for value in values + if isinstance(value, str) and value.strip().startswith("system.ai.") + } + ) + + def discover_model_catalog(workspace: str, token: str) -> dict[str, Any]: """Bucket current model services using ucode's provider precedence. @@ -387,7 +408,7 @@ def discover_model_catalog(workspace: str, token: str) -> dict[str, Any]: and entry["id"].startswith("databricks-") and not any(marker in canonical for marker in NON_CHAT_MARKERS) } - ids = sorted(set(ids) | gateway_ids) + ids = sorted(set(ids) | gateway_ids | set(configured_model_ids())) # Every served version of each offered family, newest first — not just the # newest per family. A picker that lists one model per family cannot switch # to an older opus the workspace still serves, which is the whole point of @@ -422,12 +443,15 @@ def discover_model_catalog(workspace: str, token: str) -> dict[str, Any]: for model in ids: canonical = _canonical(model) spec = specs_by_id.get(canonical) - if spec is None and any(family in canonical for family in OSS_STATIC_FAMILIES): + if spec is None and ( + any(family in canonical for family in OSS_STATIC_FAMILIES) + or canonical in OSS_OUTPUT_LIMITS + ): spec = { "id": model, - "reasoning": True, + "reasoning": canonical.startswith(("gpt-oss-", "kimi-", "glm-")), "context_window": 1_000_000 if "glm-5-2" in canonical else 128_000, - "max_tokens": 65_536 if "kimi" in canonical or "glm-5-2" in canonical else 8_192, + "max_tokens": OSS_OUTPUT_LIMITS.get(canonical, 8_192), } if spec is not None: oss.append(model) diff --git a/tests/test_configure_gateway_resources.py b/tests/test_configure_gateway_resources.py index 4ad65c40..5026952b 100644 --- a/tests/test_configure_gateway_resources.py +++ b/tests/test_configure_gateway_resources.py @@ -1,5 +1,6 @@ from configure_gateway_resources import ( discover_gateway_endpoint_names, + endpoint_model_id, gateway_resource_name, merge_resources, ) @@ -33,6 +34,10 @@ def test_discovers_only_ready_foundation_model_chat_endpoints(): ] +def test_endpoint_names_map_to_routable_model_ids(): + assert endpoint_model_id("databricks-claude-sonnet-5") == "system.ai.claude-sonnet-5" + + def test_gateway_resource_names_are_stable_and_fit_apps_limit(): name = gateway_resource_name("databricks-claude-sonnet-5") assert name == gateway_resource_name("databricks-claude-sonnet-5") @@ -50,7 +55,9 @@ def test_merge_preserves_unrelated_resources_and_replaces_managed_set(): ] merged = [resource.as_dict() for resource in merge_resources( - current, ["databricks-claude-sonnet-5"] + current, + ["databricks-claude-sonnet-5"], + catalog_secret_key="coda-model-catalog", )] by_name = {resource["name"]: resource for resource in merged} @@ -61,3 +68,8 @@ def test_merge_preserves_unrelated_resources_and_replaces_managed_set(): "name": "databricks-claude-sonnet-5", "permission": "CAN_QUERY", } + assert by_name["gateway-model-catalog"]["secret"] == { + "scope": "coda-gateway", + "key": "coda-model-catalog", + "permission": "READ", + } diff --git a/tests/test_gateway_models.py b/tests/test_gateway_models.py index 75fe89cf..64addaf3 100644 --- a/tests/test_gateway_models.py +++ b/tests/test_gateway_models.py @@ -461,6 +461,20 @@ def test_picker_lists_only_models_the_gateway_serves_for_that_dialect(monkeypatc assert catalog["gemini"] == ["system.ai.gemini-3-pro"] +def test_deployment_catalog_supplies_models_when_app_sp_cannot_list_them(monkeypatch): + monkeypatch.setenv( + "CODA_GATEWAY_MODEL_CATALOG", + '["system.ai.claude-sonnet-5","system.ai.gpt-oss-120b"]', + ) + monkeypatch.setattr(gm, "list_model_services", lambda *_a, **_kw: []) + monkeypatch.setattr(gm, "fetch_foundation_models", lambda *_a, **_kw: {}) + + catalog = gm.discover_model_catalog(WORKSPACE, "tok") + + assert catalog["anthropic"] == ["system.ai.claude-sonnet-5"] + assert catalog["oss"] == ["system.ai.gpt-oss-120b"] + + def test_gateway_catalog_supplies_models_when_app_sp_cannot_browse_uc(monkeypatch): """CAN_QUERY resources work without broad model-services browse grants.""" monkeypatch.setattr(gm, "list_model_services", lambda *_a, **_kw: []) From 9cd6d4a8315c15817d7b355c76242e1274071791 Mon Sep 17 00:00:00 2001 From: CoDA PR triage Date: Wed, 2 Sep 2026 01:04:27 +1000 Subject: [PATCH 21/23] fix(deploy): grant Unity model-service execution --- Makefile | 2 +- configure_gateway_resources.py | 25 ++++++++++++++++++++++++- 2 files changed, 25 insertions(+), 2 deletions(-) diff --git a/Makefile b/Makefile index 1133008d..7d915ab6 100644 --- a/Makefile +++ b/Makefile @@ -221,7 +221,7 @@ redeploy-git: grant-omnigent-host deploy-git ## (Re)grant Omnigent host IAM, the AUTO_CONFIGURE_GATEWAY ?= true -configure-gateway-resources: ## Grant the app SP CAN_QUERY on READY Foundation Model chat endpoints +configure-gateway-resources: ## Grant app SP UC EXECUTE (v3) + CAN_QUERY (legacy) on READY chat models @if [ "$(AUTO_CONFIGURE_GATEWAY)" = "true" ]; then \ echo "==> Configuring AI Gateway model resources for '$(APP_NAME)'..."; \ python3 configure_gateway_resources.py --profile $(PROFILE) --app $(APP_NAME); \ diff --git a/configure_gateway_resources.py b/configure_gateway_resources.py index f1b90db6..41c850dd 100644 --- a/configure_gateway_resources.py +++ b/configure_gateway_resources.py @@ -16,6 +16,7 @@ from databricks.sdk import WorkspaceClient from databricks.sdk.errors import ResourceAlreadyExists from databricks.sdk.service.apps import App, AppResource +from databricks.sdk.service.catalog import PermissionsChange, Privilege _RESOURCE_PREFIX = "coda-gw-" _CATALOG_RESOURCE = "gateway-model-catalog" @@ -102,6 +103,25 @@ def configure(profile: str, app_name: str) -> list[str]: except ResourceAlreadyExists: pass model_ids = [endpoint_model_id(name) for name in endpoint_names] + principal = app.service_principal_client_id + if not principal: + raise RuntimeError(f"app {app_name} has no service principal client ID") + w.grants.update( + "catalog", + "system", + changes=[PermissionsChange(principal=principal, add=[Privilege.USE_CATALOG])], + ) + w.grants.update( + "schema", + "system.ai", + changes=[PermissionsChange(principal=principal, add=[Privilege.USE_SCHEMA])], + ) + for model_id in model_ids: + w.grants.update( + "model_service", + model_id, + changes=[PermissionsChange(principal=principal, add=[Privilege.EXECUTE])], + ) w.secrets.put_secret( _CATALOG_SCOPE, catalog_secret_key, @@ -121,7 +141,10 @@ def main() -> None: parser.add_argument("--app", required=True) args = parser.parse_args() endpoints = configure(args.profile, args.app) - print(f"Attached {len(endpoints)} READY chat endpoints to {args.app} with CAN_QUERY:") + print( + f"Configured {len(endpoints)} READY chat models for {args.app} " + "(UC EXECUTE + legacy CAN_QUERY):" + ) for endpoint in endpoints: print(f" {endpoint}") From de53c07b9b8fdb25547e214603ede9548cf8dac8 Mon Sep 17 00:00:00 2001 From: CoDA PR triage Date: Wed, 2 Sep 2026 01:12:53 +1000 Subject: [PATCH 22/23] docs(setup): surface required Gateway model grants --- README.md | 6 +- app.py | 162 +++++++++++++++++++++++++-------- docs/deployment.md | 2 +- static/index.html | 19 +++- tests/test_auth_enforcement.py | 19 ++++ 5 files changed, 167 insertions(+), 41 deletions(-) diff --git a/README.md b/README.md index 285b8876..4406f663 100644 --- a/README.md +++ b/README.md @@ -263,10 +263,10 @@ The managed endpoints require the configured Omnigent server SP and return `404` 1. Click [**Use this template**](https://github.com/datasciencemonkey/coding-agents-databricks-apps/generate) to create your own repo 2. Go to **Databricks → Apps → Create App** 3. Choose **Custom App** and connect your new repo -4. Deploy -5. Open the app — paste a short-lived PAT when prompted on first terminal session +4. From your checkout, run `make configure-gateway-resources PROFILE= APP_NAME=` +5. Deploy, then open the app and paste a short-lived PAT if prompted -That's it. No secrets to configure, no pre-deployment setup. +> **Required:** a plain Apps UI deployment cannot grant Unity AI Gateway v3 model-service privileges. The configuration target grants the app SP `USE CATALOG`, `USE SCHEMA`, and least-privilege `EXECUTE`; without it Pi and Claude report `404 '' does not exist`. The repository's normal `make deploy`, `make redeploy`, and `make deploy-git` workflows run it automatically. [→ Full deployment guide](docs/deployment.md) — environment variables, gateway config, and advanced options. diff --git a/app.py b/app.py index 0177eba9..da179052 100644 --- a/app.py +++ b/app.py @@ -229,7 +229,21 @@ def _update_step(step_id, **kwargs): def _get_setup_state_snapshot(): with setup_lock: - return copy.deepcopy(setup_state) + snapshot = copy.deepcopy(setup_state) + warnings = [] + agents_need_gateway = any( + os.environ.get(name, "true").strip().lower() not in ("false", "0", "no") + for name in ("ENABLE_CLAUDE", "ENABLE_PI") + ) + if agents_need_gateway and not os.environ.get("CODA_GATEWAY_MODEL_CATALOG", "").strip(): + warnings.append( + "AI Gateway model access is not provisioned for this app. Deploy with " + "the repository Make targets, or run `make configure-gateway-resources " + "PROFILE= APP_NAME=` before redeploying. App resources " + "alone do not grant Unity model-service EXECUTE." + ) + snapshot["warnings"] = warnings + return snapshot # Single-user security: only the token owner can access the terminal @@ -1953,26 +1967,62 @@ def omnigent_host_lease(): from omnigents_host import acquire_lease ok, lease = acquire_lease(owner, lease_id) - return jsonify(lease), (200 if ok else 409) + # The caller needs only the generation fence. Owner identity and lease + # timestamps remain process-internal and never cross the control API. + return jsonify({"lease_id": lease.get("lease_id"), "acquired": ok}), (200 if ok else 409) + + +def _repository_workspace_args(data): + """Return optional protocol-v2 repository metadata from a control request.""" + fields = ("repo_url", "repo_branch", "repo_name") + requested = any(data.get(field) is not None for field in fields) + if requested and data.get("workspace_protocol_version") != 2: + raise ValueError("repository workspace protocol version 2 is required") + return requested, {field: data.get(field) for field in fields} @app.route("/api/omnigent-host/workspaces", methods=["POST"]) def omnigent_host_workspace(): - """Allocate a distinct session directory under the active lease.""" + """Allocate, materialize, or release a distinct session workspace.""" if not _managed_omnigent_enabled(): return jsonify({"error": "Managed OmniGENT mode is disabled"}), 404 if not _omnigent_server_request_authorized(): return jsonify({"error": "Forbidden"}), 403 data = request.get_json(silent=True) or {} - from omnigents_host import allocate_workspace + from omnigents_host import ( + WorkspaceAllocationError, + allocate_workspace, + release_workspace, + ) + lease_id = str(data.get("lease_id") or "") + session_id = str(data.get("session_id") or "") try: + if data.get("action") == "release": + released = release_workspace(lease_id, session_id) + return jsonify( + { + "released": released, + "workspace_protocol_version": 2, + } + ) + repository_requested, repository = _repository_workspace_args(data) workspace = allocate_workspace( - str(data.get("lease_id") or ""), str(data.get("session_id") or "") + lease_id, + session_id, + **repository, ) + except WorkspaceAllocationError as exc: + return jsonify({"error": str(exc)}), exc.status_code except ValueError as exc: - return jsonify({"error": str(exc)}), 409 - return jsonify({"workspace": workspace}) + return jsonify({"error": str(exc)}), 400 + return jsonify( + { + "workspace": workspace, + "workspace_protocol_version": 2, + "repository_materialized": repository_requested, + } + ) @app.route("/api/omnigent-host/runner-log/") @@ -1980,7 +2030,7 @@ def omnigent_host_runner_log(session_id): """Return a bounded runner log tail to the configured server SP.""" if not _managed_omnigent_enabled(): return jsonify({"error": "Managed OmniGENT mode is disabled"}), 404 - if not _omnigent_server_request_authorized() and get_request_user() != app_owner: + if not _omnigent_server_request_authorized(): return jsonify({"error": "Forbidden"}), 403 from omnigents_host import runner_log_tail @@ -2005,54 +2055,87 @@ def omnigent_host_connect(): if configured_server_url and server_url.rstrip("/") != configured_server_url.rstrip("/"): return jsonify({"error": "server_url does not match configured Omnigent server"}), 409 - from omnigents_host import active_lease, connect_host + from omnigents_host import ( + WorkspaceAllocationError, + active_lease, + allocate_workspace, + connect_host, + release_workspace, + ) + lease_id = str(data.get("lease_id") or "") lease = active_lease() - if lease is None or lease.get("lease_id") != data.get("lease_id"): + if lease is None or lease.get("lease_id") != lease_id: return jsonify({"error": "stale or missing lease"}), 409 host_config = data.get("host_config") if host_config is not None and not isinstance(host_config, dict): return jsonify({"error": "host_config must be an object"}), 400 - ok, status = connect_host( - server_url, - _omnigent_sp_creds, - host_token=(data.get("host_token") or None), - host_id=(data.get("host_id") or None), - host_name=(data.get("host_name") or None), - host_config=host_config, - lease_id=(data.get("lease_id") or None), - ) + allocated_session_id = None + try: + repository_requested, repository = _repository_workspace_args(data) + session_id = str(data.get("session_id") or "") + if repository_requested and not session_id: + return jsonify({"error": "repository workspace protocol upgrade required"}), 426 + if session_id: + workspace = allocate_workspace( + lease_id, + session_id, + **repository, + ) + allocated_session_id = session_id + else: + # Older servers cannot identify an isolated opening workspace. + workspace = os.environ.get("HOME", "/app/python/source_code") + ok, status = connect_host( + server_url, + _omnigent_sp_creds, + host_token=(data.get("host_token") or None), + host_id=(data.get("host_id") or None), + host_name=(data.get("host_name") or None), + host_config=host_config, + lease_id=(data.get("lease_id") or None), + ) + except WorkspaceAllocationError as exc: + return jsonify({"error": str(exc)}), exc.status_code + except ValueError as exc: + return jsonify({"error": str(exc)}), 400 + except Exception: + if allocated_session_id is not None: + try: + release_workspace(lease_id, allocated_session_id) + except WorkspaceAllocationError: + pass + raise if not ok: + if allocated_session_id is not None: + try: + release_workspace(lease_id, allocated_session_id) + except WorkspaceAllocationError: + pass code = 409 if status.get("last_error") == "host already running" else 400 - return jsonify(status), code - status["workspace"] = os.environ.get("HOME", "/app/python/source_code") + return jsonify({"error": "host connection failed"}), code + status["workspace"] = workspace + status["workspace_protocol_version"] = 2 + status["repository_materialized"] = repository_requested return jsonify(status), 202 @app.route("/api/omnigent-host/disconnect", methods=["POST"]) def omnigent_host_disconnect(): - """Stop an external host or release the matching managed lease.""" + """Release and scrub only the matching managed lease generation.""" if not _managed_omnigent_enabled(): return jsonify({"error": "Managed OmniGENT mode is disabled"}), 404 if not _omnigent_server_request_authorized(): return jsonify({"error": "Forbidden"}), 403 data = request.get_json(silent=True) or {} lease_id = str(data.get("lease_id") or "") - from omnigents_host import ( - _scrub_session_workspaces, - active_lease, - disconnect_host, - release_lease, - ) + from omnigents_host import release_managed_lease - lease = active_lease() - if lease is None or lease.get("lease_id") != lease_id: - return jsonify({"released": False, "stale": True}) - status = disconnect_host() - if data.get("scrub"): - _scrub_session_workspaces() - release_lease(lease_id) - status["released"] = True + released, status = release_managed_lease(lease_id) + if status.get("stale"): + return jsonify(status) + if not released: + return jsonify(status), 503 return jsonify(status) @@ -2597,6 +2680,13 @@ def initialize_app(local_dev=False): """One-time init: detect owner, start cleanup thread.""" global app_owner, _omnigent_sp_creds, _sp_token_broker_server + if not os.environ.get("CODA_GATEWAY_MODEL_CATALOG", "").strip(): + logger.warning( + "AI Gateway model access is not provisioned: deploy with the repository " + "Make targets or run `make configure-gateway-resources PROFILE= " + "APP_NAME=`; Apps YAML alone cannot grant model-service EXECUTE" + ) + global _sp_token_broker_atexit_registered if not _sp_token_broker_atexit_registered: atexit.register(_shutdown_sp_token_broker) diff --git a/docs/deployment.md b/docs/deployment.md index 95d23d01..d86f9c5c 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -20,7 +20,7 @@ The app pulls the code directly from Git. To update later, just re-deploy — it > **Note:** On first startup, the app automatically removes the template's `.git` history and reinitializes a clean, remote-free git repo. This prevents accidental pushes back to the template repo from the in-browser terminal. -> **AI Gateway setup:** CLI deployment targets automatically discover READY Foundation Model chat endpoints and attach them to the app with `CAN_QUERY`. This is required because Gateway visibility and invocation are identity-scoped; a model visible to the deployer may otherwise return a masked 404 to the app service principal. Set `AUTO_CONFIGURE_GATEWAY=false` only when permissions are managed externally. +> **Required AI Gateway setup:** deploy through this repository's `make deploy`, `make redeploy`, or `make deploy-git` targets. They grant the app SP `USE CATALOG` on `system`, `USE SCHEMA` on `system.ai`, and `EXECUTE` on each selected Unity AI Gateway v3 model service; they also attach legacy `CAN_QUERY` resources and publish the app-visible model inventory. A plain Apps UI deploy cannot express the v3 UC grants and leaves Pi/Claude returning a masked `404 '' does not exist`. Set `AUTO_CONFIGURE_GATEWAY=false` only when equivalent UC grants are managed externally. ## Alternative: Deploy with CLI diff --git a/static/index.html b/static/index.html index bd4f9fa8..559f2947 100644 --- a/static/index.html +++ b/static/index.html @@ -1643,6 +1643,21 @@

General

stopSpinner(); term.write('\r\x1b[K' + color + ' ' + emoji + '\x1b[0m ' + msg + '\r\n'); }; + const showSetupWarnings = (data) => { + const clean = (value) => String(value || '').replace(/[\x00-\x1f\x7f-\x9f]/g, ' '); + const messages = [...(data.warnings || [])]; + for (const step of (data.steps || [])) { + if (step.status === 'error' && step.error) { + messages.push(`${step.label}: ${step.error}`); + } + } + if (!messages.length) return; + term.write('\r\n\x1b[1;33m AI agent setup requires attention:\x1b[0m\r\n'); + for (const message of messages) { + term.write(`\x1b[33m • ${clean(message)}\x1b[0m\r\n`); + } + term.write('\r\n'); + }; // Polls /api/setup-status and rewrites the spinner line with the // currently-running step + completion count. Returns when setup is // complete or errored. Used by both bootstrap paths. @@ -1659,7 +1674,7 @@

General

return; } else if (pollData.status === 'error') { finishLine('!', '\x1b[1;33m', 'Setup completed with warnings'); - term.write('\r\n'); + showSetupWarnings(pollData); return; } const steps = pollData.steps || []; @@ -1773,6 +1788,8 @@

General

const setupData2 = await setupResp2.json(); if (setupData2.status !== 'complete' && setupData2.status !== 'error') { await waitForSetup(); + } else if (setupData2.status === 'error' || (setupData2.warnings || []).length) { + showSetupWarnings(setupData2); } var { sid, reattached } = await getOrPromptSession(term, tab.label, opts.skipPrompt); } else { diff --git a/tests/test_auth_enforcement.py b/tests/test_auth_enforcement.py index 387e621f..7b732479 100644 --- a/tests/test_auth_enforcement.py +++ b/tests/test_auth_enforcement.py @@ -69,6 +69,25 @@ def test_connect_endpoint_requires_allowlisted_server_sp(monkeypatch): assert app_module._omnigent_server_request_authorized() is False +def test_setup_snapshot_warns_when_gateway_model_grants_are_missing(monkeypatch): + app_module = _get_app_module() + monkeypatch.delenv("CODA_GATEWAY_MODEL_CATALOG", raising=False) + monkeypatch.setenv("ENABLE_CLAUDE", "true") + + snapshot = app_module._get_setup_state_snapshot() + + assert len(snapshot["warnings"]) == 1 + assert "model-service EXECUTE" in snapshot["warnings"][0] + assert "configure-gateway-resources" in snapshot["warnings"][0] + + +def test_setup_snapshot_has_no_gateway_warning_when_catalog_is_attached(monkeypatch): + app_module = _get_app_module() + monkeypatch.setenv("CODA_GATEWAY_MODEL_CATALOG", '["system.ai.claude-sonnet-5"]') + + assert app_module._get_setup_state_snapshot()["warnings"] == [] + + # 1. Session endpoints MUST enforce owner check # --------------------------------------------------------------------------- From 5171f13256ea38ed459ce92952a05568dba4cace Mon Sep 17 00:00:00 2001 From: CoDA PR triage Date: Wed, 2 Sep 2026 01:13:52 +1000 Subject: [PATCH 23/23] fix(omnigent): remove unrelated workspace changes from setup warning --- app.py | 139 +++++++++++++++------------------------------------------ 1 file changed, 35 insertions(+), 104 deletions(-) diff --git a/app.py b/app.py index da179052..1521d0d5 100644 --- a/app.py +++ b/app.py @@ -1967,62 +1967,26 @@ def omnigent_host_lease(): from omnigents_host import acquire_lease ok, lease = acquire_lease(owner, lease_id) - # The caller needs only the generation fence. Owner identity and lease - # timestamps remain process-internal and never cross the control API. - return jsonify({"lease_id": lease.get("lease_id"), "acquired": ok}), (200 if ok else 409) - - -def _repository_workspace_args(data): - """Return optional protocol-v2 repository metadata from a control request.""" - fields = ("repo_url", "repo_branch", "repo_name") - requested = any(data.get(field) is not None for field in fields) - if requested and data.get("workspace_protocol_version") != 2: - raise ValueError("repository workspace protocol version 2 is required") - return requested, {field: data.get(field) for field in fields} + return jsonify(lease), (200 if ok else 409) @app.route("/api/omnigent-host/workspaces", methods=["POST"]) def omnigent_host_workspace(): - """Allocate, materialize, or release a distinct session workspace.""" + """Allocate a distinct session directory under the active lease.""" if not _managed_omnigent_enabled(): return jsonify({"error": "Managed OmniGENT mode is disabled"}), 404 if not _omnigent_server_request_authorized(): return jsonify({"error": "Forbidden"}), 403 data = request.get_json(silent=True) or {} - from omnigents_host import ( - WorkspaceAllocationError, - allocate_workspace, - release_workspace, - ) + from omnigents_host import allocate_workspace - lease_id = str(data.get("lease_id") or "") - session_id = str(data.get("session_id") or "") try: - if data.get("action") == "release": - released = release_workspace(lease_id, session_id) - return jsonify( - { - "released": released, - "workspace_protocol_version": 2, - } - ) - repository_requested, repository = _repository_workspace_args(data) workspace = allocate_workspace( - lease_id, - session_id, - **repository, + str(data.get("lease_id") or ""), str(data.get("session_id") or "") ) - except WorkspaceAllocationError as exc: - return jsonify({"error": str(exc)}), exc.status_code except ValueError as exc: - return jsonify({"error": str(exc)}), 400 - return jsonify( - { - "workspace": workspace, - "workspace_protocol_version": 2, - "repository_materialized": repository_requested, - } - ) + return jsonify({"error": str(exc)}), 409 + return jsonify({"workspace": workspace}) @app.route("/api/omnigent-host/runner-log/") @@ -2030,7 +1994,7 @@ def omnigent_host_runner_log(session_id): """Return a bounded runner log tail to the configured server SP.""" if not _managed_omnigent_enabled(): return jsonify({"error": "Managed OmniGENT mode is disabled"}), 404 - if not _omnigent_server_request_authorized(): + if not _omnigent_server_request_authorized() and get_request_user() != app_owner: return jsonify({"error": "Forbidden"}), 403 from omnigents_host import runner_log_tail @@ -2055,87 +2019,54 @@ def omnigent_host_connect(): if configured_server_url and server_url.rstrip("/") != configured_server_url.rstrip("/"): return jsonify({"error": "server_url does not match configured Omnigent server"}), 409 - from omnigents_host import ( - WorkspaceAllocationError, - active_lease, - allocate_workspace, - connect_host, - release_workspace, - ) + from omnigents_host import active_lease, connect_host - lease_id = str(data.get("lease_id") or "") lease = active_lease() - if lease is None or lease.get("lease_id") != lease_id: + if lease is None or lease.get("lease_id") != data.get("lease_id"): return jsonify({"error": "stale or missing lease"}), 409 host_config = data.get("host_config") if host_config is not None and not isinstance(host_config, dict): return jsonify({"error": "host_config must be an object"}), 400 - allocated_session_id = None - try: - repository_requested, repository = _repository_workspace_args(data) - session_id = str(data.get("session_id") or "") - if repository_requested and not session_id: - return jsonify({"error": "repository workspace protocol upgrade required"}), 426 - if session_id: - workspace = allocate_workspace( - lease_id, - session_id, - **repository, - ) - allocated_session_id = session_id - else: - # Older servers cannot identify an isolated opening workspace. - workspace = os.environ.get("HOME", "/app/python/source_code") - ok, status = connect_host( - server_url, - _omnigent_sp_creds, - host_token=(data.get("host_token") or None), - host_id=(data.get("host_id") or None), - host_name=(data.get("host_name") or None), - host_config=host_config, - lease_id=(data.get("lease_id") or None), - ) - except WorkspaceAllocationError as exc: - return jsonify({"error": str(exc)}), exc.status_code - except ValueError as exc: - return jsonify({"error": str(exc)}), 400 - except Exception: - if allocated_session_id is not None: - try: - release_workspace(lease_id, allocated_session_id) - except WorkspaceAllocationError: - pass - raise + ok, status = connect_host( + server_url, + _omnigent_sp_creds, + host_token=(data.get("host_token") or None), + host_id=(data.get("host_id") or None), + host_name=(data.get("host_name") or None), + host_config=host_config, + lease_id=(data.get("lease_id") or None), + ) if not ok: - if allocated_session_id is not None: - try: - release_workspace(lease_id, allocated_session_id) - except WorkspaceAllocationError: - pass code = 409 if status.get("last_error") == "host already running" else 400 - return jsonify({"error": "host connection failed"}), code - status["workspace"] = workspace - status["workspace_protocol_version"] = 2 - status["repository_materialized"] = repository_requested + return jsonify(status), code + status["workspace"] = os.environ.get("HOME", "/app/python/source_code") return jsonify(status), 202 @app.route("/api/omnigent-host/disconnect", methods=["POST"]) def omnigent_host_disconnect(): - """Release and scrub only the matching managed lease generation.""" + """Stop an external host or release the matching managed lease.""" if not _managed_omnigent_enabled(): return jsonify({"error": "Managed OmniGENT mode is disabled"}), 404 if not _omnigent_server_request_authorized(): return jsonify({"error": "Forbidden"}), 403 data = request.get_json(silent=True) or {} lease_id = str(data.get("lease_id") or "") - from omnigents_host import release_managed_lease + from omnigents_host import ( + _scrub_session_workspaces, + active_lease, + disconnect_host, + release_lease, + ) - released, status = release_managed_lease(lease_id) - if status.get("stale"): - return jsonify(status) - if not released: - return jsonify(status), 503 + lease = active_lease() + if lease is None or lease.get("lease_id") != lease_id: + return jsonify({"released": False, "stale": True}) + status = disconnect_host() + if data.get("scrub"): + _scrub_session_workspaces() + release_lease(lease_id) + status["released"] = True return jsonify(status)