chore: import upstream snapshot with attribution

This commit is contained in:
wehub-resource-sync
2026-07-13 13:12:00 +08:00
commit 3de48288cb
2986 changed files with 1131193 additions and 0 deletions
+1
View File
@@ -0,0 +1 @@
"""Project development tooling (lint hooks run from .pre-commit-config.yaml)."""
+1
View File
@@ -0,0 +1 @@
"""Performance benchmarks (runnable via ``uv run``, not shipped)."""
+250
View File
@@ -0,0 +1,250 @@
# Omnigent performance benchmark
Baseline, repeatable latency/throughput numbers for key Omnigent user
journeys, so we can track them over time and catch regressions. Modeled on
MLflow's `dev/benchmarks/gateway/` workflow.
The harness boots a real `omnigent server`, drives the selected journeys under
load, prints latency/throughput tables, and writes a versioned JSON report.
Two families: **HTTP/API journeys** (server + DB, no runner/LLM — fast and
low-noise) and **full-turn journeys** (a real agent turn through the runner +
a zero-latency mock LLM). See *Journeys* below.
By default the server boots a fresh, empty SQLite DB, which gives best-case
numbers that don't move with load. For meaningful results, point it at a
**pre-seeded corpus** (`seed.py`) and, ideally, at **Postgres** — production
runs on Databricks Lakebase (Postgres), whose per-query round-trip + pooling
cost SQLite doesn't have. See *Seeding* and *Backends* below.
## Run it
```bash
# All journeys, sequential latency (100 iterations × 3 runs each).
uv run --no-sync dev/benchmarks/omnigent/run.py
# A subset, writing a report for CI artifact upload.
uv run --no-sync dev/benchmarks/omnigent/run.py \
--journeys list_sessions,load_conversation_history \
--iterations 200 --runs 3 --output bench.json
# Throughput mode: >1 concurrency drives concurrency-safe journeys as load.
uv run --no-sync dev/benchmarks/omnigent/run.py \
--requests 500 --concurrency 25 --runs 3
# CI gating: exit 1 if a threshold is breached.
uv run --no-sync dev/benchmarks/omnigent/run.py --max-p50-ms 25 --max-p99-ms 100
```
`--no-sync` runs against the already-installed venv. (A bare `uv run` may try to
rebuild the project, which fails in a git worktree without a Node web-UI build;
`OMNIGENT_SKIP_WEB_UI=true uv sync` prepares the venv once, then use
`--no-sync`.)
Key flags (`--help` for all): `--journeys A,B`, `--database-uri URI` (seeded
corpus / Postgres; default: throwaway empty SQLite), `--iterations N` (per
latency run), `--requests N` / `--concurrency N` (throughput), `--runs N`,
`--warmup N`, `--output FILE`, `--min-rps` / `--max-p50-ms` / `--max-p99-ms`
(CI thresholds).
## Journeys
### HTTP/API (server + DB, runner-free)
| Journey | Operation timed | Stressed by |
| --- | --- | --- |
| `list_sessions` | `GET /v1/sessions` — session-list read | session count |
| `create_session` | `POST /v1/sessions` then `DELETE` — session create | write path |
| `get_session` | `GET /v1/sessions/{id}` — single-session snapshot | (O(1)) |
| `load_conversation_history` | `GET /v1/sessions/{id}/items` — history read | items/session |
| `search_sessions` | `GET /v1/sessions?search_query=` — unindexed `LIKE` | total item count |
| `fork_session` | `POST /v1/sessions/{id}/fork` — fork (deep-copy items); forks deleted in teardown, untimed | items/session |
| `add_comment` | `POST /v1/sessions/{id}/comments` — create a review comment | write path |
Read journeys target a **pre-seeded** session when the DB has a corpus; against
an empty DB they self-seed a small fallback session over HTTP (the
`external_conversation_item` event — appends items without starting a task), so
they still work with no runner or LLM.
### Full-turn (runner + mock LLM)
These drive a real agent turn end-to-end — `POST …/events` → server → **runner**
→ in-process executor → mock LLM → stream back → `idle`. Selecting any of them
boots `BenchEnvironment(with_runner=True)` automatically.
Each turn costs ~1 s+ (vs. the millisecond HTTP journeys), so these journeys
cap their latency iterations (`Journey.max_iterations`, currently 5) — a large
`--iterations` tuned for the HTTP journeys is clamped down for them so the run
stays within the CI time budget, with `--runs` providing the repeats. The cap
only lowers the count, never raises it. A cold start never deletes its session,
so sessions accumulate across a run; keeping the count small also keeps that
drift negligible (~2 ms/turn).
| Journey | Operation timed |
| --- | --- |
| `session_cold_start` | Create+bind a fresh session and drive its first turn to `idle` (runner spawn + executor construction + turn) |
| `warm_turn` | Drive a turn on an already-warm session — steady-state dispatch overhead |
| `time_to_first_token` | Post a turn; time to the first streamed `output_text` delta |
| `interrupt` | Interrupt a running (gated) turn; time to cancellation |
| `read_runner_file` | `GET .../environments/default/filesystem/{path}` — server → runner filesystem read proxy |
`read_runner_file` needs a runner but does **not** drive a turn or call the LLM:
its setup plants a file via `PUT`, and the timed op is the proxied read (a
localhost round-trip). Being far cheaper than a turn, it uses a higher iteration
cap (50) than the full-turn journeys.
**Only measure what we control.** Full-turn journeys always use the
**`openai-agents`** SDK harness, which runs **in-process** (a call into the
`agents` library + an HTTP call to the mock LLM) — no vendor binary, no external
process. Native harnesses (e.g. `claude-native`) launch the real vendor CLI
into a tmux pane, whose startup we don't control, so they're deliberately
excluded. The mock LLM is zero-latency, so every number is omnigent
dispatch/streaming/cancel overhead, not model latency.
Add a journey by registering a `Journey` in `journeys.py` (set `needs_runner`
for full-turn journeys).
## Seeding a realistic corpus
`seed.py` writes a sizeable, deterministic corpus directly through the store
API (no HTTP, no runner) into the same DB the server then boots against:
```bash
# Seed 5000 sessions × 50 items into a SQLite file, then benchmark against it.
uv run --no-sync dev/benchmarks/omnigent/seed.py \
--database-uri sqlite:////abs/path/bench.db --sessions 5000 --items-per-session 50
uv run --no-sync dev/benchmarks/omnigent/run.py \
--database-uri sqlite:////abs/path/bench.db --output bench.json
```
Seeding is **idempotent**: a matching corpus (same sessions/items/schema) is
detected and reused, so re-running is a fast no-op — pass `--reseed` to force,
or a differing config to be warned. SQLite absolute paths need four slashes
(`sqlite:////abs/...`). The reuse marker records the DB's Alembic head read at
seed time, so a corpus from an older schema is automatically reseeded — no
manual revision bookkeeping. `test_seed_creates_listable_corpus` (which seeds
through the store, running migrations to the current head) is the safety net
that a schema change hasn't broken seeding.
## Backends
`--database-uri` selects the DB; the report's `backend` field (`sqlite` /
`postgres` / `mysql`) is derived from the URI scheme so results group by
backend.
- **SQLite** (default) — in-process; fast, but not prod-representative.
- **Postgres** — `postgresql+psycopg://user@host:5432/db` (the fully-qualified
`+psycopg` form; the server CLI does not normalize a bare `postgresql://`).
Requires `psycopg[binary]` (the `databricks` extra). Matches prod's
round-trip/pooling profile. Stand up a local one with
`docker run -e POSTGRES_PASSWORD=… -p 5432:5432 postgres:16`.
- **MySQL** — `mysql+mysqldb://user@host:3306/db`. Requires the `mysqlclient`
driver (`pip install mysqlclient`, which needs the `libmysqlclient-dev`
system library) — it is not in any extra. A supported backend, though prod
runs on Postgres. Stand up a local one with
`docker run -e MYSQL_ROOT_PASSWORD=… -e MYSQL_DATABASE=benchdb -p 3306:3306 mysql:8.0`.
## Output → Databricks → dashboard
The harness writes JSON only. Storage and charting live in Databricks:
```
run.py --output bench.json → GitHub Actions artifact → Databricks notebook (ETL) → Delta table → AI/BI dashboard
(this repo) (CI, follow-up) (workspace, yours)
```
The repo's contract is the **JSON schema** below. A workspace notebook (owned
outside this repo, modeled on MLflow's gateway ETL) pulls the CI artifacts via
the GitHub API, flattens each run's `summary` + `runs` + metadata, and
`saveAsTable`s into a Delta table the dashboard reads. `sample_output.json` is a
committed, faithful example so the notebook can be written against a real
document without running the harness.
### JSON schema (`schema.py`, `SCHEMA_VERSION`)
```jsonc
{
"schema_version": 1,
"generated_at": "<ISO-8601 UTC>",
"git_sha": "<HEAD sha>",
"git_branch": "<branch>",
"host": {"platform": "...", "python": "...", "cpu_count": 12},
"harness": "http-only",
"config": {"iterations": 100, "requests": 500, "concurrency": 1,
"runs": 3, "warmup": 10, "with_runner": false,
"backend": "sqlite"},
"journeys": {
"<journey name>": {
"kind": "latency" | "throughput",
"backend": "sqlite" | "postgres" | "mysql",
"runs": [ // one per --runs
{"n_success": N, "n_failures": N, "failures": {"HTTP 500": 1},
"wall_time_s": , "mean_ms": , "p50_ms": , "p95_ms": ,
"p99_ms": , "max_ms": , "rps": }
],
"summary": {"avg_mean_ms": , "avg_p50_ms": , "avg_p95_ms": ,
"avg_p99_ms": , "avg_rps": } // averaged across runs
}
}
}
```
The per-journey `summary` + `runs` shape mirrors MLflow's gateway benchmark, so
the same ETL flatten works — keyed by `journey` and `backend`. Bump
`SCHEMA_VERSION` on any breaking shape change so the notebook can branch on it.
## Layout
| File | Role |
| --- | --- |
| `run.py` | CLI orchestrator + entrypoint |
| `seed.py` | deterministic corpus seeder (store API) |
| `journeys.py` | `Journey` dataclass, latency/throughput runners, registry |
| `environment.py` | server (± runner + mock LLM) lifecycle; `--database-uri` |
| `measure.py` | `RunResult`, percentile, aggregation, thresholds, tables |
| `schema.py` | `SCHEMA_VERSION`, `build_report`, git/host metadata |
| `sample_output.json` | committed example of the JSON contract |
The smoke test is `tests/benchmarks/test_benchmark_smoke.py` (boots the server
with tiny counts + a seeded-corpus unit test; runs on the normal CI lane, no
creds).
## CI
`.github/workflows/benchmark.yml` runs nightly (and on dispatch) as a backend
matrix — `sqlite`, `postgres` (a `postgres:16` service container), and `mysql`
(a `mysql:8.0` service container; the `mysqlclient` driver is installed on that
leg only). Each leg seeds a corpus (SQLite reuses a cache keyed on the schema
head + `seed.py` + corpus config, so a migration busts the cache and forces a
reseed; Postgres and MySQL are fresh per run), runs the benchmark, and uploads
`benchmark-results-<backend>-<run_id>.json`. The workspace notebook pulls those
artifacts.
Schema changes need no manual step: the seed always targets the current
migrated schema (migrations run when the store is constructed), the reuse
marker records the head read at seed time (so old corpora auto-reseed), and
`test_seed_creates_listable_corpus` fails if a migration genuinely breaks
seeding.
## Follow-ups
- **Subagent spawn.** A planned full-turn journey (`needs_runner=True`): the
parent agent emits a `sys_session_send` tool call, the runner dispatches a
child session, and the parent auto-wakes with the collected result. It's
fully mockable with the zero-latency mock LLM (no real model) — script the
parent's queue to emit the tool call and the child's queue to return a short
reply, then poll for the child's marker. It needs the parent bundle to declare
a sub-agent under `tools:` (extend `_agent_bundle`); the pattern is in
`tests/e2e/test_coder_subagent.py`.
- **Excluded journeys** (agent-behaviour-dependent, deliberately not measured):
multi-turn and tool-calling turns (dominated by the agent's own choices) and
large-history turns (the O(N) `history_to_input_items` conversion is real app
work but only fires on a cold runner cache, so isolating it entangles with
cold-start cost).
- **CI matrix.** Runner journeys are backend-agnostic (they exercise runner
dispatch, not big DB reads), so the nightly workflow can run them on the
SQLite leg only rather than both — wire a runner `--journeys` set into
`benchmark.yml` when desired.
- **Simulated provider latency.** The mock LLM returns at ~zero latency, which
is what isolates omnigent overhead. A fixed per-response delay knob would let
turns model end-user wall-clock instead; it's a small change behind the
`configure_mock` / `set_mock_fallback` seam if that's ever wanted.
+7
View File
@@ -0,0 +1,7 @@
"""Omnigent user-journey performance benchmark.
Stands up a real server + runner against a zero-latency mock LLM, drives
key user journeys under load, and emits a versioned JSON report of latency
percentiles and throughput. See ``README.md`` for the workflow and how the
workspace ETL notebook consumes the JSON.
"""
+694
View File
@@ -0,0 +1,694 @@
"""Benchmark environment lifecycle.
:class:`BenchEnvironment` is an async context manager that stands up a real
Omnigent ``server`` with no Databricks credentials. Two modes:
- ``with_runner=False`` (default): server + SQLite DB only. Enough for the
HTTP/API journeys, which never drive an agent turn.
- ``with_runner=True``: additionally spawns a zero-latency mock LLM and a
sibling ``runner``, routes the server-side prompt-policy classifier at the
mock (via ``--config``), and sets an ALLOW fallback — everything the
full-turn journeys need.
A full env is a strict superset of the HTTP-only env, so both modes share one
class; the runner mode is gated behind the flag rather than forked into a
separate type. It mirrors the proven ``live_server`` e2e recipe
(``tests/e2e/conftest.py``) and reuses the credential-free spawn core: the
compat helpers (so subprocesses import this worktree) and
``token_bound_runner_id``.
"""
from __future__ import annotations
import asyncio
import contextlib
import io
import os
import signal
import socket
import subprocess
import sys
import tarfile
import time
import uuid
from pathlib import Path
from typing import IO
import httpx
import yaml
from omnigent.runner.identity import OMNIGENT_INTERNAL_WS_ORIGIN, token_bound_runner_id
from tests._helpers.compat import (
apply_runner_env,
apply_server_env,
compat_runner_cwd,
compat_server_cwd,
runner_executable,
server_executable,
)
_REPO_ROOT = Path(__file__).resolve().parents[3]
_MOCK_SERVER = _REPO_ROOT / "tests" / "server" / "integration" / "mock_llm_server.py"
_HEALTH_TIMEOUT_S = 90.0
_MOCK_TIMEOUT_S = 15.0
_POLL_INTERVAL_S = 0.2
_TURN_TIMEOUT_S = 180.0
# Terminal SSE events — if one arrives before any delta, the turn produced no
# streamed text (a failure for the TTFT journey).
_STREAM_TERMINAL_EVENTS = frozenset(
{"response.completed", "response.failed", "response.cancelled"}
)
# The server persists an interrupted turn as a synthetic user message whose
# text contains this marker (see tests/e2e/test_cancel_history.py).
_CANCELLATION_MARKER = "interrupted"
# Default full-turn agent (with_runner=True). The mock ignores the model for
# routing (its "default" queue serves any request), but the key is baked into
# the spec so the harness has a concrete model to send.
_DEFAULT_MODEL = "mock-bench-brain"
_DEFAULT_HARNESS = "openai-agents"
# Server-side prompt-policy classifier queue key. In runner mode we set an
# ALLOW fallback here so a classifier call (if the agent trips one) never
# blocks or returns non-verdict text.
_POLICY_LLM_KEY = "_policy_llm_"
_POLICY_ALLOW = '{"action": "allow", "reason": ""}'
def _find_free_port() -> int:
"""Bind an ephemeral port and return it (races are tolerated by retries)."""
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
sock.bind(("127.0.0.1", 0))
return int(sock.getsockname()[1])
class BenchEnvironment:
"""Async context manager owning the benchmark's server (± runner + mock).
:param with_runner: When ``False`` (default), boot the server only — the
v1 HTTP-journey path. When ``True``, also spawn the mock LLM and a
runner and wire the policy classifier at the mock — the phase-2
full-turn path.
:param database_uri: SQLAlchemy URI the server boots against. ``None``
(default) uses a fresh throwaway SQLite file in the temp dir — the
empty-DB path. Pass a pre-seeded URI (e.g. a seeded SQLite file, or a
``postgresql+psycopg://…`` instance) to benchmark against a realistic
corpus. Postgres must be the fully-qualified ``+psycopg`` form — the
server CLI does not normalize it.
:param harness: Harness for full-turn agents when ``with_runner`` (default
``openai-agents``, a base dependency needing no vendor CLI binary).
:param model: Model string baked into registered agent specs.
"""
def __init__(
self,
*,
with_runner: bool = False,
database_uri: str | None = None,
harness: str = _DEFAULT_HARNESS,
model: str = _DEFAULT_MODEL,
) -> None:
self.with_runner = with_runner
self.database_uri = database_uri
self.harness = harness
self.model = model
self.base_url = ""
self.mock_url = ""
self.runner_id = ""
self.client: httpx.AsyncClient | None = None
self._tmp = Path("/tmp") / f"omni-bench-{uuid.uuid4().hex[:8]}"
self._mock_proc: subprocess.Popen[bytes] | None = None
self._server_proc: subprocess.Popen[bytes] | None = None
self._runner_proc: subprocess.Popen[bytes] | None = None
self._log_handles: list[IO[bytes]] = []
self._agent_cache: dict[str, str] = {}
# ── lifecycle ────────────────────────────────────────────
async def __aenter__(self) -> BenchEnvironment:
await asyncio.to_thread(self._start)
self.client = httpx.AsyncClient(
base_url=self.base_url,
timeout=300.0,
headers={"Origin": OMNIGENT_INTERNAL_WS_ORIGIN},
)
if self.with_runner:
# ALLOW fallback so a server-side classifier call resolves against
# the mock (never api.openai.com) and returns a valid verdict.
await self._mock_post(
"/mock/set_fallback", {"key": _POLICY_LLM_KEY, "text": _POLICY_ALLOW}
)
return self
async def __aexit__(self, *exc: object) -> None:
if self.client is not None:
await self.client.aclose()
await asyncio.to_thread(self._stop)
def _start(self) -> None:
"""Spawn the server (± mock + runner) and block until ready."""
self._tmp.mkdir(mode=0o700, parents=True, exist_ok=True)
artifact_dir = self._tmp / "artifacts"
artifact_dir.mkdir(exist_ok=True)
if self.with_runner:
mock_port = _find_free_port()
self.mock_url = f"http://127.0.0.1:{mock_port}"
self._mock_proc = self._spawn_mock(mock_port)
self._wait_mock_ready()
port = _find_free_port()
self.base_url = f"http://localhost:{port}"
binding_token = uuid.uuid4().hex
base_env = {**os.environ}
if self.with_runner:
self.runner_id = token_bound_runner_id(binding_token)
base_env["OPENAI_API_KEY"] = "mock-key"
# The OpenAI SDK appends /responses, so include /v1 in the base.
base_env["OPENAI_BASE_URL"] = f"{self.mock_url}/v1"
# Prepend the worktree so subprocesses import this branch's source.
apply_server_env(base_env, _REPO_ROOT)
self._server_proc = self._spawn_server(port, base_env, binding_token, artifact_dir)
if self.with_runner:
self._runner_proc = self._spawn_runner(base_env, binding_token)
self._wait_ready()
def _stop(self) -> None:
"""Terminate runner, server, and mock; remove the temp dir."""
for proc in (self._runner_proc, self._server_proc, self._mock_proc):
if proc is not None and proc.poll() is None:
proc.send_signal(signal.SIGTERM)
try:
proc.wait(timeout=8)
except subprocess.TimeoutExpired:
proc.kill()
proc.wait(timeout=5)
for handle in self._log_handles:
handle.close()
import shutil
shutil.rmtree(self._tmp, ignore_errors=True)
# ── spawns ───────────────────────────────────────────────
def _log(self, name: str) -> IO[bytes]:
handle = (self._tmp / name).open("wb")
self._log_handles.append(handle)
return handle
def _spawn_mock(self, port: int) -> subprocess.Popen[bytes]:
return subprocess.Popen(
[sys.executable, str(_MOCK_SERVER), str(port)],
env={**os.environ, "PYTHONPATH": str(_REPO_ROOT)},
stdout=self._log("mock.log"),
stderr=subprocess.STDOUT,
)
def _spawn_server(
self,
port: int,
base_env: dict[str, str],
binding_token: str,
artifact_dir: Path,
) -> subprocess.Popen[bytes]:
# Pre-seeded URI when given (realistic corpus), else a throwaway SQLite
# file in the temp dir (the empty-DB path). SQLite absolute paths need
# four slashes; the temp path is absolute.
db_uri = self.database_uri or f"sqlite:///{self._tmp / 'bench.db'}"
args = [
server_executable(),
"-m",
"omnigent.cli",
"server",
"--port",
str(port),
"--database-uri",
db_uri,
"--artifact-location",
str(artifact_dir),
]
env = {**base_env}
if self.with_runner:
# Route the server-side policy-classifier LLM at the mock, mirroring
# live_server. Without this the classifier's client defaults to
# api.openai.com and errors. Server-only mode needs no llm config —
# the classifier only builds under OMNIGENT_SMART_ROUTING=1.
server_cfg = self._tmp / "server.yaml"
server_cfg.write_text(
yaml.safe_dump(
{
"llm": {
"model": _POLICY_LLM_KEY,
"connection": {
"base_url": f"{self.mock_url}/v1",
"api_key": "mock-key",
},
}
}
)
)
args.extend(["--config", str(server_cfg)])
env["OMNIGENT_RUNNER_TUNNEL_TOKEN"] = binding_token
return subprocess.Popen(
args,
env=env,
cwd=compat_server_cwd(),
stdout=self._log("server.log"),
stderr=subprocess.STDOUT,
)
def _spawn_runner(
self, base_env: dict[str, str], binding_token: str
) -> subprocess.Popen[bytes]:
# Point the runner's filesystem workspace at the temp dir so file
# writes (e.g. read_runner_file's setup) land there and are cleaned up
# on teardown, rather than in the launch cwd (its default).
workspace = self._tmp / "workspace"
workspace.mkdir(exist_ok=True)
runner_env = apply_runner_env(
{
**base_env,
"OMNIGENT_RUNNER_ID": self.runner_id,
"OMNIGENT_RUNNER_TUNNEL_BINDING_TOKEN": binding_token,
"OMNIGENT_RUNNER_PARENT_PID": str(os.getpid()),
"RUNNER_SERVER_URL": self.base_url,
"OMNIGENT_RUNNER_WORKSPACE": str(workspace),
}
)
return subprocess.Popen(
[runner_executable(), "-m", "omnigent.runner._entry"],
env=runner_env,
cwd=compat_runner_cwd(),
stdout=self._log("runner.log"),
stderr=subprocess.STDOUT,
)
# ── readiness ────────────────────────────────────────────
def _wait_mock_ready(self) -> None:
deadline = time.monotonic() + _MOCK_TIMEOUT_S
while time.monotonic() < deadline:
try:
if httpx.get(f"{self.mock_url}/stats", timeout=1).status_code == 200:
return
except httpx.HTTPError:
pass
time.sleep(0.1)
raise RuntimeError(f"mock LLM not ready within {_MOCK_TIMEOUT_S}s; logs in {self._tmp}")
def _wait_ready(self) -> None:
"""Wait for ``/health`` (and, in runner mode, the runner online)."""
deadline = time.monotonic() + _HEALTH_TIMEOUT_S
while time.monotonic() < deadline:
try:
health = httpx.get(f"{self.base_url}/health", timeout=2)
if health.status_code == 200 and self._runner_ready():
return
except httpx.HTTPError:
pass
time.sleep(_POLL_INTERVAL_S)
raise RuntimeError(f"server not ready within {_HEALTH_TIMEOUT_S}s; logs in {self._tmp}")
def _runner_ready(self) -> bool:
"""Whether the runner reports online (always ``True`` server-only)."""
if not self.with_runner:
return True
status = httpx.get(f"{self.base_url}/v1/runners/{self.runner_id}/status", timeout=2)
return status.status_code == 200 and status.json().get("online") is True
# ── mock control (runner mode only) ──────────────────────
async def _mock_post(self, path: str, body: dict[str, object]) -> None:
async with httpx.AsyncClient(timeout=5.0) as client:
resp = await client.post(f"{self.mock_url}{path}", json=body)
resp.raise_for_status()
async def configure_mock(
self,
responses: list[dict[str, object]],
*,
key: str = "default",
match: str | None = None,
) -> None:
"""Load a keyed response queue on the mock (see e2e ``configure_mock_llm``)."""
payload: dict[str, object] = {"key": key, "responses": responses}
if match is not None:
payload["match"] = match
await self._mock_post("/mock/configure", payload)
async def set_mock_fallback(
self, text: str, *, key: str = "default", stream: bool = False
) -> None:
"""Set a reset-surviving fallback response for a mock queue *key*.
:param stream: When ``True`` the fallback emits per-word
``output_text.delta`` events before completing — needed for the
time-to-first-token journey to observe streamed deltas.
"""
await self._mock_post("/mock/set_fallback", {"key": key, "text": text, "stream": stream})
# ── agent + session primitives ───────────────────────────
def _agent_bundle(self, name: str) -> bytes:
"""Build a ``spec_version: 1`` agent bundle.
In runner mode the executor is wired at the mock LLM (auth +
connection). Server-only, no LLM is ever called, so the bundle just
needs to be a valid spec the server can register and bind sessions to.
"""
executor: dict[str, object] = {
"type": "omnigent",
"model": self.model,
"config": {"harness": self.harness},
}
config: dict[str, object] = {
"spec_version": 1,
"name": name,
"prompt": "You are a helpful assistant used for performance benchmarking.",
"executor": executor,
}
if self.with_runner:
executor["auth"] = {
"type": "api_key",
"api_key": "mock-key",
"base_url": f"{self.mock_url}/v1",
}
executor["connection"] = {"base_url": f"{self.mock_url}/v1", "api_key": "mock-key"}
# A filesystem env so the runner can serve the resource endpoints
# (read_runner_file). Without os_env the runner has no primary
# environment to materialize and the filesystem proxy 404s.
# sandbox.type=none avoids needing a bwrap binary on the host.
config["os_env"] = {
"type": "caller_process",
"cwd": ".",
"sandbox": {"type": "none"},
}
buf = io.BytesIO()
with tarfile.open(fileobj=buf, mode="w:gz") as tar:
payload = yaml.safe_dump(config).encode()
info = tarfile.TarInfo("config.yaml")
info.size = len(payload)
tar.addfile(info, io.BytesIO(payload))
return buf.getvalue()
async def ensure_agent(self, name: str = "bench-agent") -> str:
"""Register the benchmark agent once, returning its name (idempotent)."""
assert self.client is not None
if name in self._agent_cache:
return name
resp = await self.client.post(
"/v1/sessions",
data={"metadata": "{}"},
files={"bundle": ("agent.tar.gz", self._agent_bundle(name), "application/gzip")},
)
if resp.status_code not in (200, 201, 409):
raise RuntimeError(f"agent register failed: {resp.status_code} {resp.text[:400]}")
self._agent_cache[name] = name
return name
async def agent_id(self, agent_name: str) -> str:
"""Resolve a registered agent's id by name."""
assert self.client is not None
listing = await self.client.get(
"/v1/sessions", params={"agent_name": agent_name, "limit": 1}
)
listing.raise_for_status()
return str(listing.json()["data"][0]["agent_id"])
async def create_session(self, agent_id: str) -> str:
"""Create an (unbound) session for *agent_id*, returning its id."""
assert self.client is not None
created = await self.client.post("/v1/sessions", json={"agent_id": agent_id})
created.raise_for_status()
return str(created.json()["id"])
async def seed_items(self, session_id: str, count: int) -> None:
"""Append *count* history items over HTTP, with no runner or LLM.
Uses the ``external_conversation_item`` event, which the server
appends "without starting or steering a task" — the runner-free path
for giving ``load_conversation_history`` something to read back.
Items are user messages: assistant messages require an ``agent`` field
the server only has after a real turn, and the read path this seeds is
role-agnostic — item count and size, not role, drive its cost.
"""
assert self.client is not None
for i in range(count):
body = {
"type": "external_conversation_item",
"data": {
"item_type": "message",
"item_data": {
"role": "user",
"content": [{"type": "input_text", "text": f"benchmark seed item {i}"}],
},
},
}
resp = await self.client.post(f"/v1/sessions/{session_id}/events", json=body)
resp.raise_for_status()
# ── runner-mode session driving (phase 2) ────────────────
async def create_bound_session(self, agent_id: str) -> str:
"""Create a session for *agent_id* and bind it to the runner."""
assert self.client is not None
if not self.with_runner:
raise RuntimeError("create_bound_session requires with_runner=True")
session_id = await self.create_session(agent_id)
bound = await self.client.patch(
f"/v1/sessions/{session_id}", json={"runner_id": self.runner_id}
)
bound.raise_for_status()
return session_id
async def write_runner_file(self, session_id: str, relative_path: str, content: str) -> None:
"""Write a file into the runner's default environment over HTTP.
The server proxies the ``PUT`` to the bound runner, which writes to its
sandboxed filesystem — so this needs a runner. Used to plant a file the
read journey can then fetch back.
:raises RuntimeError: If not in runner mode.
"""
assert self.client is not None
if not self.with_runner:
raise RuntimeError("write_runner_file requires with_runner=True")
resp = await self.client.put(
f"/v1/sessions/{session_id}/resources/environments/default/filesystem/{relative_path}",
json={"content": content, "encoding": "utf-8"},
)
resp.raise_for_status()
async def read_runner_file(self, session_id: str, relative_path: str) -> None:
"""Read a file from the runner's default environment over HTTP.
Times the server → runner filesystem proxy (a localhost round-trip); no
LLM is involved. Requires a runner — the server returns 502 without one.
:raises RuntimeError: If not in runner mode.
"""
assert self.client is not None
if not self.with_runner:
raise RuntimeError("read_runner_file requires with_runner=True")
resp = await self.client.get(
f"/v1/sessions/{session_id}/resources/environments/default/filesystem/{relative_path}",
)
resp.raise_for_status()
async def drive_turn(
self, session_id: str, text: str, *, timeout: float = _TURN_TIMEOUT_S
) -> None:
"""Post a user message and poll the session to a terminal state.
:raises RuntimeError: If not in runner mode, the turn fails, or it does
not settle within *timeout* seconds.
"""
assert self.client is not None
if not self.with_runner:
raise RuntimeError("drive_turn requires with_runner=True")
body = {
"type": "message",
"data": {"role": "user", "content": [{"type": "input_text", "text": text}]},
}
posted = await self.client.post(f"/v1/sessions/{session_id}/events", json=body)
posted.raise_for_status()
deadline = time.monotonic() + timeout
seen_running = False
while time.monotonic() < deadline:
snap = await self.client.get(f"/v1/sessions/{session_id}")
snap.raise_for_status()
status = snap.json().get("status")
if status in ("running", "waiting"):
seen_running = True
elif status == "failed":
raise RuntimeError(f"turn failed: {snap.json().get('last_task_error')}")
elif status == "idle" and seen_running:
return
await asyncio.sleep(_POLL_INTERVAL_S)
raise RuntimeError(f"turn did not settle within {timeout}s (session {session_id})")
async def _wait_idle(self, session_id: str, *, timeout: float = _TURN_TIMEOUT_S) -> None:
"""Poll until the session is ``idle`` (a prior turn has settled)."""
assert self.client is not None
deadline = time.monotonic() + timeout
while time.monotonic() < deadline:
snap = await self.client.get(f"/v1/sessions/{session_id}")
snap.raise_for_status()
if snap.json().get("status") == "idle":
return
await asyncio.sleep(_POLL_INTERVAL_S)
raise RuntimeError(f"session did not reach idle within {timeout}s ({session_id})")
async def time_to_first_delta(
self, session_id: str, text: str, *, timeout: float = _TURN_TIMEOUT_S
) -> None:
"""Post a turn and return once the first output-text delta streams back.
The session SSE stream (``GET …/stream``) is separate from the message
POST, so we subscribe first (as a concurrent task), post the turn, then
return when the first ``response.output_text.delta`` event arrives. This
times omnigent's streaming-pipeline overhead to first token — with the
zero-latency mock there is no model latency in the number.
:raises RuntimeError: If not in runner mode, or no delta / a terminal
event arrives within *timeout*.
"""
assert self.client is not None
if not self.with_runner:
raise RuntimeError("time_to_first_delta requires with_runner=True")
connected = asyncio.Event()
first_delta = asyncio.Event()
outcome: dict[str, str] = {}
async def _read_stream() -> None:
try:
async with self.client.stream( # type: ignore[union-attr]
"GET", f"/v1/sessions/{session_id}/stream", timeout=timeout
) as resp:
# Any first line means the SSE connection is live (the server
# emits a heartbeat on connect). Signalling here lets us post
# the turn only once subscribed — without a blind sleep that
# would otherwise inflate the measured time-to-first-delta.
connected.set()
async for line in resp.aiter_lines():
if not line.startswith("event:"):
continue
etype = line[len("event:") :].strip()
if etype == "response.output_text.delta":
first_delta.set()
return
if etype in _STREAM_TERMINAL_EVENTS:
outcome["terminal"] = etype
first_delta.set()
return
except httpx.HTTPError as exc:
outcome["error"] = repr(exc)
connected.set()
first_delta.set()
# Ensure any prior turn has settled so the fresh subscription's first
# terminal event can't be the previous turn completing (which would
# otherwise race ahead of this turn's delta).
await self._wait_idle(session_id, timeout=timeout)
reader = asyncio.create_task(_read_stream())
try:
# Wait until the stream is actually connected (not a fixed sleep) so
# the measured window is post → first delta, not subscription setup.
await asyncio.wait_for(connected.wait(), timeout=timeout)
posted = await self.client.post(
f"/v1/sessions/{session_id}/events",
json={
"type": "message",
"data": {"role": "user", "content": [{"type": "input_text", "text": text}]},
},
)
posted.raise_for_status()
try:
await asyncio.wait_for(first_delta.wait(), timeout=timeout)
except TimeoutError as exc:
raise RuntimeError(
f"no output_text.delta within {timeout}s (session {session_id})"
) from exc
if "error" in outcome:
raise RuntimeError(f"stream error: {outcome['error']}")
if "terminal" in outcome:
raise RuntimeError(
f"turn reached {outcome['terminal']} before any delta (session {session_id})"
)
finally:
reader.cancel()
async def drive_and_interrupt(
self, session_id: str, *, timeout: float = _TURN_TIMEOUT_S
) -> None:
"""Drive a gated turn, interrupt it mid-flight, return when cancelled.
The caller configures a ``block=True`` mock response first (see
:meth:`configure_mock`), so the turn parks in ``running`` on the
executor's LLM call. We post an ``interrupt`` once running, wait for the
server's cancellation marker, then release the gate so the runner
unwinds cleanly. Times the server → runner → executor cancel path.
:raises RuntimeError: If not in runner mode, or the interrupt is not
honored within *timeout*.
"""
assert self.client is not None
if not self.with_runner:
raise RuntimeError("drive_and_interrupt requires with_runner=True")
body = {
"type": "message",
"data": {"role": "user", "content": [{"type": "input_text", "text": "Interrupt me."}]},
}
posted = await self.client.post(f"/v1/sessions/{session_id}/events", json=body)
posted.raise_for_status()
deadline = time.monotonic() + timeout
interrupted = False
try:
while time.monotonic() < deadline:
snap = (await self.client.get(f"/v1/sessions/{session_id}")).json()
status = snap.get("status")
items = snap.get("items", [])
if status in ("running", "waiting") and not interrupted:
await self.client.post(
f"/v1/sessions/{session_id}/events", json={"type": "interrupt"}
)
interrupted = True
if _has_cancellation_marker(items):
return
if status == "idle" and interrupted:
if _has_cancellation_marker(items):
return
raise RuntimeError("turn settled without a cancellation marker")
await asyncio.sleep(_POLL_INTERVAL_S)
raise RuntimeError(f"interrupt not honored within {timeout}s (session {session_id})")
finally:
# Always release the gate so the blocked runner turn unwinds and
# teardown doesn't hang, even if the interrupt path errored above.
with contextlib.suppress(httpx.HTTPError):
await self._mock_post("/gate/release", {})
def _has_cancellation_marker(items: list[dict[str, object]]) -> bool:
"""Whether items include the synthetic 'interrupted' user message."""
for raw in items:
data = raw.get("data", raw)
if not isinstance(data, dict):
continue
if raw.get("type") == "message" and data.get("role") == "user":
content = data.get("content") or []
if isinstance(content, list) and any(
isinstance(b, dict) and _CANCELLATION_MARKER in str(b.get("text", ""))
for b in content
):
return True
return False
+559
View File
@@ -0,0 +1,559 @@
"""User-journey definitions and the runners that time them.
A :class:`Journey` names a user-facing operation, an optional per-journey
``setup`` that returns a context object, and a ``measure`` coroutine — the
timed unit. :func:`run_latency` times ``measure`` sequentially; journeys marked
``concurrency_safe`` can also be driven by :func:`run_throughput` with many
operations in flight.
v1 journeys are pure HTTP/API (server + DB, no runner, no LLM):
- ``list_sessions`` — the session-list read behind the sidebar/home.
- ``create_session`` — session creation cost (POST then DELETE).
- ``get_session`` — single-session snapshot load.
- ``load_conversation_history`` — history read, seeded runner-free via
``external_conversation_item`` (see :meth:`BenchEnvironment.seed_items`).
- ``fork_session`` — fork a session (deep-copy its items), then DELETE.
- ``add_comment`` — create a review comment on a file (DB write).
``read_runner_file`` needs a runner but no LLM turn: it plants a file in the
runner environment (setup) and times the server → runner filesystem read proxy.
The framework (``Journey`` + the two runners) is harness-agnostic and reused
verbatim by phase-2 full-turn journeys.
"""
from __future__ import annotations
import asyncio
import contextlib
import time
from collections.abc import Awaitable, Callable
from dataclasses import dataclass
from typing import Literal, cast
import httpx
from .environment import BenchEnvironment
from .measure import RunResult
# Per-journey context returned by ``setup`` and threaded to ``measure``. Its
# concrete type varies by journey (an agent id, a session id, or nothing), so
# it is opaque at the framework level; each measure op casts it as needed.
JourneyContext = object
JourneyKind = Literal["latency", "throughput"]
# Items requested per history-read page. Also the count self-seeded into a
# fallback session when the DB has no corpus (empty-DB smoke path).
_HISTORY_PAGE_LIMIT = 20
_HISTORY_SEED_ITEMS = _HISTORY_PAGE_LIMIT
@dataclass
class Journey:
"""One benchmarkable user journey.
:param name: Stable identifier used on the CLI and as the report key.
:param kind: ``"latency"`` (time each operation) or ``"throughput"``
(fixed request count under concurrency). A latency journey that is
``concurrency_safe`` can additionally be run as throughput.
:param measure: Coroutine performing exactly one timed operation, given
the environment and the setup context.
:param setup: Optional coroutine run once before timing; its return value
is passed to ``measure`` (and ``teardown``) as ``ctx``.
:param teardown: Optional coroutine run once after timing, given ``ctx``.
:param concurrency_safe: Whether many ``measure`` calls may run at once
against a shared setup (true for read-only / independent-write HTTP
journeys).
:param needs_runner: Whether this journey drives a full agent turn and so
requires ``BenchEnvironment(with_runner=True)`` (mock LLM + runner).
HTTP/DB journeys leave this ``False``.
:param max_iterations: Upper bound on latency iterations for this journey,
clamping ``--iterations`` down (never up). Full-turn journeys cost ~1s+
per op, so 100+ iterations would blow the CI time budget; they cap at a
few samples per run and lean on ``--runs`` for repeats. ``None`` (HTTP
journeys) means no cap.
:param description: Human-readable one-liner for ``--list``.
"""
name: str
kind: JourneyKind
measure: Callable[[BenchEnvironment, JourneyContext], Awaitable[None]]
setup: Callable[[BenchEnvironment], Awaitable[JourneyContext]] | None = None
teardown: Callable[[BenchEnvironment, JourneyContext], Awaitable[None]] | None = None
concurrency_safe: bool = False
needs_runner: bool = False
max_iterations: int | None = None
description: str = ""
async def run_setup(self, env: BenchEnvironment) -> JourneyContext:
return await self.setup(env) if self.setup is not None else None
async def run_teardown(self, env: BenchEnvironment, ctx: JourneyContext) -> None:
if self.teardown is not None:
await self.teardown(env, ctx)
# ── timed operation (shared by both runners) ─────────────────
async def _timed(
journey: Journey, env: BenchEnvironment, ctx: JourneyContext, result: RunResult
) -> None:
"""Run one ``measure`` op, recording its latency or a failure reason."""
start = time.perf_counter()
try:
await journey.measure(env, ctx)
except httpx.HTTPStatusError as exc:
result.record_failure(f"HTTP {exc.response.status_code}")
except Exception as exc: # noqa: BLE001 — any failure is a recorded data point
result.record_failure(exc.__class__.__name__)
else:
result.latencies_ms.append((time.perf_counter() - start) * 1000)
# ── runners ──────────────────────────────────────────────────
async def run_latency(
journey: Journey, env: BenchEnvironment, *, iterations: int, warmup: int
) -> RunResult:
"""Time *iterations* sequential operations after discarding *warmup*.
Warmup operations run through the same path but are excluded from the
result, so first-call import/JIT/connection costs don't skew the numbers.
"""
ctx = await journey.run_setup(env)
try:
for _ in range(warmup):
with contextlib.suppress(Exception): # warmup errors are non-fatal
await journey.measure(env, ctx)
result = RunResult()
wall_start = time.perf_counter()
for _ in range(iterations):
await _timed(journey, env, ctx, result)
result.wall_time = time.perf_counter() - wall_start
return result
finally:
await journey.run_teardown(env, ctx)
async def run_throughput(
journey: Journey,
env: BenchEnvironment,
*,
requests: int,
concurrency: int,
warmup: int,
) -> RunResult:
"""Fire *requests* operations with at most *concurrency* in flight.
Wall time spans from the first dispatch to the last completion, so
``throughput`` reflects sustained req/s under load (MLflow's ``_run_once``
shape, with an :class:`asyncio.Semaphore` gate).
"""
ctx = await journey.run_setup(env)
try:
sem = asyncio.Semaphore(concurrency)
async def _one(count_it: bool, result: RunResult) -> None:
async with sem:
if count_it:
await _timed(journey, env, ctx, result)
else:
with contextlib.suppress(Exception): # warmup errors are non-fatal
await journey.measure(env, ctx)
if warmup:
throwaway = RunResult()
await asyncio.gather(*[_one(False, throwaway) for _ in range(warmup)])
result = RunResult()
wall_start = time.perf_counter()
await asyncio.gather(*[_one(True, result) for _ in range(requests)])
result.wall_time = time.perf_counter() - wall_start
return result
finally:
await journey.run_teardown(env, ctx)
# ── journey implementations ──────────────────────────────────
#
# Setups return the context each measure op needs. Ops must be independent so
# concurrency-safe journeys don't interfere across in-flight calls.
# A token present in the seeded corpus (titles + item text, see seed.py
# _FRAGMENTS) so search_sessions exercises the LIKE path with real matches.
_SEARCH_TOKEN = "runner"
async def _setup_agent_id(env: BenchEnvironment) -> str:
"""Register the benchmark agent and return its id."""
name = await env.ensure_agent()
return await env.agent_id(name)
async def _setup_target_session(env: BenchEnvironment) -> str:
"""Return a session id to read: an existing corpus session if any, else make one.
Real runs target a pre-seeded corpus (``seed.py``), so we read a
representative existing session. When the DB is empty (e.g. the smoke test
against a throwaway DB), fall back to creating one with a little history so
the journey still exercises the read path.
"""
assert env.client is not None
listing = await env.client.get("/v1/sessions", params={"limit": 1})
listing.raise_for_status()
data = listing.json().get("data", [])
if data:
return str(data[0]["id"])
# Empty DB: self-seed one session over HTTP (runner-free).
name = await env.ensure_agent()
agent_id = await env.agent_id(name)
session_id = await env.create_session(agent_id)
await env.seed_items(session_id, _HISTORY_SEED_ITEMS)
return session_id
async def _measure_list_sessions(env: BenchEnvironment, _ctx: JourneyContext) -> None:
assert env.client is not None
resp = await env.client.get("/v1/sessions", params={"limit": 20})
resp.raise_for_status()
async def _measure_search_sessions(env: BenchEnvironment, _ctx: JourneyContext) -> None:
assert env.client is not None
resp = await env.client.get(
"/v1/sessions", params={"limit": 20, "search_query": _SEARCH_TOKEN}
)
resp.raise_for_status()
async def _measure_create_session(env: BenchEnvironment, ctx: JourneyContext) -> None:
assert env.client is not None
agent_id = cast(str, ctx) # _setup_agent_id
created = await env.client.post("/v1/sessions", json={"agent_id": agent_id})
created.raise_for_status()
# Delete inline so a long run doesn't accumulate unbounded sessions; the
# POST is the operation of interest and dominates the timed span.
session_id = created.json()["id"]
deleted = await env.client.delete(f"/v1/sessions/{session_id}")
deleted.raise_for_status()
async def _measure_get_session(env: BenchEnvironment, ctx: JourneyContext) -> None:
assert env.client is not None
session_id = cast(str, ctx) # _setup_target_session
resp = await env.client.get(f"/v1/sessions/{session_id}")
resp.raise_for_status()
async def _measure_load_history(env: BenchEnvironment, ctx: JourneyContext) -> None:
assert env.client is not None
session_id = cast(str, ctx) # _setup_target_session
resp = await env.client.get(
f"/v1/sessions/{session_id}/items",
params={"order": "asc", "limit": _HISTORY_PAGE_LIMIT},
)
resp.raise_for_status()
@dataclass
class _ForkContext:
"""Fork-journey context: the session to fork + the forks to clean up.
``measure`` records each fork's id here instead of deleting it inline, so
the DELETE stays out of the timed span; ``teardown`` removes them after.
"""
source_id: str
fork_ids: list[str]
async def _setup_fork_session(env: BenchEnvironment) -> _ForkContext:
"""Resolve a session to fork; start an empty fork-id collector."""
source_id = await _setup_target_session(env)
return _ForkContext(source_id=source_id, fork_ids=[])
async def _measure_fork_session(env: BenchEnvironment, ctx: JourneyContext) -> None:
assert env.client is not None
fork_ctx = cast(_ForkContext, ctx) # _setup_fork_session
forked = await env.client.post(f"/v1/sessions/{fork_ctx.source_id}/fork", json={})
forked.raise_for_status()
# Record the fork for teardown; deleting it here would fold the DELETE into
# the timed span. The fork POST (a deep-copy of the source's items) is the
# operation of interest.
fork_ctx.fork_ids.append(forked.json()["id"])
async def _teardown_fork_session(env: BenchEnvironment, ctx: JourneyContext) -> None:
"""Delete every fork created during the run (best effort, untimed)."""
assert env.client is not None
fork_ctx = cast(_ForkContext, ctx)
for fork_id in fork_ctx.fork_ids:
with contextlib.suppress(httpx.HTTPError):
await env.client.delete(f"/v1/sessions/{fork_id}")
# Anchor snapshot for the comment journey; the offsets below span it.
_COMMENT_ANCHOR = "benchmark"
async def _measure_add_comment(env: BenchEnvironment, ctx: JourneyContext) -> None:
assert env.client is not None
session_id = cast(str, ctx) # _setup_target_session
# Each POST creates an independent comment row. Unlike sessions, an
# accumulating comment skews no measured read path, so there's no cleanup.
# The file need not exist — the handler stores the path + offsets + body.
resp = await env.client.post(
f"/v1/sessions/{session_id}/comments",
json={
"path": "bench_target.py",
"body": "benchmark review comment",
"start_index": 0,
"end_index": len(_COMMENT_ANCHOR),
"anchor_content": _COMMENT_ANCHOR,
},
)
resp.raise_for_status()
# ── runner (full-turn) journeys ──────────────────────────────
#
# These drive a real agent turn through the runner + mock LLM (with_runner=True,
# openai-agents). The mock is zero-latency, so every number is omnigent dispatch
# / streaming / cancel overhead, not model latency. Short deterministic replies.
# A multi-word reply so the streaming path emits several output_text deltas.
_TURN_REPLY = "Hello there, this is a mock benchmark reply."
_TURN_PROMPT = "Say hello."
# Iteration cap for full-turn journeys. At ~1s+ per turn, matching the HTTP
# journeys' iteration count would overrun the CI time budget, so we take a few
# samples per run and lean on --runs for repeats. Sessions accumulate across a
# run (a cold start never deletes its session), so a small count also keeps that
# drift negligible.
_RUNNER_MAX_ITERATIONS = 5
# Iteration cap for the runner filesystem read. It's a proxied localhost read,
# not a full turn, so it's far cheaper than the drive-a-turn journeys — a higher
# cap gives a usable p50/p99 while staying well within the CI time budget.
_RUNNER_FS_MAX_ITERATIONS = 50
# File planted by the read-runner-file setup and fetched by its measure op.
# ~1 KB — a modest, representative source file, not a stress case.
_RUNNER_FILE_PATH = "bench_read_target.txt"
_RUNNER_FILE_CONTENT = "benchmark file content line\n" * 40
async def _setup_turn_agent(env: BenchEnvironment, *, stream: bool = False) -> str:
"""Register the agent + a reset-surviving reply; return the agent id.
The fallback survives per-call queue exhaustion, so every turn in the run
gets the same reply regardless of how many turns consume the queue. When
*stream* is set the reply emits per-word deltas (for the TTFT journey).
"""
name = await env.ensure_agent()
await env.set_mock_fallback(_TURN_REPLY, stream=stream)
return await env.agent_id(name)
async def _setup_warm_session(env: BenchEnvironment) -> str:
"""Create+bind a session and drive one warm-up turn; return the session id.
The warm-up pays the cold-start cost (runner spawn + executor construction)
so the measured op times only steady-state per-turn overhead.
"""
agent_id = await _setup_turn_agent(env)
session_id = await env.create_bound_session(agent_id)
await env.drive_turn(session_id, _TURN_PROMPT)
return session_id
async def _setup_streaming_session(env: BenchEnvironment) -> str:
"""Warm session whose mock reply streams deltas — for the TTFT journey."""
agent_id = await _setup_turn_agent(env, stream=True)
session_id = await env.create_bound_session(agent_id)
await env.drive_turn(session_id, _TURN_PROMPT)
return session_id
async def _setup_interrupt_session(env: BenchEnvironment) -> str:
"""Create+bind a session for the interrupt journey; return the session id.
Configures a ``block=True`` mock response so each turn parks in ``running``
until the gate is released — giving the interrupt something to cancel
mid-flight, deterministically.
"""
name = await env.ensure_agent()
agent_id = await env.agent_id(name)
session_id = await env.create_bound_session(agent_id)
await env.configure_mock([{"text": _TURN_REPLY, "block": True}])
return session_id
async def _measure_session_cold_start(env: BenchEnvironment, ctx: JourneyContext) -> None:
agent_id = cast(str, ctx) # _setup_turn_agent
session_id = await env.create_bound_session(agent_id)
await env.drive_turn(session_id, _TURN_PROMPT)
async def _measure_warm_turn(env: BenchEnvironment, ctx: JourneyContext) -> None:
session_id = cast(str, ctx) # _setup_warm_session
await env.drive_turn(session_id, _TURN_PROMPT)
async def _measure_time_to_first_token(env: BenchEnvironment, ctx: JourneyContext) -> None:
session_id = cast(str, ctx) # _setup_warm_session
await env.time_to_first_delta(session_id, _TURN_PROMPT)
async def _measure_interrupt(env: BenchEnvironment, ctx: JourneyContext) -> None:
session_id = cast(str, ctx) # _setup_interrupt_session
await env.drive_and_interrupt(session_id)
async def _setup_runner_file_session(env: BenchEnvironment) -> str:
"""Bind a session to the runner and plant a file to read; return its id.
No turn is driven and no mock reply is configured — the measured op is a
filesystem read proxied to the runner, which never calls the LLM.
"""
name = await env.ensure_agent()
agent_id = await env.agent_id(name)
session_id = await env.create_bound_session(agent_id)
await env.write_runner_file(session_id, _RUNNER_FILE_PATH, _RUNNER_FILE_CONTENT)
return session_id
async def _measure_read_runner_file(env: BenchEnvironment, ctx: JourneyContext) -> None:
session_id = cast(str, ctx) # _setup_runner_file_session
await env.read_runner_file(session_id, _RUNNER_FILE_PATH)
# ── registry ─────────────────────────────────────────────────
ALL_JOURNEYS: dict[str, Journey] = {
j.name: j
for j in (
Journey(
name="list_sessions",
kind="latency",
measure=_measure_list_sessions,
concurrency_safe=True,
description="GET /v1/sessions — session list read.",
),
Journey(
name="create_session",
kind="latency",
measure=_measure_create_session,
setup=_setup_agent_id,
concurrency_safe=True,
description="POST /v1/sessions then DELETE — session create.",
),
Journey(
name="get_session",
kind="latency",
measure=_measure_get_session,
setup=_setup_target_session,
concurrency_safe=True,
description="GET /v1/sessions/{id} — single-session snapshot.",
),
Journey(
name="load_conversation_history",
kind="latency",
measure=_measure_load_history,
setup=_setup_target_session,
concurrency_safe=True,
description="GET /v1/sessions/{id}/items — conversation history read.",
),
Journey(
name="search_sessions",
kind="latency",
measure=_measure_search_sessions,
concurrency_safe=True,
description="GET /v1/sessions?search_query= — unindexed LIKE over titles + items.",
),
Journey(
name="fork_session",
kind="latency",
measure=_measure_fork_session,
setup=_setup_fork_session,
teardown=_teardown_fork_session,
concurrency_safe=True,
description="POST /v1/sessions/{id}/fork — session fork (deep-copy); DELETE untimed.",
),
Journey(
name="add_comment",
kind="latency",
measure=_measure_add_comment,
setup=_setup_target_session,
concurrency_safe=True,
description="POST /v1/sessions/{id}/comments — create a review comment.",
),
# Runner (full-turn) journeys — with_runner=True, openai-agents, mock LLM.
Journey(
name="session_cold_start",
kind="latency",
measure=_measure_session_cold_start,
setup=_setup_turn_agent,
needs_runner=True,
max_iterations=_RUNNER_MAX_ITERATIONS,
description="Create+bind a fresh session and drive its first turn to idle.",
),
Journey(
name="warm_turn",
kind="latency",
measure=_measure_warm_turn,
setup=_setup_warm_session,
needs_runner=True,
max_iterations=_RUNNER_MAX_ITERATIONS,
description="Drive a turn on an already-warm session (steady-state overhead).",
),
Journey(
name="time_to_first_token",
kind="latency",
measure=_measure_time_to_first_token,
setup=_setup_streaming_session,
needs_runner=True,
max_iterations=_RUNNER_MAX_ITERATIONS,
description="Post a turn; time to the first streamed output_text delta.",
),
Journey(
name="interrupt",
kind="latency",
measure=_measure_interrupt,
setup=_setup_interrupt_session,
needs_runner=True,
max_iterations=_RUNNER_MAX_ITERATIONS,
description="Interrupt a running (gated) turn; time to cancellation.",
),
Journey(
name="read_runner_file",
kind="latency",
measure=_measure_read_runner_file,
setup=_setup_runner_file_session,
needs_runner=True,
max_iterations=_RUNNER_FS_MAX_ITERATIONS,
description="GET .../environments/default/filesystem/{path} — runner file read proxy.",
),
)
}
def resolve_journeys(names: list[str] | None) -> list[Journey]:
"""Resolve requested journey *names* (or all when ``None``/empty).
:raises KeyError: If a requested name isn't registered.
"""
if not names:
return list(ALL_JOURNEYS.values())
resolved = []
for name in names:
if name not in ALL_JOURNEYS:
raise KeyError(f"unknown journey {name!r}; known: {', '.join(ALL_JOURNEYS)}")
resolved.append(ALL_JOURNEYS[name])
return resolved
+222
View File
@@ -0,0 +1,222 @@
"""Latency/throughput measurement primitives.
Pure and I/O-free: a :class:`RunResult` accumulates per-operation latencies
and failures for one timed run, :func:`aggregate` folds several runs into the
``runs`` + ``summary`` shape the workspace ETL flattens, and
:func:`check_thresholds` gates a run in CI. Adapted from MLflow's
``dev/benchmarks/gateway/benchmark.py``.
"""
from __future__ import annotations
import math
import statistics
from dataclasses import dataclass, field
from rich.console import Console
from rich.table import Table
console = Console()
@dataclass
class RunResult:
"""Latencies and failures collected during one timed run.
:param latencies_ms: Per-operation wall-clock latency in milliseconds,
one entry per successful operation.
:param failures: Failure reason (e.g. ``"HTTP 500"`` / an exception
class name) mapped to how many times it occurred.
:param wall_time: Total elapsed seconds for the run, used for throughput.
"""
latencies_ms: list[float] = field(default_factory=list)
failures: dict[str, int] = field(default_factory=dict)
wall_time: float = 0.0
@property
def n_success(self) -> int:
"""Number of operations that completed without error."""
return len(self.latencies_ms)
@property
def n_failures(self) -> int:
"""Total failed operations across all reasons."""
return sum(self.failures.values())
@property
def throughput(self) -> float:
"""Successful operations per second over the run's wall time."""
return self.n_success / self.wall_time if self.wall_time > 0 else 0.0
def record_failure(self, reason: str) -> None:
"""Increment the count for one failure *reason*."""
self.failures[reason] = self.failures.get(reason, 0) + 1
def percentile(self, p: float) -> float:
"""Return the *p*-th percentile latency in ms (ceil-index method).
:param p: Percentile in ``[0, 100]``, e.g. ``99`` for p99.
:returns: The latency at that percentile, or ``0.0`` when no
successful operation was recorded.
"""
if not self.latencies_ms:
return 0.0
ordered = sorted(self.latencies_ms)
idx = max(0, math.ceil(p / 100 * len(ordered)) - 1)
return ordered[idx]
def mean_ms(self) -> float:
"""Mean latency in ms, or ``0.0`` when no operation succeeded."""
return statistics.mean(self.latencies_ms) if self.latencies_ms else 0.0
def max_ms(self) -> float:
"""Maximum latency in ms, or ``0.0`` when no operation succeeded."""
return max(self.latencies_ms) if self.latencies_ms else 0.0
def _run_to_dict(result: RunResult) -> dict[str, object]:
"""Flatten one :class:`RunResult` into a JSON-serializable per-run row."""
return {
"n_success": result.n_success,
"n_failures": result.n_failures,
"failures": dict(result.failures),
"wall_time_s": result.wall_time,
"mean_ms": result.mean_ms(),
"p50_ms": result.percentile(50),
"p95_ms": result.percentile(95),
"p99_ms": result.percentile(99),
"max_ms": result.max_ms(),
"rps": result.throughput,
}
def aggregate(results: list[RunResult]) -> dict[str, object]:
"""Fold per-run results into ``{"runs": [...], "summary": {...}}``.
The ``summary`` averages each metric across runs. Its keys mirror
MLflow's gateway benchmark (``avg_mean_ms`` / ``avg_p50_ms`` /
``avg_p99_ms`` / ``avg_rps``) plus ``avg_p95_ms``, so the workspace ETL
that flattens ``summary`` works unchanged.
:param results: One :class:`RunResult` per timed run (warmup excluded).
:returns: A dict with a per-run ``runs`` list and an averaged
``summary`` (empty ``summary`` when *results* is empty).
"""
runs = [_run_to_dict(r) for r in results]
if not results:
return {"runs": runs, "summary": {}}
summary = {
"avg_mean_ms": statistics.mean(r.mean_ms() for r in results),
"avg_p50_ms": statistics.mean(r.percentile(50) for r in results),
"avg_p95_ms": statistics.mean(r.percentile(95) for r in results),
"avg_p99_ms": statistics.mean(r.percentile(99) for r in results),
"avg_rps": statistics.mean(r.throughput for r in results),
}
return {"runs": runs, "summary": summary}
def check_thresholds(
results: list[RunResult],
*,
min_rps: float | None = None,
max_p50_ms: float | None = None,
max_p99_ms: float | None = None,
) -> bool:
"""Check averaged results against optional CI thresholds.
:param results: Timed runs for one journey.
:param min_rps: Fail if average throughput is below this (req/s).
:param max_p50_ms: Fail if average p50 latency exceeds this (ms).
:param max_p99_ms: Fail if average p99 latency exceeds this (ms).
:returns: ``True`` when every supplied threshold passes (vacuously
true when none are supplied or *results* is empty).
"""
if not results:
return True
avg_rps = statistics.mean(r.throughput for r in results)
avg_p50 = statistics.mean(r.percentile(50) for r in results)
avg_p99 = statistics.mean(r.percentile(99) for r in results)
passed = True
if min_rps is not None and avg_rps < min_rps:
console.print(
f" [red]THRESHOLD FAILED:[/red] avg throughput {avg_rps:.0f} req/s"
f" < minimum {min_rps:.0f} req/s"
)
passed = False
if max_p50_ms is not None and avg_p50 > max_p50_ms:
console.print(
f" [red]THRESHOLD FAILED:[/red] avg P50 {avg_p50:.1f} ms"
f" > maximum {max_p50_ms:.1f} ms"
)
passed = False
if max_p99_ms is not None and avg_p99 > max_p99_ms:
console.print(
f" [red]THRESHOLD FAILED:[/red] avg P99 {avg_p99:.1f} ms"
f" > maximum {max_p99_ms:.1f} ms"
)
passed = False
return passed
def print_results(journey_name: str, results: list[RunResult]) -> None:
"""Render per-run and averaged metrics for one journey as a rich table.
:param journey_name: Journey label used as the table title.
:param results: Timed runs to display.
"""
table = Table(
title=journey_name,
show_header=True,
header_style="bold cyan",
box=None,
padding=(0, 2),
title_justify="left",
)
table.add_column("Run", style="dim", width=5)
table.add_column("Mean ms", justify="right")
table.add_column("P50 ms", justify="right")
table.add_column("P95 ms", justify="right")
table.add_column("P99 ms", justify="right")
table.add_column("Max ms", justify="right")
table.add_column("Req/s", justify="right")
table.add_column("Failures", justify="right")
for i, r in enumerate(results):
fail_str = f"[red]{r.n_failures}[/red]" if r.n_failures else "0"
table.add_row(
str(i + 1),
f"{r.mean_ms():.1f}",
f"{r.percentile(50):.1f}",
f"{r.percentile(95):.1f}",
f"{r.percentile(99):.1f}",
f"{r.max_ms():.1f}",
f"{r.throughput:.0f}",
fail_str,
)
if len(results) > 1:
table.add_section()
table.add_row(
"[bold]avg[/bold]",
f"[bold]{statistics.mean(r.mean_ms() for r in results):.1f}[/bold]",
f"[bold]{statistics.mean(r.percentile(50) for r in results):.1f}[/bold]",
f"[bold]{statistics.mean(r.percentile(95) for r in results):.1f}[/bold]",
f"[bold]{statistics.mean(r.percentile(99) for r in results):.1f}[/bold]",
f"[bold]{statistics.mean(r.max_ms() for r in results):.1f}[/bold]",
f"[bold]{statistics.mean(r.throughput for r in results):.0f}[/bold]",
"",
)
console.print()
console.print(table)
combined: dict[str, int] = {}
for r in results:
for reason, count in r.failures.items():
combined[reason] = combined.get(reason, 0) + count
if combined:
console.print(" [red]Failure breakdown:[/red]")
for reason, count in sorted(combined.items(), key=lambda kv: -kv[1]):
console.print(f" {reason}: {count}")
+268
View File
@@ -0,0 +1,268 @@
"""Omnigent user-journey benchmark runner.
Boots a real ``omnigent server`` against a SQLite DB (no runner, no LLM),
drives the selected HTTP journeys under load, prints per-journey latency /
throughput tables, and writes a versioned JSON report. Exits non-zero if any
supplied threshold is breached.
Runs in the project venv — it imports ``omnigent`` and ``tests._helpers`` and
spawns the real server, so it is NOT a standalone PEP 723 script. Invoke with
``--no-sync`` so ``uv`` uses the existing environment instead of rebuilding the
project (which triggers a web-UI build that fails in a worktree)::
uv run --no-sync dev/benchmarks/omnigent/run.py
uv run --no-sync dev/benchmarks/omnigent/run.py --journeys list_sessions,get_session
uv run --no-sync dev/benchmarks/omnigent/run.py --requests 500 --concurrency 25 --runs 3
uv run --no-sync dev/benchmarks/omnigent/run.py --output bench.json --max-p50-ms 25
The JSON is the contract consumed by the workspace Databricks ETL notebook —
see ``README.md``.
"""
from __future__ import annotations
import argparse
import asyncio
import datetime
import json
import sys
from pathlib import Path
# Allow ``uv run <path>`` (no package context) to import the sibling modules.
sys.path.insert(0, str(Path(__file__).resolve().parents[3]))
from dev.benchmarks.omnigent.environment import BenchEnvironment
from dev.benchmarks.omnigent.journeys import (
ALL_JOURNEYS,
Journey,
resolve_journeys,
run_latency,
run_throughput,
)
from dev.benchmarks.omnigent.measure import (
RunResult,
aggregate,
check_thresholds,
console,
print_results,
)
from dev.benchmarks.omnigent.schema import build_report
# Harness label stamped in the report: HTTP/DB journeys drive no agent turn;
# runner journeys drive turns through the in-process openai-agents SDK harness.
_HTTP_HARNESS = "http-only"
_RUNNER_HARNESS = "openai-agents"
def _backend_of(database_uri: str | None) -> str:
"""Classify the DB URI into a coarse backend label for the report.
``None`` is the harness's throwaway SQLite temp file. Otherwise key off the
URI scheme so the report (and the workspace dashboard) can group by backend.
"""
if database_uri is None or database_uri.startswith("sqlite"):
return "sqlite"
if database_uri.startswith("postgres"):
return "postgres"
if database_uri.startswith("mysql"):
return "mysql"
return "other"
def _effective_iterations(journey: Journey, requested: int) -> int:
"""Clamp *requested* iterations down to the journey's ``max_iterations``.
Full-turn journeys cost ~1s+ per op and cap themselves so a large
``--iterations`` (tuned for the millisecond HTTP journeys) doesn't overrun
the CI time budget. The cap only ever lowers the count, never raises it.
"""
if journey.max_iterations is not None:
return min(requested, journey.max_iterations)
return requested
async def _run_journey(
journey: Journey, env: BenchEnvironment, args: argparse.Namespace
) -> tuple[str, list[RunResult]]:
"""Run one journey's timed runs, returning its report kind + per-run results.
A journey runs as throughput when ``--concurrency > 1`` and it is
concurrency-safe; otherwise as sequential latency.
"""
as_throughput = args.concurrency > 1 and journey.concurrency_safe
iterations = _effective_iterations(journey, args.iterations)
results: list[RunResult] = []
for _ in range(args.runs):
if as_throughput:
results.append(
await run_throughput(
journey,
env,
requests=args.requests,
concurrency=args.concurrency,
warmup=args.warmup,
)
)
else:
results.append(
await run_latency(journey, env, iterations=iterations, warmup=args.warmup)
)
return ("throughput" if as_throughput else "latency"), results
async def run_benchmark(args: argparse.Namespace) -> tuple[dict[str, object], bool]:
"""Run all selected journeys and build the report.
:returns: ``(report, passed)`` where *passed* is ``False`` if any journey
breached a supplied threshold.
"""
journeys = resolve_journeys(args.journeys)
journey_results: dict[str, dict[str, object]] = {}
passed = True
backend = _backend_of(args.database_uri)
# Any full-turn journey needs the runner + mock LLM. A full env is a
# superset — HTTP journeys still run against it — so a mixed selection just
# boots with_runner=True. The harness label reflects what drove the turns.
with_runner = any(j.needs_runner for j in journeys)
harness = _RUNNER_HARNESS if with_runner else _HTTP_HARNESS
async with BenchEnvironment(with_runner=with_runner, database_uri=args.database_uri) as env:
for journey in journeys:
console.print(f"\n[bold]Benchmarking[/bold] {journey.name} [dim]({backend})[/dim]")
kind, results = await _run_journey(journey, env, args)
print_results(journey.name, results)
block = aggregate(results)
block["kind"] = kind
block["backend"] = backend
journey_results[journey.name] = block
if not check_thresholds(
results,
min_rps=args.min_rps,
max_p50_ms=args.max_p50_ms,
max_p99_ms=args.max_p99_ms,
):
passed = False
config = {
"iterations": args.iterations,
"requests": args.requests,
"concurrency": args.concurrency,
"runs": args.runs,
"warmup": args.warmup,
"with_runner": with_runner,
"backend": backend,
}
generated_at = datetime.datetime.now(datetime.timezone.utc).isoformat()
report = build_report(
journey_results,
generated_at=generated_at,
config=config,
harness=harness,
)
return report, passed
def _parse_args(argv: list[str]) -> argparse.Namespace:
parser = argparse.ArgumentParser(
prog="omnigent-benchmark",
description=__doc__,
formatter_class=argparse.RawDescriptionHelpFormatter,
)
parser.add_argument(
"--journeys",
type=lambda s: [p.strip() for p in s.split(",") if p.strip()],
default=None,
metavar="A,B,C",
help=f"Comma-separated journeys to run. Default: all ({', '.join(ALL_JOURNEYS)}).",
)
parser.add_argument(
"--database-uri",
default=None,
metavar="URI",
help="DB the server boots against — a pre-seeded SQLite file, a "
"postgresql+psycopg://… instance, or a mysql+mysqldb://… instance "
"(see seed.py). Default: a fresh throwaway SQLite DB (empty — "
"best-case numbers). The report's `backend` field is derived from this.",
)
parser.add_argument(
"--iterations",
type=int,
default=100,
metavar="N",
help="Sequential operations per latency run (default: 100).",
)
parser.add_argument(
"--requests",
type=int,
default=500,
metavar="N",
help="Total operations per throughput run — used when --concurrency>1 (default: 500).",
)
parser.add_argument(
"--concurrency",
type=int,
default=1,
metavar="N",
help="Max in-flight operations. >1 runs concurrency-safe journeys as "
"throughput (default: 1 = sequential latency).",
)
parser.add_argument(
"--runs",
type=int,
default=3,
metavar="N",
help="Timed runs per journey; results are per-run and averaged (default: 3).",
)
parser.add_argument(
"--warmup",
type=int,
default=10,
metavar="N",
help="Warmup operations discarded before each run (default: 10).",
)
parser.add_argument(
"--output",
type=Path,
default=None,
metavar="FILE",
help="Write the JSON report to FILE (for CI artifact upload).",
)
parser.add_argument(
"--min-rps",
type=float,
default=None,
metavar="N",
help="Exit 1 if any journey's avg throughput falls below N req/s.",
)
parser.add_argument(
"--max-p50-ms",
type=float,
default=None,
metavar="N",
help="Exit 1 if any journey's avg P50 latency exceeds N ms.",
)
parser.add_argument(
"--max-p99-ms",
type=float,
default=None,
metavar="N",
help="Exit 1 if any journey's avg P99 latency exceeds N ms.",
)
return parser.parse_args(argv)
def main(argv: list[str] | None = None) -> int:
args = _parse_args(argv if argv is not None else sys.argv[1:])
report, passed = asyncio.run(run_benchmark(args))
if args.output is not None:
args.output.write_text(json.dumps(report, indent=2))
console.print(f"\n Results written to [cyan]{args.output}[/cyan]")
if not passed:
console.print("\n[red]One or more thresholds failed.[/red]")
return 1
return 0
if __name__ == "__main__":
raise SystemExit(main())
+268
View File
@@ -0,0 +1,268 @@
{
"schema_version": 1,
"generated_at": "2026-07-08T18:30:00+00:00",
"git_sha": "0000000000000000000000000000000000000000",
"git_branch": "main",
"host": {
"platform": "macOS-15.5-arm64-arm-64bit",
"python": "3.12.8",
"cpu_count": 12
},
"harness": "http-only",
"config": {
"iterations": 100,
"requests": 500,
"concurrency": 1,
"runs": 3,
"warmup": 10,
"with_runner": false,
"backend": "sqlite"
},
"journeys": {
"list_sessions": {
"runs": [
{
"n_success": 100,
"n_failures": 0,
"failures": {},
"wall_time_s": 0.6587757079978473,
"mean_ms": 6.586994149838574,
"p50_ms": 6.261250004172325,
"p95_ms": 7.65325000975281,
"p99_ms": 7.9127089702524245,
"max_ms": 38.27937500318512,
"rps": 151.79673261468648
},
{
"n_success": 100,
"n_failures": 0,
"failures": {},
"wall_time_s": 0.6238234580378048,
"mean_ms": 6.237602901528589,
"p50_ms": 6.126708001829684,
"p95_ms": 7.152916979975998,
"p99_ms": 7.425624993629754,
"max_ms": 7.667875033803284,
"rps": 160.30176280087855
},
{
"n_success": 100,
"n_failures": 0,
"failures": {},
"wall_time_s": 0.5675292499945499,
"mean_ms": 5.674769564066082,
"p50_ms": 5.528832960408181,
"p95_ms": 6.237249996047467,
"p99_ms": 7.393250009045005,
"max_ms": 11.83562504593283,
"rps": 176.20237194992208
}
],
"summary": {
"avg_mean_ms": 6.166455538477749,
"avg_p50_ms": 5.972263655470063,
"avg_p95_ms": 7.014472328592092,
"avg_p99_ms": 7.577194657642394,
"avg_rps": 162.7669557884957
},
"kind": "latency",
"backend": "sqlite"
},
"create_session": {
"runs": [
{
"n_success": 100,
"n_failures": 0,
"failures": {},
"wall_time_s": 2.4451411250047386,
"mean_ms": 24.450668342760764,
"p50_ms": 24.028874991927296,
"p95_ms": 27.25041698431596,
"p99_ms": 29.374166973866522,
"max_ms": 29.56758299842477,
"rps": 40.89743490769115
},
{
"n_success": 100,
"n_failures": 0,
"failures": {},
"wall_time_s": 2.498600291030016,
"mean_ms": 24.985182073432952,
"p50_ms": 24.459875014144927,
"p95_ms": 28.391665953677148,
"p99_ms": 29.0600000298582,
"max_ms": 34.39550002804026,
"rps": 40.02240788932922
},
{
"n_success": 100,
"n_failures": 0,
"failures": {},
"wall_time_s": 2.455423459003214,
"mean_ms": 24.553418274153955,
"p50_ms": 24.073333013802767,
"p95_ms": 27.43383398046717,
"p99_ms": 29.375000041909516,
"max_ms": 29.430416005197912,
"rps": 40.72617276394162
}
],
"summary": {
"avg_mean_ms": 24.663089563449223,
"avg_p50_ms": 24.187361006624997,
"avg_p95_ms": 27.691972306153428,
"avg_p99_ms": 29.269722348544747,
"avg_rps": 40.548671853654
},
"kind": "latency",
"backend": "sqlite"
},
"get_session": {
"runs": [
{
"n_success": 100,
"n_failures": 0,
"failures": {},
"wall_time_s": 0.5546917500323616,
"mean_ms": 5.5460037814918905,
"p50_ms": 5.360124981962144,
"p95_ms": 6.925499998033047,
"p99_ms": 7.144333969336003,
"max_ms": 7.331291970331222,
"rps": 180.28030882767922
},
{
"n_success": 100,
"n_failures": 0,
"failures": {},
"wall_time_s": 0.49645508296089247,
"mean_ms": 4.963922067545354,
"p50_ms": 4.782959003932774,
"p95_ms": 5.978292028885335,
"p99_ms": 6.7617910099215806,
"max_ms": 6.881375040393323,
"rps": 201.42809174919327
},
{
"n_success": 100,
"n_failures": 0,
"failures": {},
"wall_time_s": 0.460644083970692,
"mean_ms": 4.605880451854318,
"p50_ms": 4.526541975792497,
"p95_ms": 5.18629199359566,
"p99_ms": 5.445250018965453,
"max_ms": 5.790999974124134,
"rps": 217.08734244020465
}
],
"summary": {
"avg_mean_ms": 5.0386021002971875,
"avg_p50_ms": 4.889875320562472,
"avg_p95_ms": 6.030028006838013,
"avg_p99_ms": 6.450458332741012,
"avg_rps": 199.59858100569238
},
"kind": "latency",
"backend": "sqlite"
},
"load_conversation_history": {
"runs": [
{
"n_success": 100,
"n_failures": 0,
"failures": {},
"wall_time_s": 0.20811437495285645,
"mean_ms": 2.080691678565927,
"p50_ms": 2.037000027485192,
"p95_ms": 2.5742079596966505,
"p99_ms": 2.768124977592379,
"max_ms": 2.784749958664179,
"rps": 480.50501087516284
},
{
"n_success": 100,
"n_failures": 0,
"failures": {},
"wall_time_s": 0.19513549999101087,
"mean_ms": 1.9509158097207546,
"p50_ms": 1.9018329912796617,
"p95_ms": 2.284207963384688,
"p99_ms": 2.4481670116074383,
"max_ms": 2.5021659675985575,
"rps": 512.4644157757384
},
{
"n_success": 100,
"n_failures": 0,
"failures": {},
"wall_time_s": 0.19278304203180596,
"mean_ms": 1.9274150469573215,
"p50_ms": 1.8819589749909937,
"p95_ms": 2.2150420118123293,
"p99_ms": 2.2878749878145754,
"max_ms": 2.316958038136363,
"rps": 518.7178236532945
}
],
"summary": {
"avg_mean_ms": 1.9863408450813342,
"avg_p50_ms": 1.9402639979186158,
"avg_p95_ms": 2.3578193116312227,
"avg_p99_ms": 2.5013889923381307,
"avg_rps": 503.8957501013986
},
"kind": "latency",
"backend": "sqlite"
},
"search_sessions": {
"runs": [
{
"n_success": 100,
"n_failures": 0,
"failures": {},
"wall_time_s": 8.150427000015043,
"mean_ms": 81.50338126753923,
"p50_ms": 80.11816703947261,
"p95_ms": 90.9090840141289,
"p99_ms": 94.67683301772922,
"max_ms": 96.44366696011275,
"rps": 12.26929582950874
},
{
"n_success": 100,
"n_failures": 0,
"failures": {},
"wall_time_s": 8.152122708968818,
"mean_ms": 81.5203430026304,
"p50_ms": 79.57212498877198,
"p95_ms": 95.20591603359208,
"p99_ms": 98.80137501750141,
"max_ms": 99.93629204109311,
"rps": 12.266743714490682
},
{
"n_success": 100,
"n_failures": 0,
"failures": {},
"wall_time_s": 8.053999124967959,
"mean_ms": 80.53909589187242,
"p50_ms": 79.52124997973442,
"p95_ms": 91.39445802429691,
"p99_ms": 93.2748339837417,
"max_ms": 94.79766699951142,
"rps": 12.416192061654566
}
],
"summary": {
"avg_mean_ms": 81.18760672068068,
"avg_p50_ms": 79.73718066932634,
"avg_p95_ms": 92.50315269067262,
"avg_p99_ms": 95.58434733965744,
"avg_rps": 12.317410535217997
},
"kind": "latency",
"backend": "sqlite"
}
}
}
+89
View File
@@ -0,0 +1,89 @@
"""Benchmark report schema + metadata capture.
:func:`build_report` assembles the single JSON document the harness writes.
Its per-journey ``summary`` + ``runs`` shape mirrors MLflow's gateway
benchmark so the workspace ETL notebook flattens it unchanged — keyed by
journey (and ``harness``) instead of ``backend``. Bump :data:`SCHEMA_VERSION`
whenever the document's shape changes so the ETL can branch on it.
"""
from __future__ import annotations
import platform
import subprocess
# Incremented on any breaking change to the report document shape below.
SCHEMA_VERSION = 1
def _git(*args: str) -> str:
"""Run ``git *args`` at the repo root, returning stripped stdout or ``""``.
Never raises: a missing git, detached checkout, or non-zero exit all
surface as an empty string so a benchmark run outside a clean checkout
still produces a valid report.
"""
try:
out = subprocess.run(
["git", *args],
capture_output=True,
text=True,
timeout=10,
check=False,
)
except (OSError, subprocess.SubprocessError):
return ""
return out.stdout.strip() if out.returncode == 0 else ""
def git_sha() -> str:
"""Return the current commit SHA, or ``""`` when unavailable."""
return _git("rev-parse", "HEAD")
def git_branch() -> str:
"""Return the current branch name, or ``""`` when detached/unavailable."""
return _git("rev-parse", "--abbrev-ref", "HEAD")
def host_info() -> dict[str, object]:
"""Capture coarse host facts for cross-machine result comparison."""
import os
return {
"platform": platform.platform(),
"python": platform.python_version(),
"cpu_count": os.cpu_count(),
}
def build_report(
journey_results: dict[str, dict[str, object]],
*,
generated_at: str,
config: dict[str, object],
harness: str,
) -> dict[str, object]:
"""Assemble the full benchmark report document.
:param journey_results: Per-journey ``{"kind", "runs", "summary"}``
blocks (each ``runs``/``summary`` produced by
:func:`measure.aggregate`), keyed by journey name.
:param generated_at: ISO-8601 timestamp stamped by the caller (kept out
of this pure function so it stays deterministic under test).
:param config: The run's knobs (iterations, requests, concurrency, runs,
mock_llm) for provenance.
:param harness: Harness driving full-turn journeys, e.g.
``"openai-agents"``.
:returns: The JSON-serializable report document.
"""
return {
"schema_version": SCHEMA_VERSION,
"generated_at": generated_at,
"git_sha": git_sha(),
"git_branch": git_branch(),
"host": host_info(),
"harness": harness,
"config": config,
"journeys": journey_results,
}
+228
View File
@@ -0,0 +1,228 @@
"""Deterministic corpus seeder for the performance benchmark.
The v1 harness booted an empty DB, so the read journeys measured a best-case
near-empty table. This seeds a sizeable, realistic corpus directly through the
store API (no HTTP, no runner) so ``list_sessions`` / ``get_session`` /
``load_conversation_history`` read a production-shaped volume.
Writes to the same DB URI the server later boots against; startup migrations
are an idempotent no-op on an at-head DB. The seed is deterministic (fixed RNG,
fixed counts) so the same config always yields the same corpus — which is what
makes "seed once, reuse" sound. The reuse marker records the Alembic head read
at seed time, so a corpus from an older schema is auto-reseeded (no manual
revision bookkeeping).
Listable-corpus recipe, per session (the permission grant is the gotcha — the
loopback server resolves every request to user ``"local"`` and
``list_sessions`` filters by it):
1. ``create_session_with_agent`` — conversation + session-scoped agent row.
2. ``permission_store.grant("local", sid, LEVEL_OWNER)`` — makes it listable.
3. one batched ``append(sid, items)`` — user-role message items.
Run standalone::
uv run --no-sync dev/benchmarks/omnigent/seed.py \
--database-uri sqlite:///tmp/bench.db --sessions 5000 --items-per-session 50
"""
from __future__ import annotations
import argparse
import random
import sys
from pathlib import Path
# Allow ``uv run <path>`` (no package context) to import omnigent + siblings.
sys.path.insert(0, str(Path(__file__).resolve().parents[3]))
from omnigent.db.utils import _get_head_db_revision, generate_agent_id
from omnigent.entities import MessageData, NewConversationItem
from omnigent.server.auth import LEVEL_OWNER, RESERVED_USER_LOCAL
from omnigent.stores.conversation_store.sqlalchemy_store import SqlAlchemyConversationStore
from omnigent.stores.permission_store.sqlalchemy_store import SqlAlchemyPermissionStore
# Label key stamped on the first seeded session recording the corpus config, so
# a later run can detect an existing (and matching) seed and skip re-seeding.
_SEED_META_LABEL = "omni_bench_seed"
# Fixed identifiers so the corpus is byte-stable across runs at a given config.
_AGENT_NAME = "bench-agent"
_DEFAULT_SESSIONS = 5000
_DEFAULT_ITEMS = 50
_DEFAULT_RNG_SEED = 1234
# A pool of realistic-ish message fragments; the RNG assembles item text from
# these so search_text has lexical variety without external data.
_FRAGMENTS = (
"investigate the failing migration",
"the runner keeps disconnecting under load",
"add pagination to the sessions endpoint",
"why does the policy classifier time out",
"refactor the conversation store append path",
"benchmark the list endpoints against postgres",
"the web UI drops the last streamed token",
"trace the tunnel handshake for this runner id",
"summarize the changes in this pull request",
"reproduce the elicitation race on reconnect",
)
def _meta_value(sessions: int, items_per_session: int, rng_seed: int, head: str) -> str:
"""Serialize the corpus config into the seed-marker label value.
Includes the Alembic *head* read at seed time, so a corpus seeded under an
older schema auto-mismatches the current head and is reseeded — no
hand-maintained revision constant.
"""
return f"sessions={sessions};items={items_per_session};rng={rng_seed};rev={head}"
def _existing_seed_meta(conv: SqlAlchemyConversationStore) -> str | None:
"""Return the seed-marker label value if a bench corpus already exists.
Looks up the most recent ``bench-agent`` session and reads its
``omni_bench_seed`` label. ``None`` means no (recognizable) seed present.
"""
listing = conv.list_conversations(limit=1, agent_name=_AGENT_NAME)
if not listing.data:
return None
marked = conv.get_conversation(listing.data[0].id)
return marked.labels.get(_SEED_META_LABEL) if marked is not None else None
def _make_items(rng: random.Random, count: int) -> list[NewConversationItem]:
"""Build *count* deterministic user-role message items.
User-role only: assistant messages require an ``agent`` field the store
only assigns after a real turn, and the seeded read path is role-agnostic.
"""
items: list[NewConversationItem] = []
for i in range(count):
text = f"{rng.choice(_FRAGMENTS)} (item {i})"
items.append(
NewConversationItem(
type="message",
response_id=f"resp_seed_{i}",
data=MessageData(role="user", content=[{"type": "input_text", "text": text}]),
)
)
return items
def seed(
db_uri: str,
*,
sessions: int = _DEFAULT_SESSIONS,
items_per_session: int = _DEFAULT_ITEMS,
rng_seed: int = _DEFAULT_RNG_SEED,
reseed: bool = False,
) -> int:
"""Seed *sessions* sessions × *items_per_session* items into *db_uri*.
Idempotent: if a matching seed already exists (same config + schema
revision) it is left untouched unless *reseed* is set. Constructing the
store runs migrations to head on first init, so *db_uri* need not
pre-exist.
:param db_uri: SQLAlchemy URI the server will also boot against, e.g.
``"sqlite:///abs/bench.db"`` or ``"postgresql+psycopg://…"``.
:param sessions: Number of listable sessions to create.
:param items_per_session: Conversation items appended to each session.
:param rng_seed: Seed for the deterministic text RNG.
:param reseed: Seed even when a matching corpus is already present.
:returns: The number of sessions created (0 when a matching seed is reused).
"""
conv = SqlAlchemyConversationStore(db_uri)
perms = SqlAlchemyPermissionStore(db_uri)
# Read the current schema head at runtime (no DB contacted) and fold it into
# the reuse marker, so a corpus from an older schema is auto-reseeded.
head = _get_head_db_revision("sqlite:///:memory:")
want = _meta_value(sessions, items_per_session, rng_seed, head)
if not reseed:
existing = _existing_seed_meta(conv)
if existing == want:
print(f"seed: matching corpus already present ({want}); skipping")
return 0
if existing is not None:
print(f"seed: existing corpus differs ({existing!r} != {want!r}); pass --reseed")
return 0
perms.ensure_user(RESERVED_USER_LOCAL)
rng = random.Random(rng_seed)
last_sid = ""
for s in range(sessions):
created = conv.create_session_with_agent(
agent_id=generate_agent_id(),
agent_name=_AGENT_NAME,
agent_bundle_location="bench/seed", # never validated on the read path
agent_description=None,
title=f"bench session {s}: {rng.choice(_FRAGMENTS)}",
)
sid = created.conversation.id
last_sid = sid
perms.grant(RESERVED_USER_LOCAL, sid, LEVEL_OWNER)
if items_per_session:
conv.append(sid, _make_items(rng, items_per_session))
if sessions >= 100 and s % (sessions // 10) == 0 and s:
print(f"seed: {s}/{sessions} sessions")
# Stamp the corpus config on the LAST (newest) session — that's the one
# ``_existing_seed_meta``'s default desc listing returns, so the reuse
# check finds it regardless of corpus size.
if last_sid:
conv.set_labels(last_sid, {_SEED_META_LABEL: want})
print(f"seed: created {sessions} sessions × {items_per_session} items ({want})")
return sessions
def _parse_args(argv: list[str]) -> argparse.Namespace:
parser = argparse.ArgumentParser(
prog="omnigent-benchmark-seed",
description=__doc__,
formatter_class=argparse.RawDescriptionHelpFormatter,
)
parser.add_argument(
"--database-uri",
metavar="URI",
help="DB to seed. Required unless --print-head.",
)
parser.add_argument("--sessions", type=int, default=_DEFAULT_SESSIONS, metavar="N")
parser.add_argument("--items-per-session", type=int, default=_DEFAULT_ITEMS, metavar="N")
parser.add_argument("--rng-seed", type=int, default=_DEFAULT_RNG_SEED, metavar="N")
parser.add_argument(
"--reseed",
action="store_true",
help="Seed even if a matching corpus is already present.",
)
parser.add_argument(
"--print-head",
action="store_true",
help="Print the repo's Alembic head revision and exit (drift-check helper).",
)
return parser.parse_args(argv)
def main(argv: list[str] | None = None) -> int:
args = _parse_args(argv if argv is not None else sys.argv[1:])
if args.print_head:
print(_get_head_db_revision("sqlite:///:memory:"))
return 0
if not args.database_uri:
print("seed: --database-uri is required (unless --print-head)", file=sys.stderr)
return 2
seed(
args.database_uri,
sessions=args.sessions,
items_per_session=args.items_per_session,
rng_seed=args.rng_seed,
reseed=args.reseed,
)
return 0
if __name__ == "__main__":
raise SystemExit(main())
+1
View File
@@ -0,0 +1 @@
"""Project-specific lint hooks run from .pre-commit-config.yaml."""
+262
View File
@@ -0,0 +1,262 @@
"""Flag patches that globally clobber ``asyncio`` module attributes.
What goes wrong
---------------
When test code writes::
with patch("omnigent.tools.mcp.asyncio.sleep", new_callable=AsyncMock):
...
``unittest.mock``'s dotted-path resolver walks:
1. ``import omnigent.tools.mcp``
2. ``getattr(omnigent.tools.mcp, "asyncio")`` -> returns the
``asyncio`` module singleton (because production code did
``import asyncio`` at module top).
3. ``setattr(asyncio_module, "sleep", mock)`` -> globally clobbers
``asyncio.sleep`` for the entire process for the patch's lifetime.
Under ``pytest-xdist``, concurrent tests in other workers that touch
``asyncio.sleep`` (directly or via any of the many helpers built on
top of it) silently get the mock, causing hard-to-debug cross-test
flakiness.
The same bug shape applies to:
- ``monkeypatch.setattr("...asyncio.sleep", ...)`` (pytest)
- ``patch.object(mod.asyncio, "sleep", ...)``
What this catches
-----------------
1. ``patch("...asyncio.<name>", ...)`` / ``mock.patch(...)`` where the
first positional arg is a string literal containing ``.asyncio.``
or ending in ``.asyncio``.
2. ``monkeypatch.setattr("...asyncio.<name>", ...)`` with the same
string-literal shape.
3. ``patch.object(<expr>, "<name>", ...)`` where ``<expr>`` is an
attribute access whose leaf attribute is ``asyncio``
(e.g. ``patch.object(mcp_mod.asyncio, "sleep", ...)``).
Correct shapes
--------------
Option A (preferred, when production has few sleep sites): add a
thin ``_sleep`` indirection in the production module and patch
``module._sleep`` from tests. The real ``asyncio.sleep`` is never
touched.
Option B (when A is too invasive): replace the ``asyncio`` binding
in the target module's namespace with a ``SimpleNamespace`` that
mirrors the real module but with the desired attribute swapped.
The real ``asyncio`` module is left alone.
Exit code 1 if any hits are found.
"""
from __future__ import annotations
import ast
import sys
from pathlib import Path
def _is_global_asyncio_literal(value: str) -> bool:
"""Return True if ``value`` is a one-segment getattr off the
stdlib ``asyncio`` module name (the shape that clobbers the
singleton).
The dangerous shape is::
patch("<pkg>.<mod>.asyncio.<leaf>", ...)
where ``<leaf>`` is exactly one segment. The dotted-path resolver
imports ``<pkg>.<mod>``, calls ``getattr(_, "asyncio")`` -- which
returns the stdlib ``asyncio`` module singleton because the
consumer did ``import asyncio`` -- then sets ``<leaf>`` on it.
NOT flagged:
* ``"<pkg>.asyncio.<sub>.<leaf>"`` -- here ``asyncio`` is a
subpackage name (e.g. ``websockets.asyncio.client``) and the
resolver descends *through* a package directory, not through
a getattr on the stdlib module.
* ``"<pkg>.asyncio"`` -- the bare module path is used to *replace*
the binding (the safe Option B pattern), not to clobber an
attribute inside it.
:param value: The dotted-path string passed to ``patch`` /
``monkeypatch.setattr``.
:returns: ``True`` if the path has exactly one segment after the
last ``.asyncio.``.
"""
marker = ".asyncio."
idx = value.rfind(marker)
if idx == -1:
return False
tail = value[idx + len(marker) :]
# Exactly one segment after `.asyncio.` -> getattr on the stdlib
# module singleton. Two or more segments -> the path descends
# through a subpackage named `asyncio`.
return bool(tail) and "." not in tail
def _is_patch_func(func: ast.expr) -> bool:
"""Return True if ``func`` spells ``patch`` or ``mock.patch``.
Plain ``patch.object`` is handled separately -- this only matches
the dotted-path form ``patch("...")`` / ``mock.patch("...")``.
"""
if isinstance(func, ast.Name) and func.id == "patch":
return True
if isinstance(func, ast.Attribute) and func.attr == "patch":
return True
return False
def _is_patch_object(func: ast.expr) -> bool:
"""Return True if ``func`` spells ``patch.object`` or ``mock.patch.object``."""
if not isinstance(func, ast.Attribute) or func.attr != "object":
return False
inner = func.value
if isinstance(inner, ast.Name) and inner.id == "patch":
return True
if isinstance(inner, ast.Attribute) and inner.attr == "patch":
return True
return False
def _is_monkeypatch_setattr(func: ast.expr) -> bool:
"""Return True if ``func`` spells ``<monkeypatch>.setattr``.
Conservative: matches any attribute access where the leaf is
``setattr`` and the receiver is a ``Name`` containing
``monkeypatch`` (case-insensitive). This covers the common
``monkeypatch.setattr(...)`` and ``m.setattr(...)`` (from
``with monkeypatch.context() as m:``).
:param func: The function-position expression of a ``Call``.
:returns: ``True`` if this looks like ``monkeypatch.setattr``.
"""
if not isinstance(func, ast.Attribute) or func.attr != "setattr":
return False
base = func.value
if isinstance(base, ast.Name):
return "monkeypatch" in base.id.lower() or base.id == "m"
return False
def _attribute_ends_in_asyncio(node: ast.expr) -> bool:
"""Return True if ``node`` is an attribute access ending in ``.asyncio``.
e.g. matches ``mcp_mod.asyncio`` and ``omnigent.tools.mcp.asyncio``.
"""
return isinstance(node, ast.Attribute) and node.attr == "asyncio"
def _check_string_literal_first_arg(node: ast.Call) -> bool:
"""Return True if ``node.args[0]`` is a string literal flagged by
:func:`_is_global_asyncio_literal`."""
if not node.args:
return False
first = node.args[0]
if not isinstance(first, ast.Constant) or not isinstance(first.value, str):
return False
return _is_global_asyncio_literal(first.value)
def scan(path: Path) -> list[tuple[int, str]]:
"""
Return ``[(lineno, snippet), ...]`` for each violation in ``path``.
:param path: File to scan.
:returns: One entry per flagged call site.
"""
try:
source = path.read_text()
tree = ast.parse(source)
except (SyntaxError, UnicodeDecodeError):
return []
source_lines = source.splitlines()
hits: list[tuple[int, str]] = []
for node in ast.walk(tree):
if not isinstance(node, ast.Call):
continue
if not _is_flagged_call(node):
continue
line = node.lineno
snippet = source_lines[line - 1].strip() if 0 < line <= len(source_lines) else ""
hits.append((line, snippet))
return hits
def _is_flagged_call(node: ast.Call) -> bool:
"""Return True if ``node`` is one of the banned patch shapes.
Covers three call shapes:
1. ``patch("...asyncio.<name>", ...)`` / ``mock.patch(...)``.
2. ``monkeypatch.setattr("...asyncio.<name>", ...)``.
3. ``patch.object(<expr ending in .asyncio>, "<name>", ...)``.
:param node: A ``Call`` node from the AST.
:returns: ``True`` if the call matches a banned shape.
"""
# patch("...asyncio.<name>", ...) / mock.patch("...asyncio.<name>", ...)
if _is_patch_func(node.func) and _check_string_literal_first_arg(node):
return True
# monkeypatch.setattr("...asyncio.<name>", ...)
if _is_monkeypatch_setattr(node.func) and _check_string_literal_first_arg(node):
return True
# patch.object(<...>.asyncio, "<name>", ...)
if (
_is_patch_object(node.func)
and len(node.args) >= 1
and _attribute_ends_in_asyncio(node.args[0])
):
return True
return False
def main(argv: list[str]) -> int:
"""
Scan every file in ``argv[1:]`` and report violations.
:param argv: Command-line args (``argv[0]`` is the script name).
:returns: Exit code -- 0 on clean scan, 1 if any hits.
"""
failed = False
for arg in argv[1:]:
path = Path(arg)
if not path.is_file():
continue
hits = scan(path)
if hits:
failed = True
for line, snippet in hits:
sys.stdout.write(
f"{path}:{line}: globally-clobbering asyncio patch detected: {snippet}\n"
)
sys.stdout.write(
" hint: see dev/lint/lint_no_global_asyncio_patch.py "
"docstring for correct shapes.\n"
)
if failed:
sys.stdout.write(
"\nPatching ``module.asyncio.sleep`` (or any other asyncio "
"attribute via a dotted path that walks through the asyncio "
"module) mutates the real asyncio module singleton, which "
"leaks into every other test running in the same process "
"(critical under pytest-xdist). Prefer adding a thin "
"``_sleep`` indirection in the production module and "
"patching that, or replace the module's ``asyncio`` binding "
"wholesale with a ``SimpleNamespace`` -- see the script "
"docstring for Option A / Option B examples.\n"
)
return 1
return 0
if __name__ == "__main__":
sys.exit(main(sys.argv))
+185
View File
@@ -0,0 +1,185 @@
"""Flag unconditional ``pytest.mark.skip`` on tests.
Per CLAUDE.md and the omnigent-testing skill, skipped tests are
invisible coverage loss — they silently rot and nobody remembers to
unskip them. If a test can't pass, rewrite it or delete it. Never
defer broken tests behind ``@pytest.mark.skip``.
What this catches
-----------------
- ``@pytest.mark.skip`` as a decorator (with or without args)
- ``pytestmark = pytest.mark.skip(...)`` as module-level skip
- ``pytest.skip(...)`` called as a statement at module scope
- ``pytest.mark.skip`` imported under an alias (e.g. ``from pytest
import mark; @mark.skip``) — detected conservatively when the
attribute chain resolves to ``skip``.
What this does NOT catch
------------------------
- ``@pytest.mark.skipif(condition, reason=...)``. Conditional skip
is legitimate when the condition tests an environmental
precondition (missing binary, wrong platform, missing optional
dep). Flagging it uniformly would produce too many false
positives; reviewers judge each ``skipif`` on whether the
condition is a real precondition or a hack to hide a broken
test.
Exit code 1 if any hits are found.
"""
from __future__ import annotations
import ast
import sys
from pathlib import Path
def _is_skip_attribute(node: ast.expr) -> bool:
"""Return True if ``node`` spells ``pytest.mark.skip`` (or ``mark.skip``).
Conservative: follows the ``.skip`` attribute access. Any chain
whose leaf attribute is ``skip`` and whose root resolves to
``pytest`` or ``mark`` is flagged.
"""
if not isinstance(node, ast.Attribute):
return False
if node.attr != "skip":
return False
# Walk back up the chain.
base = node.value
while isinstance(base, ast.Attribute):
if base.attr == "mark":
return True
base = base.value
if isinstance(base, ast.Name) and base.id in ("pytest", "mark"):
return True
return False
def _decorator_calls_skip(decorator: ast.expr) -> bool:
"""Return True if ``decorator`` is ``@pytest.mark.skip`` or ``@...skip(...)``."""
if _is_skip_attribute(decorator):
return True
if isinstance(decorator, ast.Call) and _is_skip_attribute(decorator.func):
return True
return False
def _is_module_level_skip_assignment(node: ast.stmt) -> bool:
"""Return True if ``node`` is ``pytestmark = pytest.mark.skip(...)``.
Assigning to ``pytestmark`` at module scope applies the mark to
every test collected from the file — the nuclear option for
skipping.
"""
if not isinstance(node, ast.Assign):
return False
if not any(isinstance(t, ast.Name) and t.id == "pytestmark" for t in node.targets):
return False
return _decorator_calls_skip(node.value)
def _is_module_level_skip_call(node: ast.stmt) -> bool:
"""Return True if ``node`` is a bare ``pytest.skip(...)`` call.
A ``pytest.skip()`` call at module scope during collection skips
the entire module. (Inside a test body it's a conditional skip
expression — not what we flag; we only care about module-level
or top-of-class uses.)
"""
if not isinstance(node, ast.Expr):
return False
call = node.value
if not isinstance(call, ast.Call):
return False
func = call.func
if isinstance(func, ast.Attribute) and func.attr == "skip":
base = func.value
if isinstance(base, ast.Name) and base.id == "pytest":
return True
return False
def scan(path: Path) -> list[tuple[int, str]]:
"""
Return ``[(lineno, message), ...]`` for each skip site in ``path``.
:param path: File to scan.
:returns: Hits with line numbers and descriptive messages.
"""
try:
source = path.read_text()
tree = ast.parse(source)
except (SyntaxError, UnicodeDecodeError):
return []
hits: list[tuple[int, str]] = []
# Decorator-level skips on test functions / classes.
for node in ast.walk(tree):
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
for decorator in node.decorator_list:
if _decorator_calls_skip(decorator):
hits.append(
(
decorator.lineno,
f"`@pytest.mark.skip` on `{node.name}`: skipped tests "
"rot invisibly; rewrite or delete",
),
)
# Module-level `pytestmark = pytest.mark.skip(...)` and bare
# `pytest.skip(...)` calls at module scope.
for stmt in tree.body:
if _is_module_level_skip_assignment(stmt):
hits.append(
(
stmt.lineno,
"`pytestmark = pytest.mark.skip(...)` at module scope skips "
"every test in the file",
),
)
elif _is_module_level_skip_call(stmt):
hits.append(
(
stmt.lineno,
"`pytest.skip(...)` at module scope skips the entire module during collection",
),
)
return hits
def main(argv: list[str]) -> int:
"""
Scan every file in ``argv[1:]`` and report unconditional skips.
:param argv: Command-line args (``argv[0]`` is the script name).
:returns: Exit code — 0 on clean scan, 1 if any hits.
"""
failed = False
for arg in argv[1:]:
path = Path(arg)
if not path.is_file():
continue
hits = scan(path)
if hits:
failed = True
for line, msg in hits:
sys.stdout.write(f"{path}:{line}: {msg}\n")
if failed:
sys.stdout.write(
"\nSkipped tests are invisible coverage loss. If a test can't "
"pass, rewrite it for the current architecture or delete it. "
"For genuine environmental gates (missing binary, platform), "
"use ``@pytest.mark.skipif`` with a clear reason — this rule "
"only flags unconditional ``skip``.\n"
)
return 1
return 0
if __name__ == "__main__":
sys.exit(main(sys.argv))
+1169
View File
File diff suppressed because it is too large Load Diff
+34
View File
@@ -0,0 +1,34 @@
[package]
name = "omnidev"
version = "0.1.0"
edition = "2021"
description = "Per-repo dev pod supervisor TUI for the Omnigent repo"
publish = false
[[bin]]
name = "omnidev"
path = "src/main.rs"
[dependencies]
anyhow = "1"
clap = { version = "4", features = ["derive"] }
crossterm = "0.28"
ratatui = "0.29"
ansi-to-tui = "7"
notify = "8"
notify-debouncer-full = "0.5"
libc = "0.2"
serde = { version = "1", features = ["derive"] }
serde_json = "1"
toml = "0.8"
tokio = { version = "1", features = [
"rt-multi-thread",
"macros",
"process",
"io-util",
"net",
"time",
"sync",
"signal",
] }
if-addrs = "0.15"
+177
View File
@@ -0,0 +1,177 @@
# omnidev
Dev tooling for Omnigent, in one binary with two independent capabilities:
1. A per-repo dev **pod supervisor** (bare `omnidev`) — the default.
2. **Install management** (`omnidev install`/`update`/`check`) — install and
keep a git-based omnigent up to date. See
[Managing your omnigent install](#managing-your-omnigent-install). These
subcommands need no checkout and run anywhere.
## Pod supervisor
A per-repo dev **pod** supervisor, as a single long-running terminal UI. It
replaces the three-terminal local dev flow (`omnigent server`, `omnigent host`,
`npm run dev`) with one process that:
- runs each checkout in an **isolated pod** — its own state dir, database,
artifacts, logs, and auto-allocated ports — so multiple worktrees never
collide;
- **supervises** the backend server, the host daemon, and the Vite frontend,
restarting any that crash (with backoff);
- **reloads the backend** (server → host) when you edit `omnigent/**/*.py`;
the frontend self-reloads through Vite HMR;
- gives you **scrollable per-process log panes** plus a combined view.
## Build & run
Requires the repo's usual dev prerequisites (`uv` for Python, `npm` for the
web UI) plus a Rust toolchain.
```bash
cd dev/omnidev
cargo run # launches the TUI for the surrounding checkout
```
Run it from anywhere inside the checkout — it walks up to the repo root
(the `.jj`/`.git` marker) and requires `omnigent/` and
`web/` to be present. Build a release binary with `cargo build --release`
(lands at `target/release/omnidev`).
## What it starts
| Process | Command | Notes |
|---|---|---|
| server | `uv run omnigent server --host 127.0.0.1 --port <p> --database-uri … --artifact-location …` | Waited on via `GET /health`. |
| host | `uv run omnigent host --server http://127.0.0.1:<p>` | Started once the server is healthy. |
| vite | `npm run dev -- --host <host> --port <p> --strictPort` (cwd `web/`) | `OMNIGENT_URL` points its proxy at the pod's server. |
Before Vite starts (and on a manual Vite restart), omnidev runs `npm install`
in `web/` when needed — `node_modules/` is missing, or `package.json` /
`package-lock.json` is newer than it — so a fresh checkout or a new dependency
doesn't make Vite fail its dependency scan. Output streams into the `vite` pane.
Open the UI at the `ui` URL shown in the header (the Vite dev server).
## Isolation
Only Omnigent's own state is isolated per pod — enough that concurrent pods
never share a database, server pidfile, or `config.yaml` — via
`OMNIGENT_DATA_DIR`, `OMNIGENT_DATABASE_URI`, `OMNIGENT_URL`, and
`OMNIGENT_CONFIG_HOME`. Everything else (your real `HOME`, credentials, and
uv/npm caches) is inherited, because the agents Omnigent runs need it. This is
deliberately lighter than the hermetic `scripts/backend-smoke.sh` sandbox,
which repoints `HOME`/`XDG_*` to touch nothing real.
Each pod gets its own `config.yaml` under `<pod>/config/`, pointed to by
`OMNIGENT_CONFIG_HOME`. On first create it's **seeded** from your real
`~/.omnigent/config.yaml` (if present) so the pod works out of the box — it
keeps your providers — after which the two are independent: server-config edits
inside a pod (via the UI or `omnigent config`) don't touch your real config.
`--clean` wipes the pod dir, so the next run re-seeds from your real config.
The pod dir defaults to
`${XDG_CACHE_HOME:-~/.cache}/omnidev/<repo-name>-<hash>/`, keyed to the
canonical checkout path. Per-process logs are written through to
`<pod>/logs/{server,host,vite}.log` for inspection outside the TUI.
## Options
```
--server-port <N> Force the backend port (default: probe from 6767)
--vite-port <N> Force the Vite port (default: probe from 5173)
--vite-host <ADDR> Vite bind host (default: 127.0.0.1; use 0.0.0.0 for LAN access)
--trust-lan-origins Trust this machine's LAN origins (for device testing)
--pod-dir <PATH> Use a specific pod dir instead of the per-repo default
--no-vite Backend + host only (no frontend)
--clean Wipe the pod dir before starting
```
`--vite-host 0.0.0.0` exposes the Vite dev server on all interfaces for device
testing. Vite still proxies API traffic to the pod backend through `127.0.0.1`.
### Testing from a phone or tablet
`--vite-host 0.0.0.0` alone lets a device load the UI, but the backend runs in
single-user local mode, where its CSRF/CSWSH guard trusts only loopback
origins. A device loads the UI at `http://<your-lan-ip>:<vite-port>`, so its
browser stamps that non-loopback origin on every request — and the guard then
rejects multipart uploads (403) and refuses the live WebSocket stream.
`--trust-lan-origins` fixes that: omnidev enumerates this machine's LAN IPv4
addresses and trusts the matching `http://<ip>:<vite-port>` origins via the
server's `OMNIGENT_WS_ALLOWED_ORIGINS` allowlist (merged with any value you
already export). It stays exact-match — only those origins are trusted, nothing
is disabled — so it's for dev pods, not deployed servers. The trusted origins
are printed in the combined log at startup.
```bash
omnidev --vite-host 0.0.0.0 --trust-lan-origins
```
This covers IPv4 LAN addresses; mDNS `.local` hostnames and HTTPS origins are
not auto-trusted (add those to `OMNIGENT_WS_ALLOWED_ORIGINS` yourself).
## Keys
| Key | Action |
|---|---|
| `1` / `2` / `3` / `0` | Focus server / host / vite / combined pane |
| `Tab` | Cycle panes |
| `↑` `↓` `PgUp` `PgDn` | Scroll (detaches from tail) |
| `f` | Toggle follow-tail |
| `r` | Restart the focused process (server/host restart as a pair) |
| `R` | Restart the backend (server then host) |
| `c` | Clear the focused pane |
| `q` / `Ctrl-C` | Quit and tear down all processes |
## Managing your omnigent install
For people who *run* omnigent (installed from git via `uv tool install`) rather
than develop it. This wraps the fiddly PEP 508 install syntax and adds a daily
update check — filling a gap, since omnigent's own update notice only works for
PyPI-wheel installs and skips git installs.
These subcommands manage the global tool and work from **any directory** (no
checkout needed).
```
omnidev install # uv tool install omnigent from git (databricks extra, main)
omnidev update # reinstall the latest of the tracked ref/extras
omnidev check # check for an update; prompt to update on a TTY
omnidev refresh # refresh the check cache from the network (usually detached)
omnidev shell-hook # print the daily-check snippet for your shell rc
```
`install` options: `--ref <branch/tag/sha>` (default `main`), `--extra <name>`
(repeatable; defaults to `databricks`), `--no-default-extra` (install with no
extras), `--repo <url>`. The choice is saved to
`${XDG_CONFIG_HOME:-~/.config}/omnidev/install.toml` so `update` reuses it.
Installing from git **builds the web UI from source**, so Node 22+/npm must be
on PATH (the PyPI wheel ships the UI prebuilt; the git install does not).
`omnidev install` fails early with a clear message if `uv` or `npm` is missing.
### Daily update check
Append the hook to your shell rc once to be told, at most once a day, when a
newer `main` commit is available — and be offered to update on the spot:
```bash
omnidev shell-hook >> ~/.zshrc # or ~/.bashrc
```
The snippet itself guards on `command -v omnidev`, so it's a no-op in shells
where omnidev isn't on PATH — nothing to fail. (Appending the snippet is
preferred over `eval "$(omnidev shell-hook)"`: the latter would run omnidev on
every shell startup and print a "command not found" error whenever omnidev is
absent.)
On each interactive shell it runs `omnidev check --quiet`, which reads a cached
result (`${XDG_CACHE_HOME:-~/.cache}/omnidev/omnigent-check.json`) and, when
stale (>24h), refreshes it in a detached background process — so shell startup
never blocks on the network. When a newer commit is available it prints a notice
and, on a terminal, prompts `Update omnigent now? [y/N]`; on yes it runs
`omnidev update` in the foreground. Declining suppresses that same commit until a
newer one lands. Set `OMNIGENT_NO_UPDATE_CHECK` in your environment if you want
to silence omnigent's own separate notice.
+209
View File
@@ -0,0 +1,209 @@
//! Manage the user's git-based omnigent installation via `uv tool install`.
//!
//! None of this needs a local checkout: it drives `uv` and reads the installed
//! tool's metadata, and any git call targets the remote.
use std::path::PathBuf;
use std::process::Command;
use anyhow::{bail, Context, Result};
use serde::{Deserialize, Serialize};
use crate::paths;
pub const DEFAULT_REPO: &str = "https://github.com/omnigent-ai/omnigent.git";
pub const DEFAULT_REF: &str = "main";
pub const DEFAULT_EXTRA: &str = "databricks";
const PYTHON_VERSION: &str = "3.12";
/// Durable record of how the user wants omnigent installed. Persisted so
/// `update` reinstalls the same repo/ref/extras without re-specifying them.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct InstallConfig {
pub repo: String,
#[serde(rename = "ref")]
pub git_ref: String,
pub extras: Vec<String>,
}
impl Default for InstallConfig {
fn default() -> Self {
InstallConfig {
repo: DEFAULT_REPO.to_string(),
git_ref: DEFAULT_REF.to_string(),
extras: vec![DEFAULT_EXTRA.to_string()],
}
}
}
impl InstallConfig {
pub fn load() -> Result<Option<InstallConfig>> {
let path = paths::install_config_path()?;
match std::fs::read_to_string(&path) {
Ok(text) => Ok(Some(
toml::from_str(&text).with_context(|| format!("parsing {}", path.display()))?,
)),
Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(None),
Err(e) => Err(e).with_context(|| format!("reading {}", path.display())),
}
}
pub fn save(&self) -> Result<()> {
let path = paths::install_config_path()?;
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)
.with_context(|| format!("creating {}", parent.display()))?;
}
let text = toml::to_string(self).context("serializing install config")?;
std::fs::write(&path, text).with_context(|| format!("writing {}", path.display()))?;
Ok(())
}
/// The PEP 508 install spec, e.g.
/// `omnigent[databricks] @ git+https://github.com/omnigent-ai/omnigent.git@main`.
/// With no extras it collapses to the bare `git+<repo>@<ref>` URL.
pub fn spec(&self) -> String {
let source = format!("git+{}@{}", self.repo, self.git_ref);
if self.extras.is_empty() {
source
} else {
format!("omnigent[{}] @ {}", self.extras.join(","), source)
}
}
}
/// Fail early with a clear message if the toolchain a git install needs is
/// missing. Installing from git builds the web UI from source (Node/npm),
/// unlike the PyPI wheel which ships it prebuilt.
fn preflight() -> Result<()> {
if which("uv").is_none() {
bail!("`uv` is not on PATH. Install it first: https://docs.astral.sh/uv/");
}
if which("npm").is_none() {
bail!(
"`npm` is not on PATH. Installing omnigent from git builds the web UI \
from source and needs Node 22+/npm. Install Node, then retry."
);
}
Ok(())
}
/// Install omnigent from git per `config`. `reinstall` forces uv past its cache
/// so a moving ref (e.g. `main`) actually re-resolves.
pub fn run_uv_install(config: &InstallConfig, reinstall: bool) -> Result<()> {
preflight()?;
let spec = config.spec();
let mut cmd = Command::new("uv");
cmd.args(["tool", "install", "--force", "--python", PYTHON_VERSION]);
if reinstall {
cmd.arg("--reinstall");
}
cmd.arg(&spec);
eprintln!("omnidev: uv tool install {spec}");
let status = cmd
.status()
.context("running `uv tool install` (is uv installed?)")?;
if !status.success() {
bail!("`uv tool install` failed ({status})");
}
Ok(())
}
/// `install` subcommand: persist intent, install, then record the resolved sha.
pub fn install(config: &InstallConfig) -> Result<()> {
config.save()?;
run_uv_install(config, false)?;
record_installed_sha(config);
println!("omnidev: installed omnigent ({})", config.spec());
Ok(())
}
/// `update` subcommand: reinstall the latest of the persisted ref/extras. Falls
/// back to defaults when no config has been written yet.
pub fn update() -> Result<()> {
let config = InstallConfig::load()?.unwrap_or_default();
config.save()?;
run_uv_install(&config, true)?;
record_installed_sha(&config);
println!("omnidev: updated omnigent ({})", config.spec());
Ok(())
}
/// After a successful install, capture the remote sha of the tracked ref and
/// stash it in the cache so `check` has a baseline even before the dist-info
/// reader runs. Best-effort — failures here never fail the install.
fn record_installed_sha(config: &InstallConfig) {
if let Some(sha) = crate::update_check::remote_sha(&config.repo, &config.git_ref) {
let _ = crate::update_check::set_installed_sha(&sha);
}
}
/// Read the commit the installed omnigent tool was built from, via its PEP 610
/// `direct_url.json`. Returns `None` for a non-VCS install or when uv/metadata
/// can't be read. Never touches the working directory.
pub fn installed_commit() -> Option<String> {
let dir = uv_tool_dir()?;
// …/omnigent/**/omnigent-*.dist-info/direct_url.json
let omnigent_root = dir.join("omnigent");
let dist_info = find_dist_info(&omnigent_root)?;
let text = std::fs::read_to_string(dist_info.join("direct_url.json")).ok()?;
let value: serde_json::Value = serde_json::from_str(&text).ok()?;
value
.get("vcs_info")?
.get("commit_id")?
.as_str()
.map(str::to_string)
}
fn uv_tool_dir() -> Option<PathBuf> {
let output = Command::new("uv").args(["tool", "dir"]).output().ok()?;
if !output.status.success() {
return None;
}
let path = String::from_utf8(output.stdout).ok()?;
let trimmed = path.trim();
if trimmed.is_empty() {
None
} else {
Some(PathBuf::from(trimmed))
}
}
/// Find the `omnigent-*.dist-info` dir under a uv tool's environment. uv lays
/// tools out as `<tool>/lib/pythonX.Y/site-packages/<pkg>-<ver>.dist-info`, so
/// we walk rather than hardcode the python version.
fn find_dist_info(root: &std::path::Path) -> Option<PathBuf> {
let mut stack = vec![root.to_path_buf()];
while let Some(dir) = stack.pop() {
let Ok(entries) = std::fs::read_dir(&dir) else {
continue;
};
for entry in entries.flatten() {
let path = entry.path();
if !path.is_dir() {
continue;
}
let name = entry.file_name();
let name = name.to_string_lossy();
if name.starts_with("omnigent-") && name.ends_with(".dist-info") {
return Some(path);
}
stack.push(path);
}
}
None
}
/// Locate an executable on PATH (portable `which`, no external dep).
fn which(program: &str) -> Option<PathBuf> {
let path = std::env::var_os("PATH")?;
for dir in std::env::split_paths(&path) {
let candidate = dir.join(program);
if candidate.is_file() {
return Some(candidate);
}
}
None
}
+114
View File
@@ -0,0 +1,114 @@
//! LAN origin discovery for device testing.
//!
//! When Vite binds to `0.0.0.0` (`--vite-host 0.0.0.0`), a phone or tablet on
//! the same network loads the UI at `http://<lan-ip>:<vite-port>`. Its browser
//! stamps that non-loopback address as the `Origin` on every request. The
//! backend runs in local single-user mode, where the origin guard
//! (`omnigent.server.ws_origin.origin_allowed`) admits only loopback origins —
//! so multipart uploads get a 403 and the WebSocket stream is refused.
//!
//! `--trust-lan-origins` closes that gap by enumerating this machine's LAN
//! IPv4 addresses and handing the server the matching `http://<ip>:<port>`
//! origins via `OMNIGENT_WS_ALLOWED_ORIGINS` — the server's own exact-match
//! allowlist. It stays exact-match (no security disable): only the origins we
//! name are trusted.
use std::net::Ipv4Addr;
/// Whether an IPv4 address is a usable LAN address to trust as an origin.
///
/// Keeps private (RFC 1918) and link-local (169.254/16) addresses — the ones a
/// device on the same network actually reaches this machine by. Drops loopback
/// (already trusted), unspecified (`0.0.0.0`), broadcast, documentation, and
/// multicast, none of which a real device browses to.
fn is_lan_ipv4(ip: &Ipv4Addr) -> bool {
(ip.is_private() || ip.is_link_local())
&& !ip.is_loopback()
&& !ip.is_unspecified()
&& !ip.is_broadcast()
&& !ip.is_multicast()
}
/// Build the `http://<ip>:<port>` origins to trust for a given set of LAN
/// IPv4 addresses.
///
/// Split out from interface enumeration so the origin-shaping (which is all we
/// assert on) is testable without touching the host's real interfaces. The
/// input is deduplicated and the output is sorted for a stable env value.
fn origins_for_ips(ips: impl IntoIterator<Item = Ipv4Addr>, vite_port: u16) -> Vec<String> {
let mut origins: Vec<String> = ips
.into_iter()
.filter(is_lan_ipv4)
.map(|ip| format!("http://{ip}:{vite_port}"))
.collect();
origins.sort();
origins.dedup();
origins
}
/// Discover the `http://<lan-ip>:<vite-port>` origins for this machine's LAN
/// interfaces.
///
/// Returns an empty vector when no LAN interface is found (e.g. offline) — the
/// caller then simply trusts nothing extra rather than failing. Interface
/// enumeration errors are treated the same way: LAN trust is a convenience, so
/// a lookup failure must not block the pod from starting.
pub fn trusted_lan_origins(vite_port: u16) -> Vec<String> {
let ips = match if_addrs::get_if_addrs() {
Ok(ifaces) => ifaces
.into_iter()
.filter_map(|iface| match iface.addr.ip() {
std::net::IpAddr::V4(v4) => Some(v4),
std::net::IpAddr::V6(_) => None,
}),
Err(_) => return Vec::new(),
};
origins_for_ips(ips, vite_port)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn keeps_private_and_link_local_drops_loopback_and_public() {
assert!(is_lan_ipv4(&Ipv4Addr::new(192, 168, 1, 42)));
assert!(is_lan_ipv4(&Ipv4Addr::new(10, 0, 0, 5)));
assert!(is_lan_ipv4(&Ipv4Addr::new(172, 16, 3, 9)));
assert!(is_lan_ipv4(&Ipv4Addr::new(169, 254, 10, 1)));
assert!(!is_lan_ipv4(&Ipv4Addr::new(127, 0, 0, 1)));
assert!(!is_lan_ipv4(&Ipv4Addr::new(0, 0, 0, 0)));
assert!(!is_lan_ipv4(&Ipv4Addr::new(8, 8, 8, 8)));
assert!(!is_lan_ipv4(&Ipv4Addr::new(255, 255, 255, 255)));
}
#[test]
fn builds_http_origins_with_the_vite_port() {
let origins = origins_for_ips([Ipv4Addr::new(192, 168, 1, 42)], 5173);
assert_eq!(origins, vec!["http://192.168.1.42:5173"]);
}
#[test]
fn filters_and_sorts_and_dedups() {
let origins = origins_for_ips(
[
Ipv4Addr::new(10, 0, 0, 9),
Ipv4Addr::new(127, 0, 0, 1), // loopback dropped
Ipv4Addr::new(8, 8, 8, 8), // public dropped
Ipv4Addr::new(192, 168, 1, 5),
Ipv4Addr::new(10, 0, 0, 9), // duplicate collapsed
],
8080,
);
assert_eq!(
origins,
vec!["http://10.0.0.9:8080", "http://192.168.1.5:8080"]
);
}
#[test]
fn no_lan_interfaces_yields_no_origins() {
assert!(origins_for_ips([Ipv4Addr::new(127, 0, 0, 1)], 5173).is_empty());
}
}
+46
View File
@@ -0,0 +1,46 @@
//! Single-instance guard per pod.
//!
//! Two omnidev runs in the same checkout resolve to the same pod dir (the dir
//! is keyed to the canonical repo root), so their processes would fight over
//! the same ports and state. An advisory `flock` on a file in the pod dir lets
//! only the first in. The lock is held for the process lifetime and released
//! by the OS on exit or crash — no stale-file cleanup needed.
use std::fs::{File, OpenOptions};
use std::os::fd::AsRawFd;
use std::path::Path;
use anyhow::{bail, Context, Result};
/// An acquired pod lock. Dropping it (on process exit) releases the flock.
pub struct PodLock {
_file: File,
}
/// Try to take the pod's exclusive lock. Returns an error naming the pod dir if
/// another omnidev already holds it.
pub fn acquire(pod_dir: &Path) -> Result<PodLock> {
let path = pod_dir.join("omnidev.lock");
let file = OpenOptions::new()
.create(true)
.write(true)
.truncate(false)
.open(&path)
.with_context(|| format!("opening lock file {}", path.display()))?;
// Non-blocking exclusive lock: EWOULDBLOCK means a peer holds it.
let rc = unsafe { libc::flock(file.as_raw_fd(), libc::LOCK_EX | libc::LOCK_NB) };
if rc != 0 {
let err = std::io::Error::last_os_error();
if err.raw_os_error() == Some(libc::EWOULDBLOCK) {
bail!(
"another omnidev is already running for this checkout (pod {}). \
Quit it first, or run in a different worktree.",
pod_dir.display()
);
}
return Err(err).with_context(|| format!("locking {}", path.display()));
}
Ok(PodLock { _file: file })
}
+58
View File
@@ -0,0 +1,58 @@
//! Per-process bounded log buffers with write-through to disk.
use std::collections::VecDeque;
use std::fs::{File, OpenOptions};
use std::io::Write;
use std::path::Path;
const MAX_LINES: usize = 5000;
/// A bounded ring buffer of log lines for one channel, mirrored to a file so
/// the full session output survives for later inspection (`tail`, editor).
pub struct LogBuffer {
lines: VecDeque<String>,
file: Option<File>,
/// Monotonic count of lines ever appended — lets panes detect growth for
/// follow-tail without diffing the buffer.
pub total: u64,
}
impl LogBuffer {
pub fn new(path: &Path) -> Self {
let file = OpenOptions::new().create(true).append(true).open(path).ok();
LogBuffer {
lines: VecDeque::with_capacity(MAX_LINES),
file,
total: 0,
}
}
/// In-memory only channel (e.g. the synthetic "omnidev" event log).
pub fn memory() -> Self {
LogBuffer {
lines: VecDeque::with_capacity(256),
file: None,
total: 0,
}
}
pub fn push(&mut self, line: impl Into<String>) {
let line = line.into();
if let Some(f) = self.file.as_mut() {
let _ = writeln!(f, "{line}");
}
if self.lines.len() == MAX_LINES {
self.lines.pop_front();
}
self.lines.push_back(line);
self.total = self.total.saturating_add(1);
}
pub fn clear(&mut self) {
self.lines.clear();
}
pub fn iter(&self) -> impl Iterator<Item = &String> {
self.lines.iter()
}
}
+211
View File
@@ -0,0 +1,211 @@
//! omnidev — dev tooling for Omnigent.
//!
//! Two independent capabilities in one binary:
//! - **pod supervisor** (bare `omnidev`): manages an isolated dev instance for
//! the current checkout — server/host/vite, restarting the backend on Python
//! changes while Vite handles frontend HMR.
//! - **install management** (`omnidev install`/`update`/`check`/…): install and
//! keep a git-based omnigent up to date. These need no checkout and run
//! anywhere.
mod install;
mod lan;
mod lock;
mod logs;
mod paths;
mod pod;
mod ports;
mod process;
mod shellhook;
mod state;
mod supervisor;
mod tui;
mod update_check;
mod watcher;
use std::path::PathBuf;
use std::sync::Arc;
use anyhow::Result;
use clap::{Parser, Subcommand};
use tokio::sync::mpsc;
use install::InstallConfig;
use pod::Pod;
use ports::Ports;
use state::Shared;
use supervisor::{Cmd, Supervisor};
#[derive(Parser, Debug)]
#[command(name = "omnidev", about = "Dev tooling for Omnigent", version)]
struct Args {
#[command(subcommand)]
command: Option<Command>,
#[command(flatten)]
run: RunArgs,
}
/// Flags for the default (no-subcommand) pod-supervisor run.
#[derive(clap::Args, Debug)]
struct RunArgs {
/// Force the backend server port (default: probe from 6767).
#[arg(long)]
server_port: Option<u16>,
/// Force the Vite dev-server port (default: probe from 5173).
#[arg(long)]
vite_port: Option<u16>,
/// Vite dev-server bind host (default: 127.0.0.1; use 0.0.0.0 for LAN access).
#[arg(long, default_value = "127.0.0.1")]
vite_host: String,
/// Trust this machine's LAN origins so a phone/tablet on the same network
/// can use the UI (uploads + live stream). Pairs with `--vite-host 0.0.0.0`.
#[arg(long)]
trust_lan_origins: bool,
/// Use this pod directory instead of the per-repo default.
#[arg(long)]
pod_dir: Option<PathBuf>,
/// Do not start the Vite frontend (backend + host only).
#[arg(long)]
no_vite: bool,
/// Wipe the pod directory before starting.
#[arg(long)]
clean: bool,
}
#[derive(Subcommand, Debug)]
enum Command {
/// Install omnigent from git (defaults to the databricks extra, main).
Install {
/// Git ref (branch/tag/sha) to track.
#[arg(long, default_value = install::DEFAULT_REF)]
r#ref: String,
/// Extra to include (repeatable). Defaults to `databricks`.
#[arg(long = "extra")]
extras: Vec<String>,
/// Omit the default databricks extra (install with no extras).
#[arg(long)]
no_default_extra: bool,
/// Git repo URL.
#[arg(long, default_value = install::DEFAULT_REPO)]
repo: String,
},
/// Reinstall the latest of the tracked ref/extras.
Update,
/// Check for an omnigent update (the shell hook calls this).
Check {
/// Print nothing when already up to date.
#[arg(long)]
quiet: bool,
},
/// Refresh the update-check cache from the network (usually run detached).
Refresh,
/// Print a shell snippet to eval from .zshrc/.bashrc for daily checks.
ShellHook,
}
fn main() -> Result<()> {
let args = Args::parse();
// Install-management subcommands manage a global tool and must work from
// anywhere — dispatch them before any checkout discovery.
match args.command {
Some(Command::Install {
r#ref,
extras,
no_default_extra,
repo,
}) => {
let extras = if !extras.is_empty() {
extras
} else if no_default_extra {
vec![]
} else {
vec![install::DEFAULT_EXTRA.to_string()]
};
let config = InstallConfig {
repo,
git_ref: r#ref,
extras,
};
install::install(&config)
}
Some(Command::Update) => install::update(),
Some(Command::Check { quiet }) => update_check::check(quiet),
Some(Command::Refresh) => update_check::refresh(),
Some(Command::ShellHook) => {
shellhook::print();
Ok(())
}
None => run_supervisor(args.run),
}
}
/// Default path: the pod supervisor for the current checkout. This is the only
/// path that requires an Omnigent checkout.
#[tokio::main]
async fn run_supervisor(args: RunArgs) -> Result<()> {
let cwd = std::env::current_dir()?;
let repo_root = paths::find_repo_root(&cwd)?;
let pod_dir = match &args.pod_dir {
Some(p) => p.clone(),
None => paths::default_pod_dir(&repo_root)?,
};
if args.clean {
pod::clean(&pod_dir)?;
}
std::fs::create_dir_all(&pod_dir)?;
// Only one omnidev per pod — same-checkout runs share this dir and would
// otherwise fight over ports and state. Held until the process exits.
let _lock = lock::acquire(&pod_dir)?;
let ports = Ports::resolve(&pod_dir, args.server_port, args.vite_port)?;
// LAN origins are keyed to the resolved Vite port, so compute them here
// once the port is known. Empty unless `--trust-lan-origins` is set.
let trusted_origins = if args.trust_lan_origins {
lan::trusted_lan_origins(ports.vite)
} else {
Vec::new()
};
let pod = Arc::new(Pod::create(
repo_root,
pod_dir,
ports,
args.vite_host,
trusted_origins,
)?);
let shared = Shared::new(&pod);
let (cmd_tx, cmd_rx) = mpsc::unbounded_channel::<Cmd>();
// File watcher: Python changes -> Reload commands. Keep the debouncer alive
// for the whole session.
let _watcher = watcher::spawn(&pod.omnigent_dir(), cmd_tx.clone())?;
// Supervisor runs on the tokio runtime; the TUI drives it via cmd_tx.
let supervisor = Supervisor::new(
pod.clone(),
shared.clone(),
!args.no_vite,
args.trust_lan_origins,
);
let sup_handle = tokio::spawn(supervisor.run(cmd_rx));
// Run the TUI (owns the terminal) until the user quits.
let app = tui::App::new(pod.clone(), shared.clone(), cmd_tx.clone());
let result = app.run().await;
// Tear down children, then wait for the supervisor to finish shutdown.
let _ = cmd_tx.send(Cmd::Shutdown);
let _ = sup_handle.await;
result
}
+93
View File
@@ -0,0 +1,93 @@
//! Repo-root discovery and per-repo pod-directory resolution.
use std::path::{Path, PathBuf};
use anyhow::{bail, Context, Result};
/// Walk up from `start` looking for the checkout root.
///
/// The root is the first ancestor holding a `.jj/` or `.git/` marker — the VCS
/// root. We then require `web/` and `omnigent/` to be present so we fail early
/// on an unrelated repo rather than mid-spawn.
pub fn find_repo_root(start: &Path) -> Result<PathBuf> {
let start = start
.canonicalize()
.with_context(|| format!("resolving start dir {}", start.display()))?;
let mut cur: Option<&Path> = Some(&start);
while let Some(dir) = cur {
if dir.join(".jj").is_dir() || dir.join(".git").exists() {
let root = dir.to_path_buf();
if !root.join("omnigent").is_dir() || !root.join("web").is_dir() {
bail!(
"found a VCS root at {} but it lacks omnigent/ and web/ — \
run omnidev from inside an Omnigent checkout",
root.display()
);
}
return Ok(root);
}
cur = dir.parent();
}
bail!(
"could not find a checkout root above {} (no .jj or .git marker)",
start.display()
)
}
/// Stable per-repo pod directory: `${XDG_CACHE_HOME:-~/.cache}/omnidev/<slug>-<hash8>/`.
///
/// The hash of the canonical repo path keeps two worktrees on distinct pods;
/// the slug (repo basename) keeps the path human-readable.
pub fn default_pod_dir(repo_root: &Path) -> Result<PathBuf> {
let cache = cache_home()?;
let slug = repo_root
.file_name()
.map(|s| s.to_string_lossy().to_string())
.unwrap_or_else(|| "repo".to_string());
let hash = short_hash(repo_root.to_string_lossy().as_bytes());
Ok(cache.join("omnidev").join(format!("{slug}-{hash}")))
}
/// `${XDG_CACHE_HOME:-~/.cache}`.
pub fn cache_home() -> Result<PathBuf> {
if let Some(x) = std::env::var_os("XDG_CACHE_HOME") {
if !x.is_empty() {
return Ok(PathBuf::from(x));
}
}
let home = std::env::var_os("HOME").context("HOME is not set")?;
Ok(PathBuf::from(home).join(".cache"))
}
/// `${XDG_CONFIG_HOME:-~/.config}`.
pub fn config_home() -> Result<PathBuf> {
if let Some(x) = std::env::var_os("XDG_CONFIG_HOME") {
if !x.is_empty() {
return Ok(PathBuf::from(x));
}
}
let home = std::env::var_os("HOME").context("HOME is not set")?;
Ok(PathBuf::from(home).join(".config"))
}
/// `~/.config/omnidev/install.toml` — durable record of install intent.
pub fn install_config_path() -> Result<PathBuf> {
Ok(config_home()?.join("omnidev").join("install.toml"))
}
/// `~/.cache/omnidev/omnigent-check.json` — volatile update-check state.
pub fn check_cache_path() -> Result<PathBuf> {
Ok(cache_home()?.join("omnidev").join("omnigent-check.json"))
}
/// FNV-1a 64-bit, rendered as 8 hex chars. No external dep needed — we only
/// need a stable, collision-unlikely tag for a filesystem path.
fn short_hash(bytes: &[u8]) -> String {
let mut hash: u64 = 0xcbf29ce484222325;
for &b in bytes {
hash ^= b as u64;
hash = hash.wrapping_mul(0x100000001b3);
}
format!("{:08x}", (hash ^ (hash >> 32)) as u32)
}
+348
View File
@@ -0,0 +1,348 @@
//! A `Pod` = one isolated dev instance: its own state dir, ports, and the env
//! map injected into every supervised child.
use std::path::{Path, PathBuf};
use anyhow::{Context, Result};
use crate::ports::Ports;
pub struct Pod {
pub repo_root: PathBuf,
pub dir: PathBuf,
pub ports: Ports,
pub vite_host: String,
/// LAN origins to trust for device testing (`--trust-lan-origins`); empty
/// otherwise. Fed to the server as `OMNIGENT_WS_ALLOWED_ORIGINS`.
pub trusted_origins: Vec<String>,
}
impl Pod {
/// Create the pod directory tree (idempotent) and return the pod handle.
/// Only omnigent's own state is isolated (DB, artifacts, logs, config); the
/// pod inherits your real home, credentials, and caches.
pub fn create(
repo_root: PathBuf,
dir: PathBuf,
ports: Ports,
vite_host: String,
trusted_origins: Vec<String>,
) -> Result<Pod> {
for sub in ["data/omnigent", "artifacts", "logs", "config"] {
let p = dir.join(sub);
std::fs::create_dir_all(&p)
.with_context(|| format!("creating pod dir {}", p.display()))?;
}
let pod = Pod {
repo_root,
dir,
ports,
vite_host,
trusted_origins,
};
// Seed the pod's config from the developer's real one so it works out
// of the box (keeps their providers). Best-effort: a copy failure just
// starts the pod with an empty config, so warn rather than abort.
if let Some(src) = real_config_path() {
let dest = pod.config_dir().join("config.yaml");
if let Err(e) = seed_config_file(&src, &dest) {
eprintln!("omnidev: could not seed pod config: {e:#}");
}
}
Ok(pod)
}
pub fn db_uri(&self) -> String {
format!(
"sqlite:///{}",
self.dir.join("data/omnigent/chat.db").display()
)
}
pub fn artifacts_dir(&self) -> PathBuf {
self.dir.join("artifacts")
}
/// The pod's isolated config home, exposed to children as
/// `OMNIGENT_CONFIG_HOME` so its `config.yaml` is separate from the
/// developer's real `~/.omnigent/config.yaml`.
pub fn config_dir(&self) -> PathBuf {
self.dir.join("config")
}
pub fn server_url(&self) -> String {
format!("http://127.0.0.1:{}", self.ports.server)
}
/// Clickable URLs for display. Terminals linkify `localhost` but often not
/// a bare `127.0.0.1`. Functional uses (server bind, host `--server`,
/// `OMNIGENT_URL`) stay on `127.0.0.1` so we don't accidentally target IPv6
/// `localhost` (`::1`), where the server isn't listening.
pub fn server_display_url(&self) -> String {
format!("http://localhost:{}", self.ports.server)
}
pub fn vite_display_url(&self) -> String {
format!("http://localhost:{}", self.ports.vite)
}
pub fn web_dir(&self) -> PathBuf {
self.repo_root.join("web")
}
/// Whether `web/` needs `npm install` before Vite can start: either
/// `node_modules/` is absent, or the lockfile / `package.json` is newer
/// than the installed tree (a dependency was added/changed since the last
/// install — the case that makes Vite's dependency scan fail).
pub fn needs_npm_install(&self) -> bool {
let web = self.web_dir();
let modules = web.join("node_modules");
if !modules.is_dir() {
return true;
}
let mtime = |p: PathBuf| std::fs::metadata(p).and_then(|m| m.modified()).ok();
let Some(installed) = mtime(modules) else {
return true;
};
// Reinstall if either manifest is newer than node_modules.
[web.join("package-lock.json"), web.join("package.json")]
.into_iter()
.filter_map(mtime)
.any(|t| t > installed)
}
/// Directory to watch for backend source changes.
pub fn omnigent_dir(&self) -> PathBuf {
self.repo_root.join("omnigent")
}
pub fn log_file(&self, name: &str) -> PathBuf {
self.dir.join("logs").join(format!("{name}.log"))
}
/// The env overrides applied on top of the inherited parent env for every
/// child. We isolate omnigent's own state — the DB, data dir, and config
/// home — so concurrent pods don't share a database, pidfile, or
/// `config.yaml`. The rest (real `HOME`, credentials, uv/npm caches) is
/// inherited, since the agents omnigent runs need it. `OMNIGENT_URL` is the
/// seam `web/vite.config.ts` reads to point its proxy at this pod's backend;
/// `OMNIGENT_CONFIG_HOME` is where the server/host/runner read `config.yaml`.
pub fn env(&self) -> Vec<(String, String)> {
let d = |p: &str| self.dir.join(p).display().to_string();
let mut env = vec![
("OMNIGENT_DATA_DIR".into(), d("data/omnigent")),
("OMNIGENT_DATABASE_URI".into(), self.db_uri()),
("OMNIGENT_URL".into(), self.server_url()),
(
"OMNIGENT_CONFIG_HOME".into(),
self.config_dir().display().to_string(),
),
];
if let Some(allowed) = self.allowed_origins_env() {
env.push(("OMNIGENT_WS_ALLOWED_ORIGINS".into(), allowed));
}
env
}
/// The `OMNIGENT_WS_ALLOWED_ORIGINS` value to inject, or `None` to leave it
/// untouched. Merges the trusted LAN origins onto any value inherited from
/// the parent environment (comma-separated, order-preserving, deduped) so a
/// developer's own allowlist survives. Returns `None` when there are no LAN
/// origins to add — then the parent's value (if any) simply passes through.
fn allowed_origins_env(&self) -> Option<String> {
if self.trusted_origins.is_empty() {
return None;
}
let inherited = std::env::var("OMNIGENT_WS_ALLOWED_ORIGINS").unwrap_or_default();
let mut merged: Vec<String> = Vec::new();
let parts = inherited
.split(',')
.map(str::trim)
.filter(|s| !s.is_empty())
.map(str::to_string)
.chain(self.trusted_origins.iter().cloned());
for part in parts {
if !merged.contains(&part) {
merged.push(part);
}
}
Some(merged.join(","))
}
}
/// Remove a pod directory (for `--clean`). No-op if it does not exist.
pub fn clean(dir: &Path) -> Result<()> {
if dir.exists() {
std::fs::remove_dir_all(dir)
.with_context(|| format!("removing pod dir {}", dir.display()))?;
}
Ok(())
}
/// The developer's real omnigent `config.yaml` to seed a fresh pod from.
///
/// Honors `OMNIGENT_CONFIG_HOME` if the parent env sets it (nested/test
/// setups), else `~/.omnigent/config.yaml` via `HOME`. Returns `None` when the
/// file does not exist — a fresh pod then starts with an empty config, just
/// like a first-run user.
fn real_config_path() -> Option<PathBuf> {
let home = match std::env::var_os("OMNIGENT_CONFIG_HOME") {
Some(h) if !h.is_empty() => PathBuf::from(h),
_ => PathBuf::from(std::env::var_os("HOME")?).join(".omnigent"),
};
let path = home.join("config.yaml");
path.exists().then_some(path)
}
/// Copy `src` to `dest`, but only when `dest` does not already exist — a normal
/// pod restart must not clobber config the developer edited inside the pod.
/// After `--clean` the whole pod dir is gone, so `dest` is absent and this
/// re-seeds.
fn seed_config_file(src: &Path, dest: &Path) -> Result<()> {
if dest.exists() {
return Ok(());
}
std::fs::copy(src, dest)
.with_context(|| format!("seeding {} from {}", dest.display(), src.display()))?;
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Mutex;
// `real_config_path` reads process-global env; serialize the tests that
// set it so parallel runs don't observe each other's overrides.
static ENV_LOCK: Mutex<()> = Mutex::new(());
fn tempdir() -> PathBuf {
let unique = format!(
"omnidev-pod-test-{}-{}",
std::process::id(),
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos()
);
let dir = std::env::temp_dir().join(unique);
std::fs::create_dir_all(&dir).unwrap();
dir
}
fn make_pod(pod_dir: PathBuf) -> Pod {
Pod::create(
tempdir(),
pod_dir,
Ports {
server: 19191,
vite: 19292,
},
"127.0.0.1".into(),
Vec::new(),
)
.unwrap()
}
/// Point `OMNIGENT_CONFIG_HOME` at `home` for the duration of `f`, restoring
/// the previous value afterwards. Serialized against other env-touching
/// tests via `ENV_LOCK`.
fn with_config_home<T>(home: &Path, f: impl FnOnce() -> T) -> T {
let _guard = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
let prev = std::env::var_os("OMNIGENT_CONFIG_HOME");
std::env::set_var("OMNIGENT_CONFIG_HOME", home);
let out = f();
match prev {
Some(v) => std::env::set_var("OMNIGENT_CONFIG_HOME", v),
None => std::env::remove_var("OMNIGENT_CONFIG_HOME"),
}
out
}
#[test]
fn create_makes_config_dir() {
let real = tempdir(); // empty config home -> nothing to seed
let pod = with_config_home(&real, || make_pod(tempdir()));
assert!(pod.config_dir().is_dir());
}
#[test]
fn env_includes_config_home() {
let real = tempdir();
let pod = with_config_home(&real, || make_pod(tempdir()));
let env = pod.env();
let got = env
.iter()
.find(|(k, _)| k == "OMNIGENT_CONFIG_HOME")
.map(|(_, v)| v.clone());
assert_eq!(got, Some(pod.config_dir().display().to_string()));
}
#[test]
fn create_seeds_pod_config_from_real() {
let real = tempdir();
std::fs::write(real.join("config.yaml"), "providers:\n seeded: true\n").unwrap();
let pod = with_config_home(&real, || make_pod(tempdir()));
let seeded = std::fs::read_to_string(pod.config_dir().join("config.yaml")).unwrap();
assert_eq!(seeded, "providers:\n seeded: true\n");
}
#[test]
fn create_skips_seed_when_real_config_absent() {
let real = tempdir(); // no config.yaml inside
let pod = with_config_home(&real, || make_pod(tempdir()));
assert!(!pod.config_dir().join("config.yaml").exists());
}
#[test]
fn seed_does_not_overwrite_existing() {
let dir = tempdir();
let src = dir.join("src.yaml");
let dest = dir.join("dest.yaml");
std::fs::write(&src, "from: real\n").unwrap();
std::fs::write(&dest, "edited: in-pod\n").unwrap();
seed_config_file(&src, &dest).unwrap();
// Existing pod-local edits survive; the real config does not clobber them.
assert_eq!(std::fs::read_to_string(&dest).unwrap(), "edited: in-pod\n");
}
#[test]
fn real_config_path_honors_config_home() {
let real = tempdir();
std::fs::write(real.join("config.yaml"), "x: 1\n").unwrap();
let got = with_config_home(&real, real_config_path);
assert_eq!(got, Some(real.join("config.yaml")));
}
#[test]
fn real_config_path_falls_back_to_home_dot_omnigent() {
// With no OMNIGENT_CONFIG_HOME, the real config resolves under
// `$HOME/.omnigent/` — the path a normal pod run seeds from.
let _guard = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
let prev_cfg = std::env::var_os("OMNIGENT_CONFIG_HOME");
let prev_home = std::env::var_os("HOME");
let home = tempdir();
std::fs::create_dir_all(home.join(".omnigent")).unwrap();
std::fs::write(home.join(".omnigent/config.yaml"), "y: 2\n").unwrap();
std::env::remove_var("OMNIGENT_CONFIG_HOME");
std::env::set_var("HOME", &home);
let got = real_config_path();
match prev_cfg {
Some(v) => std::env::set_var("OMNIGENT_CONFIG_HOME", v),
None => std::env::remove_var("OMNIGENT_CONFIG_HOME"),
}
match prev_home {
Some(v) => std::env::set_var("HOME", v),
None => std::env::remove_var("HOME"),
}
assert_eq!(got, Some(home.join(".omnigent/config.yaml")));
}
}
+124
View File
@@ -0,0 +1,124 @@
//! Free-port probing and per-pod persistence.
use std::collections::HashSet;
use std::net::TcpListener;
use std::path::Path;
use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
pub const SERVER_PORT_BASE: u16 = 6767;
pub const VITE_PORT_BASE: u16 = 5173;
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub struct Ports {
pub server: u16,
pub vite: u16,
}
impl Ports {
/// Resolve the pod's ports: reuse the persisted pair if still available,
/// else probe upward from the preferred bases. Explicit overrides (from CLI
/// flags) are honored verbatim.
///
/// A port is "available" only if it both binds right now *and* isn't already
/// claimed by another pod. The bind check alone is racy: `resolve()` runs at
/// startup, before children spawn, so a peer pod whose server/vite hasn't
/// bound yet would leave the base port looking free and two pods would pick
/// it. We read sibling pods' persisted `pod.toml` to skip ports they've
/// already claimed, which is timing-independent.
pub fn resolve(
pod_dir: &Path,
server_override: Option<u16>,
vite_override: Option<u16>,
) -> Result<Ports> {
let persisted = load(pod_dir);
let mut taken = sibling_claims(pod_dir);
let server = match server_override {
Some(p) => p,
None => {
let reuse = persisted
.map(|p| p.server)
.filter(|&p| available(p, &taken));
reuse
.map(Ok)
.unwrap_or_else(|| probe_from(SERVER_PORT_BASE, &taken))?
}
};
// The server port is now spoken for — don't hand the same number to vite.
taken.insert(server);
let vite = match vite_override {
Some(p) => p,
None => {
let reuse = persisted.map(|p| p.vite).filter(|&p| available(p, &taken));
reuse
.map(Ok)
.unwrap_or_else(|| probe_from(VITE_PORT_BASE, &taken))?
}
};
let ports = Ports { server, vite };
save(pod_dir, &ports)?;
Ok(ports)
}
}
/// A port is usable if it isn't already claimed by a sibling pod and binds now.
fn available(port: u16, taken: &HashSet<u16>) -> bool {
!taken.contains(&port) && is_free(port)
}
/// True if the port can be bound on loopback right now.
fn is_free(port: u16) -> bool {
TcpListener::bind(("127.0.0.1", port)).is_ok()
}
/// First available port at or above `base`, skipping sibling-claimed ports.
fn probe_from(base: u16, taken: &HashSet<u16>) -> Result<u16> {
for port in base..=u16::MAX {
if available(port, taken) {
return Ok(port);
}
}
anyhow::bail!("no free port at or above {base}")
}
/// Ports claimed in other pods' `pod.toml` under the shared omnidev cache root.
/// Best-effort: unreadable/oddly-nested pod dirs just contribute nothing.
fn sibling_claims(pod_dir: &Path) -> HashSet<u16> {
let mut claimed = HashSet::new();
let Some(root) = pod_dir.parent() else {
return claimed;
};
let Ok(entries) = std::fs::read_dir(root) else {
return claimed;
};
for entry in entries.flatten() {
let dir = entry.path();
if dir == pod_dir || !dir.is_dir() {
continue;
}
if let Some(p) = load(&dir) {
claimed.insert(p.server);
claimed.insert(p.vite);
}
}
claimed
}
fn persist_path(pod_dir: &Path) -> std::path::PathBuf {
pod_dir.join("pod.toml")
}
fn load(pod_dir: &Path) -> Option<Ports> {
let text = std::fs::read_to_string(persist_path(pod_dir)).ok()?;
toml::from_str(&text).ok()
}
fn save(pod_dir: &Path, ports: &Ports) -> Result<()> {
let text = toml::to_string(ports).context("serializing pod.toml")?;
std::fs::write(persist_path(pod_dir), text).context("writing pod.toml")?;
Ok(())
}
+133
View File
@@ -0,0 +1,133 @@
//! Concrete command specs for the three supervised processes.
use std::path::PathBuf;
use crate::pod::Pod;
/// A resolved command line + working dir for one process. Env is applied by the
/// supervisor from `Pod::env()`, so it is not duplicated here.
pub struct ProcSpec {
pub program: String,
pub args: Vec<String>,
pub cwd: PathBuf,
}
impl ProcSpec {
/// `uv run omnigent server --host 127.0.0.1 --port <p> --database-uri <db>
/// --artifact-location <dir>`, from the repo root.
pub fn server(pod: &Pod) -> ProcSpec {
ProcSpec {
program: "uv".into(),
args: vec![
"run".into(),
"omnigent".into(),
"server".into(),
"--host".into(),
"127.0.0.1".into(),
"--port".into(),
pod.ports.server.to_string(),
"--database-uri".into(),
pod.db_uri(),
"--artifact-location".into(),
pod.artifacts_dir().display().to_string(),
],
cwd: pod.repo_root.clone(),
}
}
/// `uv run omnigent host --server http://127.0.0.1:<p>`, from the repo root.
pub fn host(pod: &Pod) -> ProcSpec {
ProcSpec {
program: "uv".into(),
args: vec![
"run".into(),
"omnigent".into(),
"host".into(),
"--server".into(),
pod.server_url(),
],
cwd: pod.repo_root.clone(),
}
}
/// `npm install`, from `web/`. Run before Vite when deps are missing or
/// stale so Vite's dependency scan doesn't fail on an unresolved import.
///
/// `--loglevel http` makes npm emit a line per package fetch even when its
/// stdout is piped (its progress bar is TTY-only), so the pane streams real
/// progress. `--no-fund --no-audit` trims the trailing noise.
pub fn npm_install(pod: &Pod) -> ProcSpec {
ProcSpec {
program: "npm".into(),
args: vec![
"install".into(),
"--no-fund".into(),
"--no-audit".into(),
"--loglevel".into(),
"http".into(),
],
cwd: pod.web_dir(),
}
}
/// `npm run dev -- --host <host> --port <p> --strictPort`, from `web/`.
/// `OMNIGENT_URL` (in the pod env) points Vite's proxy at this pod's backend.
pub fn vite(pod: &Pod) -> ProcSpec {
ProcSpec {
program: "npm".into(),
args: vec![
"run".into(),
"dev".into(),
"--".into(),
"--host".into(),
pod.vite_host.clone(),
"--port".into(),
pod.ports.vite.to_string(),
"--strictPort".into(),
],
cwd: pod.web_dir(),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::ports::Ports;
#[test]
fn vite_uses_configured_bind_host_but_backend_url_stays_loopback() {
let repo = tempdir();
let pod_dir = tempdir();
let pod = Pod::create(
repo,
pod_dir,
Ports {
server: 19191,
vite: 19292,
},
"0.0.0.0".into(),
Vec::new(),
)
.unwrap();
let vite = ProcSpec::vite(&pod);
let host_flag = vite.args.iter().position(|arg| arg == "--host").unwrap();
assert_eq!(vite.args[host_flag + 1], "0.0.0.0");
assert_eq!(pod.server_url(), "http://127.0.0.1:19191");
}
fn tempdir() -> std::path::PathBuf {
let unique = format!(
"omnidev-process-test-{}-{}",
std::process::id(),
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos()
);
let dir = std::env::temp_dir().join(unique);
std::fs::create_dir_all(&dir).unwrap();
dir
}
}
+19
View File
@@ -0,0 +1,19 @@
//! Emit the shell snippet that runs the daily update check.
/// The snippet to append to `.zshrc`/`.bashrc`
/// (`omnidev shell-hook >> ~/.zshrc`). All throttling and prompting live inside
/// `omnidev check`, so this stays trivial and shell-agnostic: run once per
/// interactive shell, quietly, and never fail the shell if it errors.
///
/// It self-guards on `command -v omnidev`, so it's meant to be appended to the
/// rc (a static no-op when omnidev is absent) rather than run via
/// `eval "$(omnidev shell-hook)"`, which would invoke omnidev on every shell
/// startup and error when it isn't on PATH.
const HOOK: &str = r#"# omnidev: daily omnigent update check
if [ -n "${PS1:-}" ] && command -v omnidev >/dev/null 2>&1; then
omnidev check --quiet || true
fi"#;
pub fn print() {
println!("{HOOK}");
}
+111
View File
@@ -0,0 +1,111 @@
//! Shared state between the supervisor and the TUI.
use std::sync::{Arc, Mutex};
use crate::logs::LogBuffer;
use crate::pod::Pod;
/// The three supervised processes.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ProcId {
Server,
Host,
Vite,
}
impl ProcId {
pub const ALL: [ProcId; 3] = [ProcId::Server, ProcId::Host, ProcId::Vite];
pub fn idx(self) -> usize {
match self {
ProcId::Server => 0,
ProcId::Host => 1,
ProcId::Vite => 2,
}
}
pub fn label(self) -> &'static str {
match self {
ProcId::Server => "server",
ProcId::Host => "host",
ProcId::Vite => "vite",
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ProcStatus {
Idle,
Starting,
Running(u32),
Restarting,
Crashed,
Stopped,
}
impl ProcStatus {
pub fn short(&self) -> &'static str {
match self {
ProcStatus::Idle => "idle",
ProcStatus::Starting => "starting",
ProcStatus::Running(_) => "running",
ProcStatus::Restarting => "restarting",
ProcStatus::Crashed => "crashed",
ProcStatus::Stopped => "stopped",
}
}
}
/// State the TUI renders and the supervisor mutates. Guarded by a std mutex;
/// locks are held only for the duration of a single push/read.
pub struct Shared {
pub status: [ProcStatus; 3],
pub server: LogBuffer,
pub host: LogBuffer,
pub vite: LogBuffer,
/// Combined, source-tagged view — also receives supervisor events.
pub all: LogBuffer,
}
impl Shared {
pub fn new(pod: &Pod) -> Arc<Mutex<Shared>> {
Arc::new(Mutex::new(Shared {
status: [ProcStatus::Idle, ProcStatus::Idle, ProcStatus::Idle],
server: LogBuffer::new(&pod.log_file("server")),
host: LogBuffer::new(&pod.log_file("host")),
vite: LogBuffer::new(&pod.log_file("vite")),
all: LogBuffer::memory(),
}))
}
fn buf_mut(&mut self, id: ProcId) -> &mut LogBuffer {
match id {
ProcId::Server => &mut self.server,
ProcId::Host => &mut self.host,
ProcId::Vite => &mut self.vite,
}
}
pub fn buf(&self, id: ProcId) -> &LogBuffer {
match id {
ProcId::Server => &self.server,
ProcId::Host => &self.host,
ProcId::Vite => &self.vite,
}
}
/// Append a line from a process: goes to its own pane and the combined view.
pub fn log_proc(&mut self, id: ProcId, line: String) {
self.all.push(format!("[{}] {}", id.label(), line));
self.buf_mut(id).push(line);
}
/// Append a supervisor event (starts, restarts, crashes, reloads).
pub fn event(&mut self, line: impl Into<String>) {
self.all.push(format!("[omnidev] {}", line.into()));
}
pub fn set_status(&mut self, id: ProcId, status: ProcStatus) {
self.status[id.idx()] = status;
}
}
+493
View File
@@ -0,0 +1,493 @@
//! Process supervision: spawn/stop/restart the three children, capture their
//! output, and recover from crashes.
use std::collections::HashSet;
use std::process::Stdio;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use tokio::io::{AsyncBufReadExt, BufReader};
use tokio::net::TcpStream;
use tokio::process::Command;
use tokio::sync::mpsc;
use tokio::time::{sleep, timeout};
use crate::pod::Pod;
use crate::process::ProcSpec;
use crate::state::{ProcId, ProcStatus, Shared};
/// Commands the TUI (and watcher) send to the supervisor.
#[derive(Debug, Clone)]
pub enum Cmd {
/// Restart a single process.
Restart(ProcId),
/// Restart the backend pair: server, then host after `/health`.
RestartBackend,
/// A backend reload triggered by `n` changed Python files.
Reload(usize),
/// Tear everything down and stop the supervisor loop.
Shutdown,
}
/// Reported by a per-child monitor when the child exits.
struct Exit {
id: ProcId,
generation: u64,
status: String,
}
struct Slot {
/// Group id (== leader pid) of the currently-running child, if any.
pgid: Option<i32>,
/// Generation of the current child; bumped on each spawn.
generation: u64,
/// Consecutive crash count for backoff; reset after a stable run.
crashes: u32,
started: Instant,
}
impl Default for Slot {
fn default() -> Self {
Slot {
pgid: None,
generation: 0,
crashes: 0,
started: Instant::now(),
}
}
}
pub struct Supervisor {
pod: Arc<Pod>,
shared: Arc<Mutex<Shared>>,
env: Vec<(String, String)>,
vite_enabled: bool,
/// Whether `--trust-lan-origins` was requested, so we can warn if it was
/// asked for but no LAN interface turned up any origins to trust.
trust_lan_origins: bool,
slots: [Slot; 3],
/// Generations we stopped on purpose — their exits are not crashes.
expected_stops: HashSet<(usize, u64)>,
gen_counter: u64,
exit_tx: mpsc::UnboundedSender<Exit>,
exit_rx: mpsc::UnboundedReceiver<Exit>,
}
impl Supervisor {
pub fn new(
pod: Arc<Pod>,
shared: Arc<Mutex<Shared>>,
vite_enabled: bool,
trust_lan_origins: bool,
) -> Supervisor {
let env = pod.env();
let (exit_tx, exit_rx) = mpsc::unbounded_channel();
Supervisor {
pod,
shared,
env,
vite_enabled,
trust_lan_origins,
slots: Default::default(),
expected_stops: HashSet::new(),
gen_counter: 0,
exit_tx,
exit_rx,
}
}
fn event(&self, msg: impl Into<String>) {
self.shared.lock().unwrap().event(msg);
}
fn set_status(&self, id: ProcId, status: ProcStatus) {
self.shared.lock().unwrap().set_status(id, status);
}
/// Main loop: bring everything up, then service commands and child exits
/// until `Shutdown`.
pub async fn run(mut self, mut cmds: mpsc::UnboundedReceiver<Cmd>) {
self.event(format!(
"pod {} — server :{} vite :{}",
self.pod.dir.display(),
self.pod.ports.server,
self.pod.ports.vite
));
if !self.pod.trusted_origins.is_empty() {
self.event(format!(
"trusting LAN origins for device testing: {}",
self.pod.trusted_origins.join(", ")
));
} else if self.trust_lan_origins {
self.event("--trust-lan-origins: no LAN interface found; no extra origins trusted");
}
self.start_backend().await;
if self.vite_enabled {
self.prepare_vite().await;
self.spawn(ProcId::Vite);
}
loop {
tokio::select! {
cmd = cmds.recv() => {
match cmd {
Some(Cmd::Restart(id)) => self.restart_one(id).await,
Some(Cmd::RestartBackend) => {
self.event("manual backend restart");
self.start_backend_restart().await;
}
Some(Cmd::Reload(n)) => {
self.event(format!("reloading backend ({n} file(s) changed)"));
self.start_backend_restart().await;
}
Some(Cmd::Shutdown) | None => {
self.shutdown().await;
return;
}
}
}
Some(exit) = self.exit_rx.recv() => {
self.on_exit(exit).await;
}
}
}
}
async fn start_backend(&mut self) {
self.spawn(ProcId::Server);
if self.wait_healthy().await {
self.spawn(ProcId::Host);
} else {
self.event("server did not become healthy; host not started");
}
}
/// Restart server then host, gated on `/health`. Used by manual restart and
/// by the reload path.
async fn start_backend_restart(&mut self) {
self.stop(ProcId::Host).await;
self.stop(ProcId::Server).await;
self.set_status(ProcId::Server, ProcStatus::Restarting);
self.set_status(ProcId::Host, ProcStatus::Restarting);
self.spawn(ProcId::Server);
if self.wait_healthy().await {
self.spawn(ProcId::Host);
} else {
self.event("server did not become healthy after restart");
}
}
async fn restart_one(&mut self, id: ProcId) {
match id {
// Restarting the server alone would strand the host on a dead
// backend, so treat it as a backend restart.
ProcId::Server | ProcId::Host => self.start_backend_restart().await,
ProcId::Vite => {
if self.vite_enabled {
self.event("restarting vite");
self.stop(ProcId::Vite).await;
self.prepare_vite().await;
self.spawn(ProcId::Vite);
}
}
}
}
fn spec(&self, id: ProcId) -> ProcSpec {
match id {
ProcId::Server => ProcSpec::server(&self.pod),
ProcId::Host => ProcSpec::host(&self.pod),
ProcId::Vite => ProcSpec::vite(&self.pod),
}
}
/// Spawn a child in its own process group and wire up output + exit monitor.
fn spawn(&mut self, id: ProcId) {
let spec = self.spec(id);
self.set_status(id, ProcStatus::Starting);
let mut cmd = Command::new(&spec.program);
cmd.args(&spec.args)
.current_dir(&spec.cwd)
.envs(self.env.iter().cloned())
.stdin(Stdio::null())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.kill_on_drop(false);
// Become a session/group leader so we can signal the whole tree
// (uvicorn workers, npm -> vite children) via the negative pgid.
unsafe {
cmd.pre_exec(|| {
libc::setsid();
Ok(())
});
}
let mut child = match cmd.spawn() {
Ok(c) => c,
Err(e) => {
self.shared
.lock()
.unwrap()
.log_proc(id, format!("failed to spawn {}: {e}", spec.program));
self.set_status(id, ProcStatus::Crashed);
return;
}
};
let pid = child.id().map(|p| p as i32);
self.gen_counter += 1;
let generation = self.gen_counter;
let slot = &mut self.slots[id.idx()];
slot.pgid = pid;
slot.generation = generation;
slot.started = Instant::now();
if let Some(p) = pid {
self.set_status(id, ProcStatus::Running(p as u32));
}
// Merge stdout + stderr into this process's buffer.
if let Some(out) = child.stdout.take() {
self.pump(id, out);
}
if let Some(err) = child.stderr.take() {
self.pump(id, err);
}
// Monitor: report the exit so the loop can decide crash vs expected.
let tx = self.exit_tx.clone();
tokio::spawn(async move {
let status = match child.wait().await {
Ok(s) => s.to_string(),
Err(e) => format!("wait error: {e}"),
};
let _ = tx.send(Exit {
id,
generation,
status,
});
});
}
/// Run `npm install` to completion before Vite starts, but only when deps
/// are missing or stale — otherwise Vite's dependency scan fails on an
/// unresolved import (e.g. a dep added to package.json but not installed).
/// Output streams into the Vite pane. A failed/absent install is logged but
/// non-fatal: we still let Vite try, so a transient npm hiccup doesn't block
/// the whole session.
async fn prepare_vite(&self) {
if !self.pod.needs_npm_install() {
return;
}
self.set_status(ProcId::Vite, ProcStatus::Starting);
self.shared.lock().unwrap().log_proc(
ProcId::Vite,
"web deps missing or stale — running npm install".into(),
);
let spec = ProcSpec::npm_install(&self.pod);
let mut cmd = Command::new(&spec.program);
cmd.args(&spec.args)
.current_dir(&spec.cwd)
.envs(self.env.iter().cloned())
.stdin(Stdio::null())
.stdout(Stdio::piped())
.stderr(Stdio::piped());
let mut child = match cmd.spawn() {
Ok(c) => c,
Err(e) => {
self.shared
.lock()
.unwrap()
.log_proc(ProcId::Vite, format!("failed to run npm install: {e}"));
return;
}
};
if let Some(out) = child.stdout.take() {
self.pump(ProcId::Vite, out);
}
if let Some(err) = child.stderr.take() {
self.pump(ProcId::Vite, err);
}
// `--loglevel http` streams a line per package fetch, but npm still
// goes quiet during the final tree-build/link phase. A slow heartbeat
// covers those gaps so the pane never looks frozen.
let started = Instant::now();
let mut heartbeat = tokio::time::interval(Duration::from_secs(5));
heartbeat.tick().await; // the first tick fires immediately; skip it
let status = loop {
tokio::select! {
result = child.wait() => break result,
_ = heartbeat.tick() => {
let secs = started.elapsed().as_secs();
self.shared
.lock()
.unwrap()
.log_proc(ProcId::Vite, format!("… npm install running ({secs}s)"));
}
}
};
match status {
Ok(s) if s.success() => self.event(format!(
"npm install complete ({}s)",
started.elapsed().as_secs()
)),
Ok(s) => self.event(format!("npm install exited {s} — starting Vite anyway")),
Err(e) => self.event(format!("npm install wait error: {e}")),
}
}
/// Spawn a task that streams one pipe into the shared buffer, line by line.
fn pump<R>(&self, id: ProcId, reader: R)
where
R: tokio::io::AsyncRead + Unpin + Send + 'static,
{
let shared = self.shared.clone();
tokio::spawn(async move {
let mut lines = BufReader::new(reader).lines();
while let Ok(Some(line)) = lines.next_line().await {
shared.lock().unwrap().log_proc(id, line);
}
});
}
/// SIGTERM the process group, wait briefly, then SIGKILL. Marks the current
/// generation as an expected stop so its exit is not counted as a crash.
async fn stop(&mut self, id: ProcId) {
let (pgid, generation) = {
let slot = &self.slots[id.idx()];
(slot.pgid, slot.generation)
};
let Some(pgid) = pgid else {
self.set_status(id, ProcStatus::Stopped);
return;
};
self.expected_stops.insert((id.idx(), generation));
unsafe {
libc::kill(-pgid, libc::SIGTERM);
}
// Give the tree up to ~5s to exit on SIGTERM.
for _ in 0..50 {
if unsafe { libc::kill(-pgid, 0) } != 0 {
break;
}
sleep(Duration::from_millis(100)).await;
}
if unsafe { libc::kill(-pgid, 0) } == 0 {
unsafe {
libc::kill(-pgid, libc::SIGKILL);
}
}
self.slots[id.idx()].pgid = None;
self.set_status(id, ProcStatus::Stopped);
}
/// Handle a child exit: distinguish an expected stop from a crash and
/// schedule a backoff restart for crashes.
async fn on_exit(&mut self, exit: Exit) {
let key = (exit.id.idx(), exit.generation);
if self.expected_stops.remove(&key) {
return; // we stopped it on purpose
}
// Ignore exits from a generation we already replaced.
if self.slots[exit.id.idx()].generation != exit.generation {
return;
}
self.slots[exit.id.idx()].pgid = None;
self.set_status(exit.id, ProcStatus::Crashed);
self.event(format!(
"{} exited unexpectedly ({})",
exit.id.label(),
exit.status
));
// Reset the crash counter if the process had been stable for a while.
let crashes = {
let slot = &mut self.slots[exit.id.idx()];
if slot.started.elapsed() > Duration::from_secs(20) {
slot.crashes = 0;
}
slot.crashes += 1;
slot.crashes
};
let backoff = backoff_secs(crashes);
self.event(format!(
"restarting {} in {backoff}s (attempt {crashes})",
exit.id.label(),
));
sleep(Duration::from_secs(backoff)).await;
// A server crash takes the host with it — restart the pair.
match exit.id {
ProcId::Server => self.start_backend_restart().await,
ProcId::Host => {
if self.wait_healthy().await {
self.spawn(ProcId::Host);
} else {
self.start_backend_restart().await;
}
}
ProcId::Vite => {
if self.vite_enabled {
self.spawn(ProcId::Vite);
}
}
}
}
/// Poll the server's `/health` until it returns 200 (up to ~30s).
async fn wait_healthy(&self) -> bool {
let addr = format!("127.0.0.1:{}", self.pod.ports.server);
for _ in 0..120 {
if health_ok(&addr).await {
return true;
}
sleep(Duration::from_millis(250)).await;
}
false
}
async fn shutdown(&mut self) {
self.event("shutting down");
self.stop(ProcId::Host).await;
self.stop(ProcId::Vite).await;
self.stop(ProcId::Server).await;
}
}
fn backoff_secs(attempt: u32) -> u64 {
// 0.5s effectively rounds to 1s here; cap at 30s.
match attempt {
0 | 1 => 1,
2 => 2,
3 => 4,
4 => 8,
5 => 16,
_ => 30,
}
}
/// Minimal HTTP/1.0 `GET /health` returning true on a `200` status line. Avoids
/// pulling an HTTP client dependency just for a readiness probe.
async fn health_ok(addr: &str) -> bool {
let Ok(Ok(mut stream)) = timeout(Duration::from_secs(1), TcpStream::connect(addr)).await else {
return false;
};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
let req = format!("GET /health HTTP/1.0\r\nHost: {addr}\r\n\r\n");
if stream.write_all(req.as_bytes()).await.is_err() {
return false;
}
let mut buf = [0u8; 128];
let Ok(Ok(n)) = timeout(Duration::from_secs(1), stream.read(&mut buf)).await else {
return false;
};
let head = String::from_utf8_lossy(&buf[..n]);
head.starts_with("HTTP/1.") && head.contains(" 200")
}
+216
View File
@@ -0,0 +1,216 @@
//! Terminal UI: renders pod status + per-process log panes and turns key
//! presses into supervisor commands.
mod render;
use std::io::{self, Stdout};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use anyhow::Result;
use crossterm::event::{self, Event, KeyCode, KeyEvent, KeyEventKind, KeyModifiers};
use crossterm::execute;
use crossterm::terminal::{
disable_raw_mode, enable_raw_mode, EnterAlternateScreen, LeaveAlternateScreen,
};
use ratatui::backend::CrosstermBackend;
use ratatui::Terminal;
use tokio::sync::mpsc;
use crate::pod::Pod;
use crate::state::{ProcId, Shared};
use crate::supervisor::Cmd;
/// Which log channel is focused. `All` is the combined, source-tagged view.
#[derive(Clone, Copy, PartialEq, Eq)]
pub enum View {
Server,
Host,
Vite,
All,
}
impl View {
fn proc(self) -> Option<ProcId> {
match self {
View::Server => Some(ProcId::Server),
View::Host => Some(ProcId::Host),
View::Vite => Some(ProcId::Vite),
View::All => None,
}
}
}
pub struct App {
pod: Arc<Pod>,
shared: Arc<Mutex<Shared>>,
cmds: mpsc::UnboundedSender<Cmd>,
view: View,
/// Lines scrolled up from the bottom; 0 == pinned to tail.
scroll_back: usize,
follow: bool,
should_quit: bool,
}
impl App {
pub fn new(pod: Arc<Pod>, shared: Arc<Mutex<Shared>>, cmds: mpsc::UnboundedSender<Cmd>) -> App {
App {
pod,
shared,
cmds,
view: View::All,
scroll_back: 0,
follow: true,
should_quit: false,
}
}
/// Run the render + input loop until the user quits. On return, the caller
/// sends `Shutdown` and the terminal is already restored.
pub async fn run(mut self) -> Result<()> {
let mut terminal = setup_terminal()?;
let mut input = spawn_input();
let mut tick = tokio::time::interval(Duration::from_millis(80));
let result = loop {
if let Err(e) = terminal.draw(|f| render::draw(f, &self)) {
break Err(e.into());
}
if self.should_quit {
break Ok(());
}
tokio::select! {
_ = tick.tick() => {}
key = input.recv() => {
match key {
Some(key) => self.on_key(key),
None => break Ok(()),
}
}
}
};
restore_terminal(&mut terminal);
result
}
fn on_key(&mut self, key: KeyEvent) {
if key.kind != KeyEventKind::Press {
return;
}
let page = 20;
match (key.code, key.modifiers) {
(KeyCode::Char('c'), KeyModifiers::CONTROL) => self.should_quit = true,
(KeyCode::Char('q'), _) => self.should_quit = true,
(KeyCode::Char('1'), _) => self.set_view(View::Server),
(KeyCode::Char('2'), _) => self.set_view(View::Host),
(KeyCode::Char('3'), _) => self.set_view(View::Vite),
(KeyCode::Char('0'), _) => self.set_view(View::All),
(KeyCode::Tab, _) => self.cycle_view(),
(KeyCode::Up, _) => self.scroll(1),
(KeyCode::Down, _) => self.scroll_down(1),
(KeyCode::PageUp, _) => self.scroll(page),
(KeyCode::PageDown, _) => self.scroll_down(page),
(KeyCode::Char('f'), _) => {
self.follow = !self.follow;
if self.follow {
self.scroll_back = 0;
}
}
(KeyCode::Char('r'), _) => {
if let Some(id) = self.view.proc() {
let _ = self.cmds.send(Cmd::Restart(id));
} else {
let _ = self.cmds.send(Cmd::RestartBackend);
}
}
(KeyCode::Char('R'), _) => {
let _ = self.cmds.send(Cmd::RestartBackend);
}
(KeyCode::Char('c'), _) => self.clear_current(),
_ => {}
}
}
fn set_view(&mut self, v: View) {
self.view = v;
self.scroll_back = 0;
}
fn cycle_view(&mut self) {
self.view = match self.view {
View::All => View::Server,
View::Server => View::Host,
View::Host => View::Vite,
View::Vite => View::All,
};
self.scroll_back = 0;
}
fn scroll(&mut self, n: usize) {
// Scrolling up detaches from the tail.
self.follow = false;
self.scroll_back = self.scroll_back.saturating_add(n);
}
fn scroll_down(&mut self, n: usize) {
self.scroll_back = self.scroll_back.saturating_sub(n);
if self.scroll_back == 0 {
self.follow = true;
}
}
fn clear_current(&mut self) {
let mut s = self.shared.lock().unwrap();
match self.view {
View::Server => s.server.clear(),
View::Host => s.host.clear(),
View::Vite => s.vite.clear(),
View::All => s.all.clear(),
}
self.scroll_back = 0;
}
/// Total line count of the focused channel, for the status readout.
pub fn line_count(&self) -> usize {
let s = self.shared.lock().unwrap();
match self.view {
View::Server => s.buf(ProcId::Server).iter().count(),
View::Host => s.buf(ProcId::Host).iter().count(),
View::Vite => s.buf(ProcId::Vite).iter().count(),
View::All => s.all.iter().count(),
}
}
}
fn setup_terminal() -> Result<Terminal<CrosstermBackend<Stdout>>> {
enable_raw_mode()?;
let mut stdout = io::stdout();
execute!(stdout, EnterAlternateScreen)?;
Ok(Terminal::new(CrosstermBackend::new(stdout))?)
}
fn restore_terminal(terminal: &mut Terminal<CrosstermBackend<Stdout>>) {
let _ = disable_raw_mode();
let _ = execute!(terminal.backend_mut(), LeaveAlternateScreen);
let _ = terminal.show_cursor();
}
/// Read crossterm key events on a dedicated thread and forward them; the async
/// loop selects on this alongside the render tick.
fn spawn_input() -> mpsc::UnboundedReceiver<KeyEvent> {
let (tx, rx) = mpsc::unbounded_channel();
std::thread::spawn(move || loop {
if event::poll(Duration::from_millis(200)).unwrap_or(false) {
if let Ok(Event::Key(key)) = event::read() {
if tx.send(key).is_err() {
break;
}
}
}
});
rx
}
+253
View File
@@ -0,0 +1,253 @@
//! Frame rendering. Minimal chrome: no boxes — regions are separated by a
//! light neutral background bar instead. The header and footer share the
//! "chrome" bar; the log body sits on the terminal's default background so
//! ANSI log colors render naturally on either a light or dark theme.
use ansi_to_tui::IntoText;
use ratatui::layout::{Alignment, Constraint, Direction, Layout, Rect};
use ratatui::style::{Color, Modifier, Style};
use ratatui::text::{Line, Span};
use ratatui::widgets::{Paragraph, Tabs};
use ratatui::Frame;
use super::{App, View};
use crate::state::{ProcId, ProcStatus};
// Palette calibrated (Solarized accents) to stay legible on both light and
// dark terminals. The chrome bars use a light neutral background with dark
// text; the log body keeps the terminal default background so ANSI log colors
// render naturally on either theme. Accent hues are mid-tone so they read on
// the light bar and on both a black and a white body background.
const CHROME_BG: Color = Color::Rgb(238, 232, 213); // light neutral bar
const CHROME_FG: Color = Color::Rgb(60, 70, 72); // dark text on the bar
const MUTED: Color = Color::Rgb(120, 132, 133); // de-emphasized labels
const SERVER: Color = Color::Rgb(38, 139, 210); // blue
const HOST: Color = Color::Rgb(42, 161, 152); // cyan
const VITE: Color = Color::Rgb(211, 54, 130); // magenta
const EVENT: Color = Color::Rgb(181, 137, 0); // amber (omnidev channel)
const OK: Color = Color::Rgb(133, 153, 0); // green (running)
const WARN: Color = Color::Rgb(203, 75, 22); // orange (starting/restarting)
const ERR: Color = Color::Rgb(220, 50, 47); // red (crashed)
/// Style for the header/footer chrome bars.
fn chrome() -> Style {
Style::default().bg(CHROME_BG).fg(CHROME_FG)
}
pub fn draw(f: &mut Frame, app: &App) {
let chunks = Layout::default()
.direction(Direction::Vertical)
.constraints([
Constraint::Length(1), // pod path
Constraint::Length(1), // urls
Constraint::Length(1), // status chips
Constraint::Length(1), // tabs + scroll status
Constraint::Min(1), // body
Constraint::Length(1), // footer
])
.split(f.area());
draw_pod(f, app, chunks[0]);
draw_urls(f, app, chunks[1]);
draw_chips(f, app, chunks[2]);
draw_tabs_row(f, app, chunks[3]);
draw_body(f, app, chunks[4]);
draw_footer(f, chunks[5]);
}
fn draw_pod(f: &mut Frame, app: &App, area: Rect) {
let line = Line::from(vec![
Span::styled(" pod ", Style::default().fg(MUTED)),
Span::raw(app.pod.dir.display().to_string()),
]);
f.render_widget(Paragraph::new(line).style(chrome()), area);
}
fn draw_urls(f: &mut Frame, app: &App, area: Rect) {
let line = Line::from(vec![
Span::styled(" server ", Style::default().fg(MUTED)),
Span::styled(
app.pod.server_display_url(),
Style::default().fg(proc_color(ProcId::Server)),
),
Span::styled(" ui ", Style::default().fg(MUTED)),
Span::styled(
app.pod.vite_display_url(),
Style::default().fg(proc_color(ProcId::Vite)),
),
]);
f.render_widget(Paragraph::new(line).style(chrome()), area);
}
fn draw_chips(f: &mut Frame, app: &App, area: Rect) {
let status = app.shared.lock().unwrap().status.clone();
let mut chips: Vec<Span> = vec![Span::raw(" ")];
for id in ProcId::ALL {
let st = &status[id.idx()];
chips.push(Span::styled(
id.label(),
Style::default()
.fg(proc_color(id))
.add_modifier(Modifier::BOLD),
));
chips.push(Span::raw(" "));
chips.push(Span::styled(
st.short(),
Style::default().fg(status_color(st)),
));
chips.push(Span::raw(" "));
}
f.render_widget(Paragraph::new(Line::from(chips)).style(chrome()), area);
}
fn draw_tabs_row(f: &mut Frame, app: &App, area: Rect) {
// Split the row: tabs on the left, scroll/follow status right-aligned.
let cols = Layout::default()
.direction(Direction::Horizontal)
.constraints([Constraint::Min(0), Constraint::Length(24)])
.split(area);
let entries = [
("server", View::Server, Some(ProcId::Server)),
("host", View::Host, Some(ProcId::Host)),
("vite", View::Vite, Some(ProcId::Vite)),
("all", View::All, None),
];
let selected = entries
.iter()
.position(|(_, v, _)| *v == app.view)
.unwrap_or(3);
let titles: Vec<Line> = entries
.iter()
.map(|(name, _, id)| {
let color = id.map(proc_color).unwrap_or(CHROME_FG);
Line::from(Span::styled(*name, Style::default().fg(color)))
})
.collect();
let tabs = Tabs::new(titles)
.select(selected)
.style(chrome())
.divider(Span::styled("·", Style::default().fg(MUTED)))
.highlight_style(Style::default().add_modifier(Modifier::REVERSED | Modifier::BOLD));
f.render_widget(tabs, cols[0]);
let total = app.line_count();
let status = if app.follow {
format!("{total} ln · follow ")
} else {
format!("{total} ln · ↑{} ", app.scroll_back)
};
f.render_widget(
Paragraph::new(Line::from(Span::styled(status, Style::default().fg(MUTED))))
.alignment(Alignment::Right)
.style(chrome()),
cols[1],
);
}
fn draw_body(f: &mut Frame, app: &App, area: Rect) {
let all_view = app.view == View::All;
let shared = app.shared.lock().unwrap();
let lines: Vec<String> = match app.view {
View::Server => shared.buf(ProcId::Server).iter().cloned().collect(),
View::Host => shared.buf(ProcId::Host).iter().cloned().collect(),
View::Vite => shared.buf(ProcId::Vite).iter().cloned().collect(),
View::All => shared.all.iter().cloned().collect(),
};
drop(shared);
let height = area.height as usize;
let total = lines.len();
let max_back = total.saturating_sub(height);
let back = app.scroll_back.min(max_back);
let end = total.saturating_sub(back);
let start = end.saturating_sub(height);
let rendered: Vec<Line> = lines[start..end]
.iter()
.map(|l| render_line(l, all_view))
.collect();
f.render_widget(Paragraph::new(rendered), area);
}
fn draw_footer(f: &mut Frame, area: Rect) {
let hint = " 1/2/3/0 view · Tab cycle · ↑↓/PgUp/PgDn scroll · f follow · r restart · R backend · c clear · q quit ";
f.render_widget(
Paragraph::new(Line::from(Span::styled(
hint,
Style::default().fg(CHROME_FG),
)))
.style(chrome()),
area,
);
}
/// Turn one stored log line into a styled `Line`. In the combined view the
/// leading `[service]` tag is colored per service and the rest keeps its ANSI
/// colors; per-service panes just pass their ANSI through.
fn render_line(raw: &str, all_view: bool) -> Line<'static> {
if all_view {
if let Some(rest) = raw.strip_prefix('[') {
if let Some(end) = rest.find(']') {
let label = &rest[..end];
let body = &rest[end + 1..];
let mut spans = vec![Span::styled(
format!("[{label}]"),
Style::default()
.fg(label_color(label))
.add_modifier(Modifier::BOLD),
)];
spans.extend(ansi_spans(body));
return Line::from(spans);
}
}
}
Line::from(ansi_spans(raw))
}
/// Parse a single line of possibly-ANSI text into owned spans, falling back to
/// the raw string if it doesn't parse.
fn ansi_spans(s: &str) -> Vec<Span<'static>> {
match s.into_text() {
Ok(text) => text
.lines
.into_iter()
.next()
.map(|l| l.spans)
.unwrap_or_default(),
Err(_) => vec![Span::raw(s.to_string())],
}
}
fn proc_color(id: ProcId) -> Color {
match id {
ProcId::Server => SERVER,
ProcId::Host => HOST,
ProcId::Vite => VITE,
}
}
/// Color for a `[label]` prefix in the combined view — the three services plus
/// the synthetic "omnidev" supervisor channel.
fn label_color(label: &str) -> Color {
match label {
"server" => proc_color(ProcId::Server),
"host" => proc_color(ProcId::Host),
"vite" => proc_color(ProcId::Vite),
"omnidev" => EVENT,
_ => MUTED,
}
}
fn status_color(st: &ProcStatus) -> Color {
match st {
ProcStatus::Running(_) => OK,
ProcStatus::Starting | ProcStatus::Restarting => WARN,
ProcStatus::Crashed => ERR,
ProcStatus::Stopped => VITE,
ProcStatus::Idle => MUTED,
}
}
+228
View File
@@ -0,0 +1,228 @@
//! Daily update check for a git-installed omnigent.
//!
//! Fills a real gap: omnigent's own update notice only works for PyPI-wheel
//! installs and bails on VCS installs. The hot path (`check`) never blocks on
//! the network — it reads a cache and spawns a detached `refresh` when stale.
use std::io::{IsTerminal, Write};
use std::process::{Command, Stdio};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use crate::install::{self, InstallConfig};
use crate::paths;
const STALE_SECS: u64 = 24 * 60 * 60;
const LS_REMOTE_TIMEOUT_SECS: u64 = 5;
/// Volatile update-check state cached between runs.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct CheckCache {
#[serde(default)]
pub last_checked: u64,
#[serde(default)]
pub remote_sha: Option<String>,
#[serde(default)]
pub installed_sha: Option<String>,
/// The remote sha we already prompted about, so a declined update isn't
/// re-nagged until a newer commit lands.
#[serde(default)]
pub last_prompted_sha: Option<String>,
}
impl CheckCache {
pub fn load() -> CheckCache {
let Ok(path) = paths::check_cache_path() else {
return CheckCache::default();
};
std::fs::read_to_string(&path)
.ok()
.and_then(|t| serde_json::from_str(&t).ok())
.unwrap_or_default()
}
pub fn save(&self) -> Result<()> {
let path = paths::check_cache_path()?;
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)
.with_context(|| format!("creating {}", parent.display()))?;
}
let text = serde_json::to_string_pretty(self).context("serializing check cache")?;
std::fs::write(&path, text).with_context(|| format!("writing {}", path.display()))?;
Ok(())
}
}
/// Whether the cache indicates an update the user hasn't already declined.
/// Pure so it can be unit-tested without touching disk or the network.
///
/// `installed` is the best-known installed commit (dist-info first, else the
/// cached `installed_sha`). An update is available when we have a remote sha
/// that differs from what's installed and that we haven't already prompted for.
pub fn update_available(cache: &CheckCache, installed: Option<&str>) -> bool {
let Some(remote) = cache.remote_sha.as_deref() else {
return false;
};
if Some(remote) == installed {
return false;
}
if cache.last_prompted_sha.as_deref() == Some(remote) {
return false;
}
true
}
/// Whether `last_checked` is older than the staleness window.
pub fn is_stale(cache: &CheckCache, now: u64) -> bool {
now.saturating_sub(cache.last_checked) > STALE_SECS
}
fn now_epoch() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_secs())
.unwrap_or(0)
}
/// The remote HEAD sha of `git_ref` in `repo`, via `git ls-remote` (targets the
/// remote, so no local checkout is needed). `None` on any failure/timeout.
pub fn remote_sha(repo: &str, git_ref: &str) -> Option<String> {
// `timeout` isn't portable (absent on macOS by default), so bound the call
// with git's own connect timeout and a wait guard instead.
let mut child = Command::new("git")
.args(["ls-remote", repo, git_ref])
.env("GIT_TERMINAL_PROMPT", "0")
.stdout(Stdio::piped())
.stderr(Stdio::null())
.stdin(Stdio::null())
.spawn()
.ok()?;
let deadline = SystemTime::now() + Duration::from_secs(LS_REMOTE_TIMEOUT_SECS);
loop {
match child.try_wait().ok()? {
Some(_) => break,
None => {
if SystemTime::now() > deadline {
let _ = child.kill();
return None;
}
std::thread::sleep(Duration::from_millis(100));
}
}
}
let output = child.wait_with_output().ok()?;
if !output.status.success() {
return None;
}
let text = String::from_utf8(output.stdout).ok()?;
// First whitespace-delimited token of the first line is the sha.
text.lines()
.next()
.and_then(|l| l.split_whitespace().next())
.map(str::to_string)
}
/// Record the installed sha into the cache (called after install/update).
pub fn set_installed_sha(sha: &str) -> Result<()> {
let mut cache = CheckCache::load();
cache.installed_sha = Some(sha.to_string());
cache.save()
}
/// `refresh` subcommand: hit the network, update `remote_sha` + `last_checked`.
/// Invoked detached by `check`, but also runnable directly.
pub fn refresh() -> Result<()> {
let config = InstallConfig::load()?.unwrap_or_default();
let mut cache = CheckCache::load();
cache.remote_sha = remote_sha(&config.repo, &config.git_ref);
cache.last_checked = now_epoch();
cache.save()
}
/// Best-known installed commit: the tool's dist-info first (authoritative),
/// else the sha we recorded at install time.
fn installed_commit(cache: &CheckCache) -> Option<String> {
install::installed_commit().or_else(|| cache.installed_sha.clone())
}
/// `check` subcommand: the fast hook primitive. Never blocks on the network.
///
/// - Stale cache ⇒ spawn a detached `refresh` and return.
/// - An available update ⇒ notice; on a TTY, prompt and update in the
/// foreground on yes, else record the decline.
/// - `quiet` suppresses the "up to date" path so shell startup stays silent.
pub fn check(quiet: bool) -> Result<()> {
let cache = CheckCache::load();
if is_stale(&cache, now_epoch()) {
spawn_detached_refresh();
// Still evaluate against whatever we already had cached.
}
let installed = installed_commit(&cache);
if !update_available(&cache, installed.as_deref()) {
if !quiet {
println!("omnigent is up to date.");
}
return Ok(());
}
let remote = cache.remote_sha.clone().unwrap_or_default();
let short = |s: &str| s.chars().take(8).collect::<String>();
let installed_desc = installed
.as_deref()
.map(short)
.unwrap_or_else(|| "unknown".to_string());
eprintln!(
"omnigent update available: {}{} (git)",
installed_desc,
short(&remote),
);
// Only prompt on an interactive terminal; scripts/CI just see the notice.
if !(std::io::stdin().is_terminal() && std::io::stderr().is_terminal()) {
return Ok(());
}
if prompt_yes_no("Update omnigent now? [y/N] ") {
install::update()?;
} else {
// Don't re-nag for this same commit.
let mut cache = CheckCache::load();
cache.last_prompted_sha = Some(remote);
cache.save()?;
}
Ok(())
}
/// Prompt on the controlling terminal. Reads from `/dev/tty` so it works even
/// when the hook's stdin is redirected. Any read failure ⇒ treated as "no".
fn prompt_yes_no(prompt: &str) -> bool {
use std::io::BufRead;
let Ok(tty) = std::fs::OpenOptions::new().read(true).open("/dev/tty") else {
return false;
};
eprint!("{prompt}");
let _ = std::io::stderr().flush();
let mut line = String::new();
if std::io::BufReader::new(tty).read_line(&mut line).is_err() {
return false;
}
matches!(line.trim().to_ascii_lowercase().as_str(), "y" | "yes")
}
/// Launch `omnidev refresh` fully detached so shell startup never waits on it.
fn spawn_detached_refresh() {
let Ok(exe) = std::env::current_exe() else {
return;
};
let _ = Command::new(exe)
.arg("refresh")
.stdin(Stdio::null())
.stdout(Stdio::null())
.stderr(Stdio::null())
.spawn();
}
+56
View File
@@ -0,0 +1,56 @@
//! Watches the backend source tree and asks the supervisor to reload on
//! Python changes. Frontend files are deliberately not watched — Vite HMR
//! handles those.
use std::path::Path;
use std::time::Duration;
use anyhow::{Context, Result};
use notify::RecursiveMode;
use notify_debouncer_full::new_debouncer;
use tokio::sync::mpsc;
use crate::supervisor::Cmd;
/// Start watching `omnigent_dir` for `*.py` changes. Coalesced bursts become a
/// single `Cmd::Reload(n)` on `cmd_tx`. The returned debouncer must be kept
/// alive for the watch to persist.
pub fn spawn(
omnigent_dir: &Path,
cmd_tx: mpsc::UnboundedSender<Cmd>,
) -> Result<impl Send + 'static> {
// The debouncer coalesces rapid saves; we still filter to *.py and skip
// caches so editor churn and __pycache__ writes don't trigger reloads.
let mut debouncer = new_debouncer(
Duration::from_millis(500),
None,
move |result: notify_debouncer_full::DebounceEventResult| {
let Ok(events) = result else { return };
let mut changed = 0usize;
for event in &events {
for path in &event.paths {
if is_relevant(path) {
changed += 1;
}
}
}
if changed > 0 {
let _ = cmd_tx.send(Cmd::Reload(changed));
}
},
)
.context("creating file watcher")?;
debouncer
.watch(omnigent_dir, RecursiveMode::Recursive)
.with_context(|| format!("watching {}", omnigent_dir.display()))?;
Ok(debouncer)
}
fn is_relevant(path: &Path) -> bool {
if path.extension().and_then(|e| e.to_str()) != Some("py") {
return false;
}
!path.components().any(|c| c.as_os_str() == "__pycache__")
}
+151
View File
@@ -0,0 +1,151 @@
//! Exercises install-management logic without network or a real install:
//! spec building, config round-trip, and the update-availability/staleness
//! decisions.
use std::sync::{Mutex, MutexGuard};
// These modules reference each other via `crate::`, so declare the whole set at
// the test crate root. Each test target exercises only part of the included
// source, so allow dead code rather than chase per-item warnings.
#[allow(dead_code)]
#[path = "../src/install.rs"]
mod install;
#[allow(dead_code)]
#[path = "../src/paths.rs"]
mod paths;
#[allow(dead_code)]
#[path = "../src/update_check.rs"]
mod update_check;
use install::InstallConfig;
use update_check::{is_stale, update_available, CheckCache};
/// Tests here mutate process-global `XDG_*` env vars; serialize them.
static ENV_LOCK: Mutex<()> = Mutex::new(());
fn lock_env() -> MutexGuard<'static, ()> {
ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner())
}
#[test]
fn spec_default_has_databricks_extra_and_main() {
let c = InstallConfig::default();
assert_eq!(
c.spec(),
"omnigent[databricks] @ git+https://github.com/omnigent-ai/omnigent.git@main"
);
}
#[test]
fn spec_no_extras_is_bare_git_url() {
let c = InstallConfig {
repo: "https://github.com/omnigent-ai/omnigent.git".into(),
git_ref: "main".into(),
extras: vec![],
};
assert_eq!(
c.spec(),
"git+https://github.com/omnigent-ai/omnigent.git@main"
);
}
#[test]
fn spec_reflects_custom_ref_and_extras() {
let c = InstallConfig {
repo: "https://example.com/x.git".into(),
git_ref: "dev".into(),
extras: vec!["databricks".into(), "kubernetes".into()],
};
assert_eq!(
c.spec(),
"omnigent[databricks,kubernetes] @ git+https://example.com/x.git@dev"
);
}
#[test]
fn config_round_trips_through_disk() {
let _guard = lock_env();
let tmp = tempdir();
std::env::set_var("XDG_CONFIG_HOME", &tmp);
let c = InstallConfig {
repo: "https://github.com/omnigent-ai/omnigent.git".into(),
git_ref: "main".into(),
extras: vec!["databricks".into()],
};
c.save().unwrap();
let loaded = InstallConfig::load().unwrap().expect("config present");
assert_eq!(c, loaded);
std::env::remove_var("XDG_CONFIG_HOME");
}
#[test]
fn missing_config_loads_as_none() {
let _guard = lock_env();
let tmp = tempdir();
std::env::set_var("XDG_CONFIG_HOME", &tmp);
assert!(InstallConfig::load().unwrap().is_none());
std::env::remove_var("XDG_CONFIG_HOME");
}
#[test]
fn update_available_logic() {
let cache = CheckCache {
remote_sha: Some("bbbb".into()),
..Default::default()
};
// Remote differs from installed and wasn't prompted → available.
assert!(update_available(&cache, Some("aaaa")));
// Installed already matches remote → not available.
assert!(!update_available(&cache, Some("bbbb")));
// No remote sha known → not available.
assert!(!update_available(&CheckCache::default(), Some("aaaa")));
// Declining a commit (last_prompted_sha == remote) suppresses it.
let declined = CheckCache {
remote_sha: Some("bbbb".into()),
last_prompted_sha: Some("bbbb".into()),
..Default::default()
};
assert!(!update_available(&declined, Some("aaaa")));
}
#[test]
fn staleness_window() {
let now = 1_000_000u64;
let day = 24 * 60 * 60;
let fresh = CheckCache {
last_checked: now - 10,
..Default::default()
};
assert!(!is_stale(&fresh, now));
let old = CheckCache {
last_checked: now - day - 1,
..Default::default()
};
assert!(is_stale(&old, now));
// Never checked (last_checked == 0) → stale.
assert!(is_stale(&CheckCache::default(), now));
}
/// Minimal unique temp dir without pulling a dev-dependency.
fn tempdir() -> std::path::PathBuf {
let base = std::env::temp_dir();
let unique = format!(
"omnidev-mgmt-{}-{}",
std::process::id(),
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos()
);
let dir = base.join(unique);
std::fs::create_dir_all(&dir).unwrap();
dir
}
+244
View File
@@ -0,0 +1,244 @@
//! Exercises the non-TUI setup path: repo detection, pod dir tree, ports.
use std::fs;
use std::sync::{Mutex, MutexGuard};
// The crate is a binary, so pull in the modules under test directly. Each test
// target uses only part of the included source, so allow dead code.
#[allow(dead_code)]
#[path = "../src/lock.rs"]
mod lock;
#[allow(dead_code)]
#[path = "../src/paths.rs"]
mod paths;
#[allow(dead_code)]
#[path = "../src/pod.rs"]
mod pod;
#[allow(dead_code)]
#[path = "../src/ports.rs"]
mod ports;
use pod::Pod;
use ports::Ports;
/// A fake checkout (.git + omnigent/ + web/) is recognized as a root, and a
/// nested subdir resolves up to it.
#[test]
fn finds_repo_root_from_subdir() {
let tmp = tempdir();
fs::create_dir_all(tmp.join(".git")).unwrap();
fs::create_dir_all(tmp.join("omnigent/server")).unwrap();
fs::create_dir_all(tmp.join("web/src")).unwrap();
let root = paths::find_repo_root(&tmp.join("omnigent/server")).unwrap();
assert_eq!(root, tmp.canonicalize().unwrap());
}
/// A VCS root without omnigent/+web/ is rejected.
#[test]
fn rejects_non_omnigent_project() {
let tmp = tempdir();
fs::create_dir_all(tmp.join(".git")).unwrap();
assert!(paths::find_repo_root(&tmp).is_err());
}
/// Two different repo paths get distinct pod dirs; the same path is stable.
#[test]
fn pod_dir_is_per_repo_and_stable() {
let a1 = paths::default_pod_dir(std::path::Path::new("/repos/one")).unwrap();
let a2 = paths::default_pod_dir(std::path::Path::new("/repos/one")).unwrap();
let b = paths::default_pod_dir(std::path::Path::new("/repos/two")).unwrap();
assert_eq!(a1, a2);
assert_ne!(a1, b);
}
/// npm install is needed when node_modules is missing, and when a manifest is
/// newer than it; not needed when node_modules is up to date.
#[test]
fn needs_npm_install_tracks_manifests() {
let repo = tempdir();
let web = repo.join("web");
fs::create_dir_all(&web).unwrap();
fs::write(web.join("package.json"), "{}").unwrap();
let pod = Pod {
repo_root: repo.clone(),
dir: repo.join("pod"),
ports: Ports {
server: 6767,
vite: 5173,
},
vite_host: "127.0.0.1".into(),
trusted_origins: Vec::new(),
};
// No node_modules yet → install needed.
assert!(pod.needs_npm_install());
// Fresh node_modules created after the manifest → up to date.
fs::create_dir_all(web.join("node_modules")).unwrap();
assert!(!pod.needs_npm_install());
// A manifest touched after node_modules → stale, install needed.
// (Sleep briefly so the mtime is observably newer on coarse filesystems.)
std::thread::sleep(std::time::Duration::from_millis(10));
fs::write(web.join("package-lock.json"), "{}").unwrap();
assert!(pod.needs_npm_install());
}
/// Ports probe to bindable values and persist/reuse across calls.
#[test]
fn ports_resolve_and_persist() {
let tmp = tempdir();
let p1 = Ports::resolve(&tmp, None, None).unwrap();
assert_ne!(p1.server, p1.vite);
assert!(tmp.join("pod.toml").is_file());
// A second resolve reuses the persisted pair (both still free).
let p2 = Ports::resolve(&tmp, None, None).unwrap();
assert_eq!(p1.server, p2.server);
assert_eq!(p1.vite, p2.vite);
// Explicit overrides win.
let p3 = Ports::resolve(&tmp, Some(19191), Some(19292)).unwrap();
assert_eq!(p3.server, 19191);
assert_eq!(p3.vite, 19292);
}
/// Two sibling pods under the same cache root never collide, even before their
/// processes have bound anything — the second reads the first's pod.toml.
#[test]
fn sibling_pods_get_distinct_ports() {
let root = tempdir();
let pod_a = root.join("repo-aaaa");
let pod_b = root.join("repo-bbbb");
fs::create_dir_all(&pod_a).unwrap();
fs::create_dir_all(&pod_b).unwrap();
// Pod A resolves and persists first (no process is ever spawned).
let a = Ports::resolve(&pod_a, None, None).unwrap();
// Pod B must avoid A's ports purely from A's persisted claim.
let b = Ports::resolve(&pod_b, None, None).unwrap();
assert_ne!(a.server, b.server);
assert_ne!(a.vite, b.vite);
assert_ne!(a.server, b.vite);
assert_ne!(a.vite, b.server);
}
/// A pod admits one holder; a second acquire fails until the first is dropped.
#[test]
fn pod_lock_is_exclusive() {
let pod = tempdir();
let held = lock::acquire(&pod).expect("first acquire succeeds");
assert!(
lock::acquire(&pod).is_err(),
"second acquire must fail while the first is held"
);
drop(held);
lock::acquire(&pod).expect("acquire succeeds again after release");
}
const ALLOWED_ORIGINS_ENV: &str = "OMNIGENT_WS_ALLOWED_ORIGINS";
/// Tests that read/write the process-global allowlist env var; serialize them.
static ENV_LOCK: Mutex<()> = Mutex::new(());
fn lock_env() -> MutexGuard<'static, ()> {
ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner())
}
/// Run `body` with `OMNIGENT_WS_ALLOWED_ORIGINS` set to `value` (or unset when
/// `None`), restoring the prior value afterward so tests don't leak env state.
fn with_allowlist_env(value: Option<&str>, body: impl FnOnce()) {
let _guard = lock_env();
let prev = std::env::var(ALLOWED_ORIGINS_ENV).ok();
match value {
Some(v) => std::env::set_var(ALLOWED_ORIGINS_ENV, v),
None => std::env::remove_var(ALLOWED_ORIGINS_ENV),
}
body();
match prev {
Some(v) => std::env::set_var(ALLOWED_ORIGINS_ENV, v),
None => std::env::remove_var(ALLOWED_ORIGINS_ENV),
}
}
fn pod_with_trusted(trusted: Vec<String>) -> Pod {
Pod {
repo_root: std::path::PathBuf::from("/repo"),
dir: std::path::PathBuf::from("/pod"),
ports: Ports {
server: 6767,
vite: 5173,
},
vite_host: "0.0.0.0".into(),
trusted_origins: trusted,
}
}
fn allowlist_from_env(pod: &Pod) -> Option<String> {
pod.env()
.into_iter()
.find(|(k, _)| k == ALLOWED_ORIGINS_ENV)
.map(|(_, v)| v)
}
/// With no trusted origins, the pod leaves the allowlist var untouched — even
/// when the developer's shell already exports one (it passes through inherited).
#[test]
fn no_trusted_origins_does_not_set_allowlist() {
with_allowlist_env(Some("https://dev.example.com"), || {
let pod = pod_with_trusted(Vec::new());
assert_eq!(allowlist_from_env(&pod), None);
});
}
/// Trusted origins with no inherited value produce exactly those origins.
#[test]
fn trusted_origins_populate_allowlist() {
with_allowlist_env(None, || {
let pod = pod_with_trusted(vec!["http://192.168.1.42:5173".into()]);
assert_eq!(
allowlist_from_env(&pod).as_deref(),
Some("http://192.168.1.42:5173")
);
});
}
/// A developer's inherited allowlist is preserved and the LAN origins are
/// appended (order-preserving, deduped) rather than clobbered.
#[test]
fn trusted_origins_merge_with_inherited_allowlist() {
with_allowlist_env(
Some("https://dev.example.com, http://192.168.1.42:5173"),
|| {
let pod = pod_with_trusted(vec![
"http://192.168.1.42:5173".into(), // already inherited → not duplicated
"http://10.0.0.9:5173".into(),
]);
assert_eq!(
allowlist_from_env(&pod).as_deref(),
Some("https://dev.example.com,http://192.168.1.42:5173,http://10.0.0.9:5173")
);
},
);
}
/// Minimal unique temp dir without pulling a dev-dependency.
fn tempdir() -> std::path::PathBuf {
let base = std::env::temp_dir();
let unique = format!(
"omnidev-test-{}-{}",
std::process::id(),
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos()
);
let dir = base.join(unique);
std::fs::create_dir_all(&dir).unwrap();
dir
}