Files
Shivam Mittal dc00417841 Add WebSocket load test (dev/loadtest/) + run-load-test skill (#3591)
* Add WebSocket load test (dev/loadtest/) + run-load-test skill

Adds a Locust load test that opens N concurrent WebSocket connections to
WS /v1/sessions/updates and holds them open, measuring the server's
WebSocket fan-out (handshake, origin/auth gating, watch-set diffing,
heartbeat) under concurrency — no runner, LLM, or agent turns.

- dev/loadtest/ws_load_test.py: the locustfile (SessionUpdatesUser).
- dev/loadtest/run.py: runner taking server + host + load params, runs
  locust headless, and writes a result set (summary.md, CSV, HTML, config).
- loadtest extra (locust + websocket-client) in pyproject.toml + uv.lock.
- .claude/skills/run-load-test: skill that gathers inputs, runs, and
  explains the latency results.

Co-authored-by: Isaac
Signed-off-by: Shivam Mittal <shivam.mittal@databricks.com>

* Launch locust via sys.executable -m locust in the load-test runner

run.py launched locust as a bare `locust` command, which resolves through
PATH and can pick up a stale/broken locust from a different Python (e.g. a
~/.local 3.10 install missing gevent's zope.event) even when run.py itself
runs under a venv — crashing the run with ModuleNotFoundError before locust
starts. Launch it as `sys.executable -m locust` so it always uses the same
interpreter + site-packages that run.py runs under. Preflight now checks
importlib.util.find_spec (the actual interpreter) instead of shutil.which
(PATH), and --web execs sys.executable too.

Co-authored-by: Isaac
Signed-off-by: Shivam Mittal <shivam.mittal@databricks.com>

* Genericize --mount-prefix docs to reverse-proxy sub-paths

Replace deployment-specific mount-prefix details with a provider-neutral
"behind a reverse proxy at a sub-path" framing (neutral /omnigent example)
across the README, the run-load-test skill, and the run.py / ws_load_test.py
help + docstrings. The --mount-prefix flag itself is unchanged.

Co-authored-by: Isaac
Signed-off-by: Shivam Mittal <shivam.mittal@databricks.com>

* Add runner-level turn load test (real multi-turn conversations, mocked LLM)

turn_load.py drives real agent turns through the runner — the full
POST .../events → server → runner → executor → LLM → stream → idle loop —
under concurrency, with the LLM mocked (zero latency) so the numbers isolate
Omnigent's own per-turn / history-handling overhead. Runs N concurrent
conversations of M sequential turns each on one durable session, so history
grows across the turns (a real long conversation, not N one-shots).

It boots the whole stack itself (server + zero-latency mock LLM + runner) by
reusing the benchmark harness's BenchEnvironment, using the in-process
openai-agents harness — no vendor CLI, no real API key — so it runs from a repo
checkout with no server to point at. Concurrency is asyncio (the runner stack
is async), not Locust. Writes the same summary.md / run_config.json result
format as the WS runner.

Documents both scenarios (WebSocket fan-out vs runner turns) in the README and
the run-load-test skill.

Co-authored-by: Isaac
Signed-off-by: Shivam Mittal <shivam.mittal@databricks.com>

* Address review: fix socket leak, URL/timeout edge cases, double-count; add tests

Copilot review follow-ups on the load-test harness:

- ws_load_test: assign self.ws before the send/recv steps so a post-create
  failure closes the socket instead of leaking it.
- ws_load_test: _ws_url now treats a schemeless host (localhost:8000, which
  Locust accepts) as ws:// rather than emitting an invalid URL.
- ws_load_test: _read_until_snapshot caps each recv to the remaining deadline
  so a late frame can't overrun by a full read timeout.
- ws_load_test: WS_READ_TIMEOUT falls back to the default on a non-numeric
  value instead of raising in on_start.
- run.py: preflight websocket-client as well as locust; rename _fmt_ms ->
  _fmt_num (it also formats Requests/s).
- run.py: _write_summary skips locust's Aggregated row when totaling, which was
  double-counting the headline request/failure counts.
- docs: the scenario reads AUTH_TOKEN from the environment; drop the wrong
  `-e AUTH_TOKEN` locust-flag examples (AUTH_TOKEN=... locust ...).
- tests: add tests/loadtest unit tests for the pure helpers (URL/env/argv
  wiring, summary formatting, timeout parsing) — deterministic, no server boot.

Co-authored-by: Isaac
Signed-off-by: Shivam Mittal <shivam.mittal@databricks.com>

* Redesign as one load test: each user is a real host driving real turns

Collapse the two scenarios (ws_load_test.py + turn_load.py) into a single
load test where each Locust user IS a real omnigent host. Each user spawns a
real `omnigent host` subprocess (unique identity + per-host $HOME so the
host-daemon singleton guard doesn't collide), registers it over the host
tunnel, then creates host-bound sessions and drives real multi-turn
conversations — every turn is a genuine post→idle loop through a runner the
host spawns, with the LLM mocked (zero latency). `-u N` scales the number of
hosts; Locust does the concurrency.

run.py boots the whole stack (server + mock LLM via BenchEnvironment),
registers one agent, sets the mock reply, then runs Locust against it — there
is no --server to pass, since mocking the LLM requires a stack we control. It
reuses the CSV→summary.md machinery (Aggregated-row dedupe kept).

Capacity-limited by design: N hosts × M sessions = N×M real runner processes on
the load box, so it drives genuine end-to-end turns rather than faking the
runner, but does not scale to hundreds on one machine (documented). Removes the
websocket-client dep (no longer used); needs [loadtest,dev,agents-sdk]. README,
skill, and tests updated for the single scenario.

Verified locally: 5 hosts × 3 turns → 133 turns, 0 failures.

Co-authored-by: Isaac
Signed-off-by: Shivam Mittal <shivam.mittal@databricks.com>

---------

Signed-off-by: Shivam Mittal <shivam.mittal@databricks.com>
Co-authored-by: Shivam Mittal <shivam.mittal@databricks.com>
2026-08-03 10:29:45 -07:00
..