chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1 @@
|
||||
"""Project development tooling (lint hooks run from .pre-commit-config.yaml)."""
|
||||
@@ -0,0 +1 @@
|
||||
"""Performance benchmarks (runnable via ``uv run``, not shipped)."""
|
||||
@@ -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.
|
||||
@@ -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.
|
||||
"""
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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}")
|
||||
@@ -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())
|
||||
@@ -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"
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
}
|
||||
@@ -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())
|
||||
@@ -0,0 +1 @@
|
||||
"""Project-specific lint hooks run from .pre-commit-config.yaml."""
|
||||
@@ -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))
|
||||
@@ -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))
|
||||
Generated
+1169
File diff suppressed because it is too large
Load Diff
@@ -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"
|
||||
@@ -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.
|
||||
@@ -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
|
||||
}
|
||||
@@ -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());
|
||||
}
|
||||
}
|
||||
@@ -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 })
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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")));
|
||||
}
|
||||
}
|
||||
@@ -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(())
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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}");
|
||||
}
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
@@ -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();
|
||||
}
|
||||
@@ -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__")
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
Reference in New Issue
Block a user