From cf50bd5bf656293ce270351e56e9cacf59fb03e5 Mon Sep 17 00:00:00 2001 From: CoDA PR triage Date: Fri, 7 Aug 2026 13:34:14 +1000 Subject: [PATCH 01/11] 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 c053300a..c8cef9d1 100644 --- a/app.py +++ b/app.py @@ -1714,16 +1714,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 fe0f5a44..372050c6 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 8386e76293db9d32710896d49a8eec2e4ff9bd88 Mon Sep 17 00:00:00 2001 From: CoDA PR triage Date: Fri, 7 Aug 2026 14:37:08 +1000 Subject: [PATCH 02/11] 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 c8cef9d1..0d0581b3 100644 --- a/app.py +++ b/app.py @@ -1535,7 +1535,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() @@ -1743,6 +1743,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.""" @@ -1753,8 +1786,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 @@ -1770,14 +1806,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"]) @@ -2285,8 +2338,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 fee11e50c2afac42f3e51a664ca006ea744d912c Mon Sep 17 00:00:00 2001 From: CoDA PR triage Date: Fri, 7 Aug 2026 15:11:35 +1000 Subject: [PATCH 03/11] 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 372050c6..4c47470a 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 53a142564a666dfbf569127fd06bef33d9561d86 Mon Sep 17 00:00:00 2001 From: CoDA PR triage Date: Fri, 7 Aug 2026 15:33:35 +1000 Subject: [PATCH 04/11] 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 0d0581b3..b5d5b0d8 100644 --- a/app.py +++ b/app.py @@ -1709,7 +1709,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 ab78387cdf476a3167d5237125e2053182ab9633 Mon Sep 17 00:00:00 2001 From: CoDA PR triage Date: Fri, 7 Aug 2026 16:40:06 +1000 Subject: [PATCH 05/11] 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 b5d5b0d8..276f4fc0 100644 --- a/app.py +++ b/app.py @@ -1778,6 +1778,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 644ebecc94dabc98a8e4abf549ff0b5a257f64e3 Mon Sep 17 00:00:00 2001 From: CoDA PR triage Date: Fri, 7 Aug 2026 16:50:41 +1000 Subject: [PATCH 06/11] 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 05131c2eaadecf2181a0ed07c3f1c5d9bfd03f9e Mon Sep 17 00:00:00 2001 From: CoDA PR triage Date: Sat, 8 Aug 2026 01:51:31 +1000 Subject: [PATCH 07/11] fix(omnigent): keep active runner leases alive Signed-off-by: CoDA PR triage --- pyproject.toml | 1 + requirements.lock | 29 +++++++++++++++++++++++++++-- requirements.txt | 4 ++++ tests/test_omnigents_host.py | 36 ++++++++++++++++++++++++++++++++++++ 4 files changed, 68 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index d87ed654..d5fb4da0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,6 +11,7 @@ dependencies = [ "databricks-sdk>=0.106.0", "mlflow-skinny==3.14.0", "requests", + "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 dbffe07a..ee039eb2 100644 --- a/requirements.lock +++ b/requirements.lock @@ -1,5 +1,5 @@ # This file was autogenerated by uv via the following command: -# uv pip compile requirements.txt -o /tmp/lock-check.txt --generate-hashes +# uv pip compile requirements.txt -o requirements.lock --generate-hashes annotated-doc==0.0.4 \ --hash=sha256:571ac1dc6991c450b25a9c2d84a3705e2ae7a53467b5d111c24fa8baabbed320 \ --hash=sha256:fbcda96e87e9c92ad167c2e53839e57503ecfda18804ea28102353485033faa4 @@ -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 @@ -988,7 +1011,9 @@ smmap==5.0.3 \ sniffio==1.3.1 \ --hash=sha256:2f6da418d1f1e0fddd844478f41680e794e6051915791a034ff65e5f100525a2 \ --hash=sha256:f4324edc670a0f49750a81b895f35c3adb843cca46f0530f79fc1babb23789dc - # via claude-agent-sdk + # via + # -r requirements.txt + # claude-agent-sdk sqlparse==0.5.5 \ --hash=sha256:12a08b3bf3eec877c519589833aed092e2444e68240a3577e8e26148acc7b1ba \ --hash=sha256:e20d4a9b0b8585fdf63b10d30066c7c94c5d7a7ec47c889a2d83a3caa93ff28e diff --git a/requirements.txt b/requirements.txt index 77f32c44..37fc948b 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) @@ -190,6 +192,8 @@ simple-websocket==1.1.0 # python-engineio smmap==5.0.3 # via gitdb +sniffio==1.3.1 + # via claude-agent-sdk sqlparse==0.5.5 # via mlflow-skinny sse-starlette==3.3.4 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 f155aadd13b5649898aecd275f4a487ffa9153ea Mon Sep 17 00:00:00 2001 From: CoDA PR triage Date: Sat, 8 Aug 2026 03:05:31 +1000 Subject: [PATCH 08/11] 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 0c1ece575f430396839b27f92ba05a3f7b6becc2 Mon Sep 17 00:00:00 2001 From: CoDA PR triage Date: Sat, 8 Aug 2026 04:12:33 +1000 Subject: [PATCH 09/11] 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 4c47470a..821cce37 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 b12b59462b6c6e9f0d0b8d51877d21ef70d0372e Mon Sep 17 00:00:00 2001 From: CoDA PR triage Date: Sat, 8 Aug 2026 06:22:26 +1000 Subject: [PATCH 10/11] 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 b30dd5bd591498811bdc3d0c4711df7a5d1e2582 Mon Sep 17 00:00:00 2001 From: CoDA PR triage Date: Sat, 8 Aug 2026 06:36:08 +1000 Subject: [PATCH 11/11] 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 276f4fc0..4995f500 100644 --- a/app.py +++ b/app.py @@ -1725,10 +1725,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 @@ -1753,8 +1753,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()