diff --git a/.gitignore b/.gitignore index 419a0480..f7b5f20e 100644 --- a/.gitignore +++ b/.gitignore @@ -49,9 +49,10 @@ htmlcov-report experiment_service_pb2.py experiment_service_pb2_grpc.py settings.json -laumch.json +launch.json pkg_main.py _version.py +.wl_opencode.json # Ignore extensions *.onnx diff --git a/AGENTS.md b/AGENTS.md deleted file mode 100644 index b7b63b3e..00000000 --- a/AGENTS.md +++ /dev/null @@ -1,279 +0,0 @@ -# WeightsLab — agent context for users & debugging - -This file is a **portable context for AI coding agents** (Claude Code, etc.) and -the humans driving them. Its job is to let you — or an agent helping you — -**install, configure, run, and debug WeightsLab and Weights Studio** without -having to reverse-engineer the system first. - -It deliberately covers only the two shipped repositories: - -- **weightslab** — the Python backend / core (training instrumentation, data - ledger, gRPC service, the shared proto). -- **weights_studio** — the browser frontend (the studio UI that inspects and - edits a *running* experiment). - -> File/line references drift as the code evolves — treat them as starting points -> and verify against the current source before relying on them. Environment -> variable names and defaults are the most stable thing here; when in doubt the -> authoritative reference is `weightslab/docs/configuration.rst`. - ---- - -## 0. How to load this guide into Claude Code - -So an agent actually *has* this context when you ask it for help: - -- **Working inside a checkout of the repo** (`git clone`): this guide is - committed as `AGENTS.md`; the repo keeps a gitignored `CLAUDE.md` copy of it at - the root so Claude Code auto-loads it every session. Nothing to do. (Claude - Code also loads `~/.claude/CLAUDE.md` global memory and any parent-dir - `CLAUDE.md`.) -- **You only ran `pip install weightslab`** (no checkout — the package lives in - `site-packages`): absolute `@import` paths are fragile because the path - changes per venv/OS. The robust pattern is a small **skill** that locates the - installed file at runtime. Create `~/.claude/skills/weightslab/SKILL.md`: - - ```yaml - --- - name: weightslab - description: Load the WeightsLab debugging & configuration guide when helping with weightslab or weights_studio problems (connection, TLS, env vars, training hangs, rendering). - --- - !`python -c "import weightslab, os; print(open(os.path.join(os.path.dirname(weightslab.__file__), 'AGENTS.md')).read())"` - - Use the guide above to diagnose the user's weightslab / weights_studio issue. - ``` - - Then run `/weightslab` (or let Claude auto-invoke it). This requires the guide - to be **shipped as package data** inside the installed package (see §7); the - copy at the repo root is for contributors working in a checkout. -- **Quick-and-dirty:** copy this file to `~/.claude/WEIGHTSLAB.md` and add - `@~/.claude/WEIGHTSLAB.md` to your `~/.claude/CLAUDE.md`. - ---- - -## 1. What it is and how the pieces connect - -A user wraps their own PyTorch training script with WeightsLab so a running -experiment becomes inspectable/editable; Weights Studio is the UI for that. - -**Wire path (the thing that breaks most often):** - -``` -Browser → weightslab start :8080 (grpc-web → grpc proxy) → Python gRPC servicer → training loop -``` - -- `weightslab start` is a pure-Python HTTP server that serves the bundled SPA - and translates grpc-web (browser) to raw gRPC (backend). No Docker, no Envoy. - If `weightslab start` is not running, the browser has no UI to load. -- The gRPC servicer and the training loop run in the **same process, different - threads**, coordinated by locks in - `weightslab/weightslab/components/global_monitoring.py`. -- One proto is the single source of truth: - `weightslab/weightslab/proto/experiment_service.proto`. - ---- - -## 2. Install & run (the happy path) - -```bash -pip install weightslab -``` - -In your training script: - -```python -import weightslab as wl -# wrap your objects so the studio can see/edit them (see §3), then: -wl.serve(serving_grpc=True, serving_cli=True) # background threads, same process -# ... your training loop ... -wl.keep_serving() # keep the process alive for the UI -``` - -Then start the UI in another terminal and open it in a browser: - -```bash -weightslab start # serves at http://localhost:8080 by default -``` - -Working starting points live in -`weightslab/weightslab/examples/{PyTorch,Lightning,Usecases}//` -(each is a `main.py` + `config.yaml`) — find the closest example and mirror it. - -UI deployment details (port, TLS, certs) are documented in -`weightslab/docs/weights_studio.rst`. TLS is opt-in: run `weightslab se` once, -then `weightslab start --certs`. - ---- - -## 3. The integration API (`import weightslab as wl`) - -How a user's script plugs in. Wrap each training object with -`wl.watch_or_edit(obj, flag=...)`; the returned tracked proxy is registered in -the global ledger (`weightslab/weightslab/backend/ledgers.py`, -`GLOBAL_LEDGER` — the hub everything reads/mutates through). - -- `flag="hyperparameters"` (dict), `flag="model"` (nn.Module, `device=…`), - `flag="optimizer"`, `flag="data"` (Dataset → tracked DataLoader: `loader_name`, - `batch_size`, `is_training`, `collate_fn`, …), `flag="loss"` (a - `reduction="none"` criterion, called with `(preds_raw, targets, batch_ids=ids, - preds=preds)`), `flag="metric"`. - -Conventions that matter for correctness: - -- Wrap the train step in `with guard_training_context:` and eval in - `with guard_testing_context:` (from - `weightslab.components.global_monitoring`). This is how pause/resume and - train/test separation work — **skip it and pause/resume or stats will misbehave.** -- Use `model.get_age()` (steps actually trained; survives checkpoint reloads), - not the raw loop counter. -- `task_type` on the dataset/model selects rendering: `classification`, - `segmentation`, `detection`, `detection_pointcloud`. -- **Hyperparameter handle access:** the registered hyperparameters proxy - supports both `hp.get("lr")` and `hp["lr"]` (subscript == `.get`), and stays - live — reads reflect in-place updates and re-registration. - ---- - -## 4. Configuration (environment variables) - -WeightsLab and Weights Studio are configured almost entirely through env vars. -**Authoritative reference: `weightslab/docs/configuration.rst`.** The high-signal -ones when debugging: - -**Backend (Python):** - -| Variable | Default | Why you touch it | -|---|---|---| -| `WEIGHTSLAB_LOG_LEVEL` | `INFO` | Set `DEBUG` to see what's happening. (`WATCHDOG` level sits between WARNING/ERROR.) | -| `GRPC_BACKEND_HOST` / `GRPC_BACKEND_PORT` | `0.0.0.0` / `50051` | Backend gRPC bind address. | -| `GRPC_TLS_ENABLED` | `0` | TLS on the gRPC socket. Set `1` with `weightslab start --certs`. | -| `GRPC_TLS_REQUIRE_CLIENT_AUTH` | `0` | mTLS. Must match what `weightslab start --certs` presents. | -| `WEIGHTSLAB_CERTS_DIR` | `~/.weightslab-certs` | Where cert files are looked up (single source of truth). | -| `GRPC_AUTH_TOKEN` | *(unset)* | Optional metadata-token auth on top of mTLS. | -| `GRPC_MAX_MESSAGE_BYTES` | `268435456` (256 MB) | Raise it if large tensors/image batches fail. | -| `WEIGHTSLAB_DISABLE_WATCHDOGS` | `0` | Set `1` when debugging with breakpoints (see §5). | -| `GRPC_WATCHDOG_STUCK_SECONDS` | `60` | Lock/RPC stuck threshold + lock-acquire timeout. | - -**Frontend (Weights Studio) — runtime-injected `window.*` globals:** - -| Variable | Default | Why you touch it | -|---|---|---| -| `WS_SERVER_HOST` / `WS_SERVER_PORT` / `WS_SERVER_PROTOCOL` | `localhost` / `8080` / `http` | How the browser reaches the `weightslab start` server. The #1 connection-issue knob. | -| `WS_HISTOGRAM_MAX_BINS` | `512` | Cap on metadata histogram bars. | -| `BB_THUMB_RENDER` | `10` | Max bounding boxes drawn per **thumbnail**, per overlay (GT and PRED capped independently). | -| `BB_MODAL_RENDER` | `100` | Max bounding boxes drawn per **modal** image, per overlay. A `?` button in the modal shows the active limit. | -| `ENABLE_PLOTS` | `1` | `0`/`false` removes the plots board + Signals card and stops plot auto-refresh. | -| `ENABLE_DATA_EXPLORATION` | `1` | `0`/`false` removes the data grid + metadata/details panel and stops the data/metadata auto-refresh. | -| `ENABLE_HYPERPARAMETERS_OPTIMIZATION` | `1` | `0`/`false` removes the Hyperparameters section, makes HP inputs read-only, and stops the HP poll. | -| `ENABLE_AGENT` | `1` | `0`/`false` removes the agent chat bar + history panel and stops the agent health poll. | -| `ENABLE_NOTEBOOK` | `1` | `0`/`false` removes the notebook button (left of the logo) + notebook window. The notebook runs Python in a shared in-process kernel against the live experiment (`df`, model, checkpoints), persisted as `notebook.ipynb` under `root_log_dir`; `>`-prefixed cells ask the agent to propose code. | - -> **VITE_ vs WS_/BB_/ENABLE_:** `VITE_*` variables are baked at **build time** -> (changing them needs a frontend rebuild). `WS_*` / `BB_*` / `ENABLE_*` are -> injected into `config.js` at `weightslab start` time and read as `window.*` -> globals — changing them needs only a restart + browser reload. Each `ENABLE_*` -> defaults to on; set it to `0`/`false`/`no`/`off` to disable. Full reference: -> `weightslab/docs/configuration.rst` (“Feature toggles”). - ---- - -## 5. Troubleshooting — symptom → cause → fix - -This is the core of the guide. Each entry is a real failure mode (several are -distilled from issues hit in development). - -**UI loads but the sample grid is empty / "failed to fetch" / gRPC errors.** -The wire path (§1) is broken somewhere. Check in order: (1) backend actually -serving on `0.0.0.0:50051`; (2) `weightslab start` is running and the browser -can reach it on `:8080`; (3) **TLS mismatch** if using `--certs` — run -`weightslab se` first and export `WEIGHTSLAB_CERTS_DIR`. For local debugging -drop TLS entirely (omit `--certs`; `GRPC_TLS_ENABLED=0`). - -**Changed an env var, restarted, but the UI still uses the old value.** -- `VITE_*` is build-time → you must **rebuild** the frontend, not just restart. -- `WS_*` / `BB_*` / `ENABLE_*` are injected at `weightslab start` time → you - must **restart `weightslab start`** then reload the tab. - -**Sample grid flashes empty cells when auto-refresh fires.** -An auto-refresh (timer or manual) that lands while a `GetDataSamples` grid fetch -is still in flight used to clear the cache mid-render. The fix in -`weights_studio/src/grid_data/gridDataManager.ts` is `isFetchInProgress()`: -refreshes are skipped while a grid fetch is ongoing. If you see this, confirm -you're on a build that has that guard. - -**Detection overlays are slow or unreadably cluttered.** -Dense detection samples can carry hundreds of boxes. Cap rendering with -`BB_THUMB_RENDER` (thumbnails) and `BB_MODAL_RENDER` (modal); each is applied -separately to GT and to predictions. Render-only — no sample data is dropped. - -**Training appears hung; RPCs return `RESOURCE_EXHAUSTED`; server "restarts".** -A watchdog monitors the global rlock and in-flight RPCs. If a lock/RPC is held -longer than `GRPC_WATCHDOG_STUCK_SECONDS` (60s) it's flagged; locks get -interrupted, and after `GRPC_WATCHDOG_RESTART_THRESHOLD` unhealthy polls the -gRPC server restarts. When **debugging with breakpoints** that intentionally -pause longer than that, set `WEIGHTSLAB_DISABLE_WATCHDOGS=1`. If RPCs fail with -`RESOURCE_EXHAUSTED`, a handler couldn't acquire the lock within the window — -something else is holding it; check for a long/blocking train or eval step. - -**Pause/resume doesn't work, or train vs test stats are mixed up.** -The train step isn't wrapped in `guard_training_context` (or eval in -`guard_testing_context`). See §3 — these context managers are how the system -gates and separates phases. - -**Large weights/images fail to transfer.** Raise `GRPC_MAX_MESSAGE_BYTES`. - -**The agent bar says it's unconfigured.** The LLM agent needs a provider: a -local **Ollama** server (`provider: ollama`, available immediately) or **cloud -OpenRouter** initialized from the UI via `/init` (then `/model` to switch, -`/reset` to clear). See `weightslab/docs/weights_studio.rst`. - ---- - -## 6. Where things live (for deeper digging) - -**Backend (`weightslab/weightslab/`):** -- `src.py` — the public verbs (`watch_or_edit`, `serve`, `keep_serving`, - `tag_samples`, `query_*`, decorators) re-exported from `__init__.py`. -- `trainer/services/` — `experiment_service.py` (gRPC servicer) delegating to - `{model,data,agent}_service.py`; `data_image_utils.py` (preview/mask encoding). -- `components/` — `global_monitoring.py` (locks, `guard_*` contexts, pause), - `checkpoint_manager.py`, `evaluation_controller.py`. -- `data/` — `dataframe_manager.py`, `data_samples_with_ops.py`, `sample_stats.py`, - H5 storage (`h5_dataframe_store.py`, `h5_array_store.py`, `array_proxy.py`). -- `backend/` — `ledgers.py` (`GLOBAL_LEDGER`), `logger.py`, `audit_logger.py`, `cli.py`. -- `security/` (`CertAuthManager`), `proto/`, `examples/`, `docs/`. - -**Frontend (`weights_studio/src/`):** -- `main.ts` — bootstrap; builds the grpc-web transport from `WS_SERVER_*`. -- `experiment_service.client.ts` / `experiment_service.ts` — generated client - (regenerate with `npm run generate-proto:data`; do not hand-edit). -- `grid_data/` — grid + modal rendering (`GridCell.ts`, `DataImageService.ts`, - `gridDataManager.ts`, `BboxRenderer.ts`, `SegmentationRenderer.ts`, - `PointCloudViewer.ts`). -- `ui/` — `server.py` (pure-Python HTTP + gRPC-Web proxy), `static/` (bundled SPA), - `utils/` (cert-generation scripts, sync-frontend helper). - -**Docs:** `weightslab/docs/` (Sphinx) — `configuration.rst` (all env vars), -`weights_studio.rst` (studio deploy + agent), `quickstart.rst`, `grpc/`. - ---- - -## 7. For contributors (working in a checkout) - -- **The two repos must sit side by side** (`…/weightslab`, `…/weights_studio`); - proto codegen scripts reach across by relative path. -- **Editing the proto is cross-repo** — do all of: edit - `experiment_service.proto`; regenerate Python stubs from the repo root; run - `npm run generate-proto:data` in weights_studio. Editing one side only leaves - the build broken. -- **Tests:** backend `python -m pytest weightslab/tests/...`; frontend unit - `npm run test` (vitest); E2E/user-simulation Playwright lives in - **weights_studio** (`test:realtime:*`, `test:e2e:*`), not here. -- **CI on a custom branch:** pushes to non-`main`/`dev` branches only run CI when - the commit message contains `[force ci]` (both repos). -- **TLS/auth in the bundled UI** is decided by cert presence under - `WEIGHTSLAB_CERTS_DIR` (single source of truth) — don't hardcode secure/insecure. -- **To make this guide available to pip users**, ship it as package data inside - the installed package (e.g. as `weightslab/weightslab/AGENTS.md`) so the §0 - skill can locate it; keep the root `AGENTS.md` (mirrored as the gitignored - `CLAUDE.md`) as the contributor-facing source. diff --git a/README.md b/README.md index 7297a035..776001c5 100644 --- a/README.md +++ b/README.md @@ -321,6 +321,20 @@ Find our documentation [online](https://grayboxtech.github.io/weightslab/latest/ +
+Agent: chat with your training run (OpenCode) + +
+ +WeightsLab ships two distinct agent surfaces — the backend SDK agent for data-manipulation +queries, and a local [OpenCode](https://opencode.ai)-backed agent with a full bash/file +toolset that can restart training, edit your code, and run recurring `/loop` monitoring +jobs. See the [Agent docs](https://grayboxtech.github.io/weightslab/latest/agent.html) for +how the two connect, how to point either one at a local model, and the full `/loop` +reference. + +
+
diff --git a/agent_config.yaml b/agent_config.yaml index ca1ef1d5..3a3aa2f9 100644 --- a/agent_config.yaml +++ b/agent_config.yaml @@ -4,25 +4,20 @@ # or directly in the config file. # Config file values will override env. variables if both are set. # If cloned, env. variables can be defined in a .env file at the root of the repository. +# +# OpenCode (opencode.ai) is the only supported agent backend: a local OpenCode +# server backs every LLM call. There is no API key here -- the credential +# lives in OpenCode's own config, entered once via `opencode auth login` or +# the Weights Studio landing page's login modal. agent: - # Select the model provider. - # Local: 'ollama' - # Remote: 'openrouter' - # provider: openrouter # Default to OpenRouter if API key is provided, otherwise fallback to local Ollama. This can be overridden by env variable PREFERRED_PROVIDER. + # URL of the local OpenCode server (can also be set as env variable + # OPENCODE_URL). Defaults to http://127.0.0.1:4096. This is the SAME shared + # root env var the frontend reads, so set it once and both sides point at + # one server. + opencode_url: http://127.0.0.1:4096 - # Local Settings - fallback_to_local: false - ollama_model: llama3.2:3b - - # # Remote Model Selection - # Default is a fast flash-class model. The intent-planning task is simple JSON - # generation, so a small/fast model responds in ~2-4s where a 70B model took - # ~15-30s for no accuracy gain. Switch back to a large model here if you see - # accuracy issues (speed/accuracy tradeoff). - openrouter_model: bytedance-seed/seed-2.0-lite # Open router model name (can also be set as env variable OPENROUTER_MODEL). Fast alternatives: openai/gpt-4o-mini, meta-llama/llama-3.1-8b-instruct. Accurate/slow: ~google/gemini-flash-latest - # openrouter_api_key: # Open router API key (can also be set as env variable OPENROUTER_API_KEY) - openrouter_base_url: https://openrouter.ai/api/v1 # Open router base URL (can also be set as env variable OPENROUTER_BASE_URL) - openrouter_request_timeout: 60.0 # Timeout for OpenRouter API requests in seconds (can also be set as env variable OPENROUTER_REQUEST_TIMEOUT) - openrouter_max_tokens: 2048 # Max completion length. OpenRouter reserves max_tokens*price against the key budget BEFORE generating, so an uncapped value can 402 ("more credits, or fewer max_tokens") on a credit/weekly-limited key. Raise only if responses get truncated (env: OPENROUTER_MAX_TOKENS) - openrouter_provider_sort: throughput # Bias OpenRouter's upstream routing to avoid slow providers: 'throughput' | 'latency' | 'price'. Empty string = let OpenRouter choose (env: OPENROUTER_PROVIDER_SORT) - openrouter_structured_output: false # Ask the model for a schema-validated plan directly (skips free-form JSON + regex repair). More reliable, but only works on models whose OpenRouter route supports structured/JSON-schema output (e.g. Gemini, GPT-4o). env: OPENROUTER_STRUCTURED_OUTPUT=1 + # OpenCode model, "providerID/modelID" (can also be set as env variable + # OPENCODE_MODEL). Empty string self-heals to whatever OpenCode's own + # config was last set to, or a configured provider default, falling back + # to the free-tier "opencode/deepseek-v4-flash-free" if neither resolves. + opencode_model: "opencode/deepseek-v4-flash-free" diff --git a/docs/agent.rst b/docs/agent.rst index 9f00ed89..2d062847 100644 --- a/docs/agent.rst +++ b/docs/agent.rst @@ -16,6 +16,52 @@ leaves the process except the prompt text sent to the configured LLM provider. Describe what you want in plain English; the agent translates it into a safe, reviewable plan of dataframe and model operations and executes it. +Two agent surfaces, one OpenCode server +----------------------------------------- + +WeightsLab's agent capability is backed entirely by `OpenCode +`_ — a local ``opencode serve`` process that WeightsLab +starts (or reuses) for you. There is no separate OpenRouter/Ollama +integration to configure: OpenCode itself is the provider layer, and its own +config (``opencode auth login``, or the login modal described below) holds +whatever credentials you use — OpenRouter, Anthropic, a local Ollama model, +anything OpenCode supports. + +That one server backs **two very different agent surfaces**, and knowing +which one you're talking to matters — everything on the rest of this page +describes the first one: + +.. list-table:: + :header-rows: 1 + :widths: 20 40 40 + + * - + - Backend SDK agent + - "Frontend" / OpenCode agent + * - Drives + - The normal query bar and chat-history-panel conversation + (``DataManipulationAgent``, ``weightslab/trainer/services/agent/agent.py``) + - The landing-page chat (pre-experiment) and ``/loop`` (during an experiment) + * - Toolset + - None — every mutating tool (``write``/``edit``/``patch``/``bash``) is + explicitly disabled on every call (``opencode_chat.py``'s + ``_MUTATING_TOOLS``) + - Full toolset — bash, file read/write/edit/patch + * - Memory + - ``self.history``, cleared/summarized by ``/clear`` and ``/compact`` + - An OpenCode session (server-side); cleared/summarized the same way, via + OpenCode's own session delete/summarize endpoints + * - Talks to OpenCode + - In one-shot mode: send a prompt, get text back, no side effects + - Interactively: it can restart training, edit your code, discard/tag + data, run reports + +**During an active experiment, the only way to reach the frontend/OpenCode +agent is** ``/loop`` **from the experiment agent bar.** The landing-page chat +only exists pre-experiment — once you're connected to a running experiment, +that surface is gone, and ``/loop`` (see the "``/loop`` reference" section +near the end of this page) is the sole entry point to the same kind of agent. + What the agent can do --------------------- @@ -171,12 +217,15 @@ dataframe state: .. code-block:: bash - export UTEST_AGENT_PROMPT_EVALUATION=sk-or-... # OpenRouter API key - export UTEST_AGENT_PROMPT_EVALUATION_MODEL=openai/gpt-4o-mini # optional + # Requires a local OpenCode server already running and authenticated + # (opencode has no API-key env var of its own -- see "Initializing the + # agent" below). + export UTEST_AGENT_PROMPT_EVALUATION=1 + export OPENCODE_MODEL=openrouter/anthropic/claude-opus-4.6 # optional pytest weightslab/tests/trainer/services/test_agent_live_prompt_evaluation.py -v Without ``UTEST_AGENT_PROMPT_EVALUATION`` set, the suite logs a note and -skips entirely (it never runs by accident in CI or consumes API credits +skips entirely (it never runs by accident in CI or against a real model unintentionally). A small always-on sanity check for the harness itself (fixture shape, op-runner correctness) still runs regardless. @@ -241,15 +290,47 @@ scenario end-to-end against a real model. Initializing the agent ---------------------- -The agent needs an LLM provider before it can serve requests. Two provider -families are supported: +The agent needs a reachable OpenCode server before it can serve requests -- +OpenCode is the only supported backend. Nothing to install beyond WeightsLab +itself: ``opencode-ai``'s bundled binary ships with the UI's dependencies, and +the UI server (``weightslab/ui/server.py``) starts an ``opencode serve`` child +process on first use, rooted at your experiment directory, tearing it down +when the UI server exits. -- **OpenRouter** — cloud-hosted models (recommended; interactive onboarding in - the UI). -- **Ollama** — local inference, available immediately at backend startup when - configured in ``agent_config.yaml``. +Both agent surfaces (see above) converge on the **same** OpenCode server via +one shared environment variable: -You can initialize it three ways. +.. code-block:: bash + + export OPENCODE_URL=http://127.0.0.1:4096 # or wherever your own `opencode serve` is running + +If ``OPENCODE_URL`` is set and reachable, the UI server adopts it directly +instead of spawning a child; the backend SDK agent reads the same variable +(``agent.py``'s ``_load_config``) — set it once and both sides talk to the one +server, so a model you authenticate once is available everywhere. +``OPENCODE_MODEL`` (or ``agent_config.yaml``'s ``agent.opencode_model``) picks +the default model for the backend SDK agent, as an OpenCode +``providerID/modelID`` string (e.g. ``openrouter/anthropic/claude-opus-4.6``). +Leave it unset to fall back, in order, to: whatever model OpenCode's own +``/config`` was last set to (the model picker's own pick, e.g. from the +Weights Studio landing page), then whichever provider default OpenCode +reports via ``/config/providers``, and finally the free-tier +``opencode/deepseek-v4-flash-free`` if neither of those resolves to anything +(a fresh OpenCode install with no provider credentials configured at all). + +Credentials and provider setup live in OpenCode itself, never in WeightsLab: + +.. code-block:: bash + + opencode auth login # OpenRouter, Anthropic, a local Ollama endpoint, anything OpenCode supports + +or, from the browser, the landing page's login modal drives the same flow +without a terminal. For a fully local setup, point OpenCode's own config at +Ollama (or any other local provider it supports) — WeightsLab needs no +changes on its side; it just asks OpenCode for whichever model you've +selected. + +You can initialize the backend SDK agent three ways. Option 1 — Weights Studio UI (recommended) ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -265,13 +346,12 @@ Type one of these commands into the chat bar: * - Command - Effect * - ``/init`` - - Opens the OpenRouter onboarding modal. Choose **A — Enter OpenRouter API - key** (paste an ``sk-or-…`` key) or **B — Get API key from OpenRouter** - (OAuth flow), then pick a model. On success the placeholder switches to a + - Connects to the OpenCode server (see ``OPENCODE_URL`` above) and lets + you pick a model. On success the placeholder switches to a ready-to-use example query. * - ``/model`` - - Opens the model browser to switch the active OpenRouter model without - re-entering the API key. + - Opens the model browser to switch the active OpenCode model without + reconnecting. * - ``/reset`` - Clears the current connection and returns the agent to the uninitialized state. @@ -294,9 +374,9 @@ the interactive console exposes an ``agent`` verb: .. code-block:: text agent status # Is the agent available? - agent init --api-key sk-or-... --model openai/gpt-4o-mini [--timeout 20] - agent models # List available OpenRouter models - agent model ~google/gemini-flash-latest # Switch model + agent init [--model openrouter/anthropic/claude-opus-4.6] + agent models # List available OpenCode models + agent model openrouter/openai/gpt-5 # Switch model agent reset # Clear the connection agent query # Run a natural-language request query # Shortcut for `agent query` @@ -314,40 +394,33 @@ UI). For example: Option 3 — Startup configuration file ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -To have a provider ready the moment the backend starts (no ``/init`` needed), -configure ``agent_config.yaml`` and/or environment variables. This is the only -way to enable the local Ollama provider. +To have the agent ready the moment the backend starts (no ``/init`` needed), +configure ``agent_config.yaml`` and/or environment variables. .. code-block:: yaml # agent_config.yaml (repo root, package root, cwd, or $AGENT_CONFIG_PATH) agent: - provider: openrouter # or "ollama" - openrouter_model: ~google/gemini-flash-latest - fallback_to_local: false - # Local Ollama alternative: - ollama_model: llama3.2:3b - ollama_host: localhost - ollama_port: 11435 + opencode_url: http://127.0.0.1:4096 + opencode_model: "" # empty = use OpenCode's own configured default .. code-block:: bash - # Prefer secrets via environment variables over YAML. - export OPENROUTER_API_KEY=your_openrouter_key + # Equivalent environment variables (config file wins if both are set). + export OPENCODE_URL=http://127.0.0.1:4096 + export OPENCODE_MODEL=openrouter/anthropic/claude-opus-4.6 See :doc:`configuration` for the full list of agent environment variables, the ``agent_config.yaml`` lookup order, and every supported YAML key. .. note:: - "Available" means the credentials were actually confirmed to work, not - just that a client object was constructed. A key configured via - ``agent_config.yaml``/environment variables (Option 3) is probed once at - backend startup exactly like the ``/init`` UI flow already does, and if a - live query ever gets rejected with 401, the connection is immediately - marked unavailable rather than continuing to report "ready" until the - next restart. If health checks and real requests ever disagree, that's a - bug — the two are kept in sync by design. + "Available" only means a client object was constructed against + ``OPENCODE_URL`` -- OpenCode's own constructor never eagerly connects, so + there is nothing to probe at backend startup the way a cloud API key + needed a connectivity check. Actual unreachability (server not running, or + later restarted) surfaces on the first real query instead, which is + reported through the normal "Internal Agent Error"/reconnect path. Using the agent effectively ---------------------------- @@ -646,6 +719,57 @@ expression it builds for you. comparisons are ``False``) — the query never errors, it just excludes those rows. +``/loop`` reference +---------------------- + +``/loop``, typed into the **experiment agent bar**, is the other agent +surface described at the top of this page: it starts a recurring check-in +against a dedicated OpenCode session — the same kind of session the +landing-page chat uses, with the same full toolset. It never touches the +backend SDK agent directly. + +.. code-block:: text + + /loop 30m Watch the training loss and loss_shape trends; if the run stalls or diverges, pause it and tell me why + /loop list + /loop stop + +- **Syntax**: ``/loop m|h `` to start (minimum interval: 60s), + ``/loop list`` to see running jobs, ``/loop stop `` to cancel one. +- **What it can do**: the loop's OpenCode session is told about the local + ``weightslab`` CLI, reachable over bash against the live training process: + + - ``weightslab pause`` / ``weightslab resume`` — freeze/resume weight updates + - ``weightslab discard `` — discard a sample by id + - ``weightslab agent query ""`` — hands the request to the + **backend SDK agent's** own intent pipeline, e.g. ``weightslab agent + query "discard samples where loss > 5 and tag them hard_examples"``. This + is how the loop reaches back into the database/history: it can't ask the + backend agent directly, but it can drive it through the CLI. + - ``weightslab status`` — a snapshot of hyperparameters/model/training state + + These four are what the loop's system prompt explicitly calls out, but bash + access means any other ``weightslab`` CLI verb is reachable too — e.g. + ``weightslab report`` to generate a narrative report for the loop to read + and act on. It may also read/edit training code directly and attempt to + restart a crashed process via bash — this is best-effort (no supervisor or + PID handoff): it looks for the process, stops it if still running, and + re-launches from whatever it can determine (shell history, a run script, + logs). There is no dedicated restart command. +- **Concurrency cap**: at most 3 loops at once, shared across both chat + surfaces (they hit the same registry). A 4th ``/loop start`` is rejected + with an error rather than silently stopping an older job — stop one first + with ``/loop stop ``. +- **Managing running jobs**: a panel pinned at the top of the chat-history + window lists every running job with a live countdown to its next check-in, + and lets you edit a job's prompt/interval in place or stop it — no need to + remember ``/loop stop `` if the panel is in view. ``/loop list``/``/loop + stop`` also work from the landing-page chat pre-experiment, hitting the + same registry. +- **Persistence**: a loop is tied to the running ``weightslab start`` process, + not the browser tab — it survives a page reload or closed tab, but not a + full restart of the UI server. + Workflow pattern ---------------- diff --git a/docs/configuration.rst b/docs/configuration.rst index cd006a70..0078b45c 100644 --- a/docs/configuration.rst +++ b/docs/configuration.rst @@ -683,10 +683,15 @@ Evaluation Mode ``0`` disables the absolute override and uses the dynamic formula only. -AI / LLM API Keys -~~~~~~~~~~~~~~~~~ +AI / LLM Configuration +~~~~~~~~~~~~~~~~~~~~~~ -These keys are required only when using the agentic data-query features. +OpenCode (`opencode.ai `_) is the agent's only supported +backend -- there is no API key to set here. The credential lives in +OpenCode's own config (``opencode auth login``, or the Weights Studio landing +page's login modal), which can point at OpenRouter, Anthropic, a local Ollama +endpoint, or anything else OpenCode supports; WeightsLab itself only needs to +know which server to talk to. .. list-table:: :header-rows: 1 @@ -695,19 +700,25 @@ These keys are required only when using the agentic data-query features. * - Variable - Default - Description - * - ``OPENROUTER_API_KEY`` + * - ``OPENCODE_URL`` + - ``http://127.0.0.1:4096`` + - URL of the local OpenCode server. Shared with the frontend (Weights + Studio's landing-page chat and ``/loop``) via the same variable, so + set it once and both sides talk to the one server. + * - ``OPENCODE_MODEL`` - *(empty)* - - OpenRouter API key ? required for cloud agent setup in Weights Studio. + - Default model for the backend SDK agent, as an OpenCode + ``providerID/modelID`` string (e.g. + ``openrouter/anthropic/claude-opus-4.6``). Empty falls back, in + order, to OpenCode's own last-picked model, then a configured + provider default, then the free-tier + ``opencode/deepseek-v4-flash-free``. Agent Configuration ~~~~~~~~~~~~~~~~~~~ These variables control how the data-query agent finds its YAML configuration. -The agent supports two provider families: - -- ``ollama`` for local inference -- ``openrouter`` for cloud-hosted models .. list-table:: :header-rows: 1 @@ -744,7 +755,7 @@ Agent Provider Setup ~~~~~~~~~~~~~~~~~~~~ The runtime agent is configured from ``agent_config.yaml`` plus optional -environment variables such as ``OPENROUTER_API_KEY``. +environment variables such as ``OPENCODE_URL``. Supported YAML keys ^^^^^^^^^^^^^^^^^^^ @@ -756,73 +767,35 @@ Supported YAML keys * - Key - Example - Description - * - ``agent.provider`` - - ``ollama`` - - Active provider. Common values: ``ollama`` or ``openrouter``. - * - ``agent.ollama_model`` - - ``llama3.2:3b`` - - Local Ollama model name. - * - ``agent.ollama_host`` - - ``localhost`` - - Ollama host. - * - ``agent.ollama_port`` - - ``11435`` - - Ollama HTTP port used by WeightsLab. - * - ``agent.openrouter_model`` - - ``~google/gemini-flash-latest`` - - Default OpenRouter model. - * - ``agent.openrouter_base_url`` - - ``https://openrouter.ai/api/v1`` - - OpenRouter-compatible base URL. - * - ``agent.openrouter_request_timeout`` - - ``15.0`` - - Request timeout in seconds for OpenRouter calls. - * - ``agent.openrouter_api_key`` - - *(secret)* - - Optional API key in YAML. Prefer environment variables or UI init when possible. - * - ``agent.fallback_to_local`` - - ``false`` - - If enabled, WeightsLab also tries the local Ollama provider as fallback. - -Local Ollama example -^^^^^^^^^^^^^^^^^^^^ - -Use this mode when you want the agent available immediately at backend startup. - -.. code-block:: yaml - - agent: - provider: ollama - ollama_model: llama3.2:3b - ollama_host: localhost - ollama_port: 11435 - fallback_to_local: false - -Operational steps: + * - ``agent.opencode_url`` + - ``http://127.0.0.1:4096`` + - URL of the local OpenCode server. + * - ``agent.opencode_model`` + - ``openrouter/anthropic/claude-opus-4.6`` + - Default model, as an OpenCode ``providerID/modelID`` string. Empty + falls back, in order, to OpenCode's own last-picked model, then a + configured provider default, then the free-tier + ``opencode/deepseek-v4-flash-free``. -1. Install Ollama. -2. Pull a model, for example ``ollama pull llama3.2:3b``. -3. Start the Ollama server. -4. Start WeightsLab. -5. Open Weights Studio and query the agent directly. - -Cloud OpenRouter example -^^^^^^^^^^^^^^^^^^^^^^^^ - -Use this mode when you want hosted models and interactive setup from Weights Studio. +Example +^^^^^^^ .. code-block:: yaml agent: - provider: openrouter - openrouter_model: ~google/gemini-flash-latest - fallback_to_local: false - -Recommended secret handling: + opencode_url: http://127.0.0.1:4096 + opencode_model: "" # empty = self-heal to OpenCode's own default, or "opencode/deepseek-v4-flash-free" -.. code-block:: bash +Setup steps: - export OPENROUTER_API_KEY=your_openrouter_key +1. Have a local OpenCode server running (WeightsLab starts one for you on + first use; see :doc:`agent`), or point ``OPENCODE_URL`` at your own. +2. Authenticate it once: ``opencode auth login`` (or the login modal from + Weights Studio's landing page) -- OpenRouter, Anthropic, a local Ollama + endpoint, anything OpenCode supports. +3. Start WeightsLab. +4. Open Weights Studio and query the agent directly, or type ``/init`` first + to pick a specific model. Weights Studio commands ^^^^^^^^^^^^^^^^^^^^^^^ @@ -830,11 +803,10 @@ Weights Studio commands When using Weights Studio, the agent bar supports these runtime commands: 1. ``/init`` - Opens the OpenRouter onboarding flow. - Users can enter an API key manually or use the OAuth flow, then select a model. + Connects to the OpenCode server and lets you pick a model. 2. ``/model`` - Opens the model browser and switches the active OpenRouter model without - requiring a full reinitialization. + Opens the model browser and switches the active OpenCode model without + reconnecting. 3. ``/reset`` Clears the current runtime connection state and returns the agent to the uninitialized status. @@ -842,13 +814,13 @@ When using Weights Studio, the agent bar supports these runtime commands: Notes ^^^^^ -- The default OpenRouter model is ``~google/gemini-flash-latest``. -- The model browser fetches the available models from OpenRouter using the - configured API key. +- The model browser fetches the available models from the OpenCode server's + own provider catalog. - Connection and model-change actions are recorded in the agent history as log-style entries. -- ``/reset`` clears the current runtime agent state. If your startup config is - local-only and you want that provider back immediately, restart the backend. +- ``/reset`` clears the current runtime agent state. If your startup config + points at a server that's still running, ``/init`` reconnects immediately; + otherwise restart the backend once the server is back. Testing diff --git a/docs/usage/parameters.rst b/docs/usage/parameters.rst index c7fb22f0..254fa0f9 100644 --- a/docs/usage/parameters.rst +++ b/docs/usage/parameters.rst @@ -463,6 +463,9 @@ Audit logging LLM / agent integration (optional) ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +The agent is backed entirely by a local OpenCode server (see :doc:`../agent`) +— there is no API key here; the credential lives in OpenCode's own config. + .. list-table:: :header-rows: 1 :widths: 35 15 50 @@ -470,27 +473,15 @@ LLM / agent integration (optional) * - Variable - Default - Description - * - ``OPENROUTER_API_KEY`` - - *(unset)* - - API key for OpenRouter. Required only when using WeightsLab's - LLM-assisted analysis features. - * - ``OPENROUTER_MODEL`` - - *(unset)* - - Model identifier forwarded to OpenRouter (e.g. - ``"openai/gpt-4o"``). - * - ``OPENROUTER_REQUEST_TIMEOUT`` + * - ``OPENCODE_URL`` + - ``http://127.0.0.1:4096`` + - URL of the local OpenCode server. Shared with the frontend, so set it + once and both sides talk to the one server. + * - ``OPENCODE_MODEL`` - *(unset)* - - Per-request timeout in seconds for OpenRouter calls. - * - ``OPENROUTER_MAX_TOKENS`` - - ``2048`` - - Maximum completion length requested from OpenRouter. OpenRouter - pre-authorizes ``max_tokens × completion_price`` against the key's - remaining budget *before* generating, so leaving this uncapped makes - the model request its full output window and can fail with a ``402`` - ("requires more credits, or fewer max_tokens") on a credit- or - weekly-limited key — even though the model is otherwise usable. The - default is ample for intent planning; raise it only if you see - truncated responses. + - Default model, as an OpenCode ``providerID/modelID`` string (e.g. + ``"openrouter/anthropic/claude-opus-4.6"``). Unset uses OpenCode's own + configured default. Telemetry ~~~~~~~~~~ diff --git a/docs/user_commands.rst b/docs/user_commands.rst index 8eebbdbe..9ef44d1f 100644 --- a/docs/user_commands.rst +++ b/docs/user_commands.rst @@ -487,9 +487,9 @@ sub-verb reference, examples, and setup: see :doc:`agent`. .. code-block:: text agent status - agent init --api-key sk-or-... --model openai/gpt-4o-mini --timeout 20 + agent init --model openrouter/anthropic/claude-opus-4.6 agent models - agent model google/gemini-flash-latest + agent model openrouter/openai/gpt-5 ask tag train samples with loss > 1.2 as goldset Experiment report diff --git a/docs/weights_studio.rst b/docs/weights_studio.rst index b90e5fd1..3f66bbab 100644 --- a/docs/weights_studio.rst +++ b/docs/weights_studio.rst @@ -195,45 +195,41 @@ Agent Usage in Weights Studio ------------------------------ Weights Studio includes an agent bar and an expandable agent history window. -The agent can run with either: +The agent is backed entirely by a local OpenCode server (`opencode.ai +`_) — see :doc:`agent` for the full setup story and the +distinction between this chat-bar agent and the separate ``/loop``/landing-page +OpenCode agent. -- a local Ollama provider configured on the backend -- a cloud OpenRouter provider configured at startup or initialized from the UI - -Local Ollama workflow -~~~~~~~~~~~~~~~~~~~~~ +OpenCode workflow +~~~~~~~~~~~~~~~~~~ -If the backend is configured with ``provider: ollama`` and the Ollama server is -running, the agent is available immediately after backend startup. +WeightsLab starts (or reuses) a local ``opencode serve`` process for you, so +there's normally nothing to configure before the agent is available. If the +backend isn't connected to it yet, Weights Studio shows the agent as +unconfigured and the input placeholder instructs the user to type ``/init``. -Typical local setup: +Typical setup: -1. Start Ollama. +1. Authenticate OpenCode once, if you haven't already: ``opencode auth + login`` (or the landing page's login modal) — OpenRouter, Anthropic, a + local Ollama endpoint, anything OpenCode supports. 2. Start WeightsLab (``wl.serve(serving_grpc=True)``). 3. Start Weights Studio (``weightslab start``). -4. Ask questions in the agent bar. - -Cloud OpenRouter workflow -~~~~~~~~~~~~~~~~~~~~~~~~~ - -If the backend is not initialized with a cloud key yet, Weights Studio shows -the agent as unconfigured and the input placeholder instructs the user to type -``/init``. +4. Ask questions in the agent bar, or type ``/init`` first to pick a specific + model. ``/init`` flow: 1. Type ``/init`` in the agent input. -2. Choose manual API key entry or the OpenRouter OAuth flow. +2. Weights Studio connects to the OpenCode server. 3. Select a model from the available model list. 4. Confirm to initialize the runtime connection. -The default cloud model is ``~google/gemini-flash-latest``. - Available agent commands ~~~~~~~~~~~~~~~~~~~~~~~~ -- ``/init`` — initialize OpenRouter from the UI -- ``/model`` — open the model chooser to switch the active OpenRouter model +- ``/init`` — connect to the OpenCode server from the UI +- ``/model`` — open the model chooser to switch the active OpenCode model - ``/reset`` — clear the current agent runtime connection and status History behavior diff --git a/pyproject.toml b/pyproject.toml index 170624f9..c4f99f08 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -81,10 +81,12 @@ dependencies = [ # Environment variable loading (used by agent service) "python-dotenv>=1,<2", - # Agent service runtime deps (imported by default module graph) + # Agent service runtime deps (imported by default module graph). Only + # langchain-core is needed -- it provides the Runnable/ChatPromptTemplate + # abstraction OpenCodeChat.as_runnable() plugs into. No OpenCode-specific + # SDK package exists; its client is hand-rolled in opencode_chat.py via + # stdlib urllib/threading. "langchain-core>=0.3,<2", - "langchain-ollama>=0.2,<2", - "langchain-openai>=0.2,<2", # Jupyter "ipykernel>=6.29,<7", @@ -152,6 +154,11 @@ weightslab = [ "ui/static/**/*", # Logo used to brand generated experiment reports (weightslab/reporting.py). "assets/**/*", + # weightslab-integration grounding for the landing chat's preset prompts + # and for agents dropped directly into a workspace (see server.py's + # _ensure_workspace_agents_md) -- lives inside the package, not just at the + # repo root, so a `pip install weightslab` ships it too. + "AGENTS.md", ] [tool.setuptools.exclude-package-data] diff --git a/tests/backend/test_cli_additional_unit.py b/tests/backend/test_cli_additional_unit.py index 0cbe947c..aa60576e 100644 --- a/tests/backend/test_cli_additional_unit.py +++ b/tests/backend/test_cli_additional_unit.py @@ -71,29 +71,27 @@ def test_add_tag_uses_sample_id_helper_for_multiple_samples(self): self.assertTrue(result["ok"]) tag_mock.assert_called_once_with(sample_ids=["sample_001", "sample_002", "sample_003"], tag="goldset", mode="add") - def test_agent_init_accepts_api_key_model_and_timeout(self): + def test_agent_init_accepts_a_model(self): agent = MagicMock() - agent.openrouter_request_timeout = 15.0 - agent.openrouter_model = "initial-model" - agent.initialize_with_cloud_key.return_value = (True, "Agent initialized successfully. Ready to help you.") + agent.opencode_model = "initial-model" + agent.initialize_with_cloud_key.return_value = (True, "Agent initialized successfully via OpenCode. Ready to help you.") cli_backend.set_cli_agent(agent) - result = _handle_command("agent init --api-key test-key --model openai/gpt-4o-mini --timeout 22") + result = _handle_command("agent init --model openrouter/openai/gpt-5") self.assertTrue(result["ok"]) - agent.initialize_with_cloud_key.assert_called_once_with("test-key", "openrouter", "openai/gpt-4o-mini") - self.assertEqual(agent.openrouter_request_timeout, 22.0) + agent.initialize_with_cloud_key.assert_called_once_with("", "opencode", "openrouter/openai/gpt-5") def test_agent_model_command_switches_model(self): agent = MagicMock() - agent.openrouter_model = "google/gemini-2.5-flash" + agent.opencode_model = "openrouter/anthropic/claude-opus-4.6" agent.change_model.return_value = (True, "Model switched") cli_backend.set_cli_agent(agent) - result = _handle_command("agent model google/gemini-2.5-flash") + result = _handle_command("agent model openrouter/anthropic/claude-opus-4.6") self.assertTrue(result["ok"]) - agent.change_model.assert_called_once_with("google/gemini-2.5-flash") + agent.change_model.assert_called_once_with("openrouter/anthropic/claude-opus-4.6") def test_agent_query_uses_data_service_when_available(self): mock_response = MagicMock( diff --git a/tests/gRPC/test_grpc_user_actions.py b/tests/gRPC/test_grpc_user_actions.py index 530d8e7d..6f2fcafe 100644 --- a/tests/gRPC/test_grpc_user_actions.py +++ b/tests/gRPC/test_grpc_user_actions.py @@ -290,7 +290,7 @@ def _make_real_data_service(self): # vectorized path and doesn't need it, but GetDataSamples does). ds._data_executor = ThreadPoolExecutor(max_workers=2) ds._agent = MagicMock() - ds._agent.is_ollama_available.return_value = True + ds._agent.is_available.return_value = True ds.audit_logger = MagicMock() return ds, df_manager diff --git a/tests/test_opencode_process.py b/tests/test_opencode_process.py new file mode 100644 index 00000000..1098cd59 --- /dev/null +++ b/tests/test_opencode_process.py @@ -0,0 +1,187 @@ +"""Tests for weightslab/opencode_process.py -- the cross-process discovery/ +spawn handshake that lets the backend SDK agent (agent.py's +DataManipulationAgent) and the UI server's _OpencodeSession (server.py, +backing the browser landing-page chat and /loop jobs) converge on ONE +OpenCode server for a given workspace directory, regardless of which one +needs it first. + +CI has no real `opencode` binary, so the real-spawn tests point +resolve_opencode_argv at a tiny stand-in HTTP server (started via +`python -c`, same pattern tests/ui/test_server_agent.py already uses) that +answers /global/health the way the real OpenCode server does -- everything +else here is exercised via that real subprocess + real lock-file I/O in a +temp directory, not mocked away. +""" + +import os +import sys +import tempfile +import unittest +from unittest.mock import patch + +from weightslab import opencode_process + + +_FAKE_OPENCODE_SRC = r""" +import sys, json +from http.server import BaseHTTPRequestHandler, HTTPServer + +def _port(): + for i, a in enumerate(sys.argv): + if a == "--port": + return int(sys.argv[i + 1]) + return 4096 + +class H(BaseHTTPRequestHandler): + def log_message(self, *a): + pass + def do_GET(self): + if self.path == "/global/health": + body = json.dumps({"version": "0.0.0-fake"}).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + else: + self.send_response(404) + self.end_headers() + +HTTPServer(("127.0.0.1", _port()), H).serve_forever() +""" + +_FAKE_ARGV = [sys.executable, "-c", _FAKE_OPENCODE_SRC] + + +class TestLockFileRoundtrip(unittest.TestCase): + def setUp(self): + self.tmp = tempfile.mkdtemp() + + def test_absent_lock_reads_as_none(self): + self.assertIsNone(opencode_process.read_lock(self.tmp)) + + def test_write_then_read_roundtrips(self): + opencode_process.write_lock(self.tmp, "http://127.0.0.1:9999", pid=1234) + lock = opencode_process.read_lock(self.tmp) + self.assertEqual(lock["url"], "http://127.0.0.1:9999") + self.assertEqual(lock["pid"], 1234) + + def test_malformed_lock_file_reads_as_none_not_an_exception(self): + with open(opencode_process.lock_path(self.tmp), "w") as f: + f.write("{not json") + self.assertIsNone(opencode_process.read_lock(self.tmp)) + + def test_lock_file_with_no_url_reads_as_none(self): + with open(opencode_process.lock_path(self.tmp), "w") as f: + f.write('{"pid": 1}') + self.assertIsNone(opencode_process.read_lock(self.tmp)) + + +class TestResolveOrSpawnUnit(unittest.TestCase): + """The precedence chain (env > lockfile > spawn), mocked so it runs in + milliseconds -- the real-subprocess path is covered separately below.""" + + def setUp(self): + self.tmp = tempfile.mkdtemp() + os.environ.pop("OPENCODE_URL", None) + + def test_healthy_env_var_wins_outright(self): + os.environ["OPENCODE_URL"] = "http://127.0.0.1:1111" + try: + with patch.object(opencode_process, "opencode_healthy", return_value=True): + result = opencode_process.resolve_or_spawn_opencode(self.tmp) + finally: + del os.environ["OPENCODE_URL"] + self.assertEqual(result, {"ok": True, "url": "http://127.0.0.1:1111", "source": "env"}) + + def test_unhealthy_env_var_is_ignored_in_favor_of_the_lockfile(self): + os.environ["OPENCODE_URL"] = "http://127.0.0.1:1111" + opencode_process.write_lock(self.tmp, "http://127.0.0.1:2222") + try: + with patch.object(opencode_process, "opencode_healthy", + side_effect=lambda url, timeout=1.5: url == "http://127.0.0.1:2222"): + result = opencode_process.resolve_or_spawn_opencode(self.tmp) + finally: + del os.environ["OPENCODE_URL"] + self.assertEqual(result, {"ok": True, "url": "http://127.0.0.1:2222", "source": "lockfile"}) + + def test_healthy_lockfile_is_adopted_without_spawning(self): + opencode_process.write_lock(self.tmp, "http://127.0.0.1:3333") + with patch.object(opencode_process, "opencode_healthy", return_value=True), \ + patch("subprocess.Popen") as popen: + result = opencode_process.resolve_or_spawn_opencode(self.tmp) + popen.assert_not_called() + self.assertEqual(result, {"ok": True, "url": "http://127.0.0.1:3333", "source": "lockfile"}) + + def test_stale_lockfile_is_ignored_and_a_fresh_server_is_spawned(self): + # The process the file names is gone -- health check on ITS url fails, + # but the newly-spawned one's succeeds. + opencode_process.write_lock(self.tmp, "http://127.0.0.1:4444") + with patch.object(opencode_process, "opencode_healthy", + side_effect=lambda url, timeout=1.5: url != "http://127.0.0.1:4444"), \ + patch.object(opencode_process, "resolve_opencode_argv", return_value=_FAKE_ARGV), \ + patch.object(opencode_process, "free_port", return_value=15000): + result = opencode_process.resolve_or_spawn_opencode(self.tmp) + self.assertTrue(result["ok"], result) + self.assertEqual(result["source"], "spawned") + # The stale entry was overwritten with the freshly-spawned server. + self.assertEqual(opencode_process.read_lock(self.tmp)["url"], result["url"]) + + def test_no_opencode_or_npx_available_reports_a_clear_error(self): + with patch.object(opencode_process, "resolve_opencode_argv", return_value=None): + result = opencode_process.resolve_or_spawn_opencode(self.tmp) + self.assertFalse(result["ok"]) + self.assertIn("opencode", result["error"]) + + +class TestResolveOrSpawnRealSubprocess(unittest.TestCase): + """Exercises the actual subprocess.Popen + health-poll + lock-file-write + path against a real (fake-OpenCode) child process, not a mock.""" + + def setUp(self): + self.tmp = tempfile.mkdtemp() + os.environ.pop("OPENCODE_URL", None) + self._timeout_patch = patch.object(opencode_process, "OPENCODE_START_TIMEOUT", 5.0) + self._timeout_patch.start() + + def tearDown(self): + self._timeout_patch.stop() + + def test_first_caller_spawns_and_writes_the_lock_file(self): + with patch.object(opencode_process, "resolve_opencode_argv", return_value=_FAKE_ARGV): + result = opencode_process.resolve_or_spawn_opencode(self.tmp, origin="http://localhost:5173") + self.assertTrue(result["ok"], result) + self.assertEqual(result["source"], "spawned") + lock = opencode_process.read_lock(self.tmp) + self.assertEqual(lock["url"], result["url"]) + self.assertTrue(opencode_process.opencode_healthy(result["url"])) + + def test_second_caller_for_the_same_workspace_adopts_instead_of_spawning(self): + """The actual cross-process handshake this feature exists for: call + it once (simulating whichever side -- backend agent or UI server -- + happens to start first), then again for the SAME workspace_dir + (simulating the other side starting later) and confirm the second + call adopts the first call's server rather than spawning a second + one.""" + with patch.object(opencode_process, "resolve_opencode_argv", return_value=_FAKE_ARGV): + first = opencode_process.resolve_or_spawn_opencode(self.tmp) + self.assertEqual(first["source"], "spawned") + + with patch("subprocess.Popen") as popen: + second = opencode_process.resolve_or_spawn_opencode(self.tmp) + popen.assert_not_called() + + self.assertEqual(second["source"], "lockfile") + self.assertEqual(second["url"], first["url"]) + + def test_never_becoming_healthy_times_out_and_reports_an_error(self): + hanging_argv = [sys.executable, "-c", "import time; time.sleep(60)"] + with patch.object(opencode_process, "resolve_opencode_argv", return_value=hanging_argv): + result = opencode_process.resolve_or_spawn_opencode(self.tmp) + self.assertFalse(result["ok"]) + self.assertIn("did not come up", result["error"]) + self.assertIsNone(opencode_process.read_lock(self.tmp)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_opencode_shared_server_integration.py b/tests/test_opencode_shared_server_integration.py new file mode 100644 index 00000000..d62ee742 --- /dev/null +++ b/tests/test_opencode_shared_server_integration.py @@ -0,0 +1,131 @@ +"""End-to-end proof that the backend SDK agent (OpenCodeChat, reached via the +gRPC query bar) and the UI server's _OpencodeSession (backing the browser +landing-page chat and /loop jobs) converge on ONE OpenCode server for a +shared workspace directory -- regardless of which one needs a server first. + +Each side's own half of this handshake already has focused unit coverage +(tests/test_opencode_process.py for the shared resolve_or_spawn_opencode +precedence chain; tests/ui/test_server_agent.py and +tests/trainer/services/test_opencode_chat.py for each side's own call into +it). This file instead drives BOTH real classes together against one real +(fake-OpenCode) subprocess, proving the actual scenario end to end rather +than trusting that the separately-tested pieces compose correctly. +""" + +import sys +import tempfile +import unittest +from unittest.mock import patch + +from weightslab.trainer.services.agent.opencode_chat import OpenCodeChat +from weightslab.ui import server as ui_server + +_FAKE_OPENCODE_SRC = r""" +import sys, json +from http.server import BaseHTTPRequestHandler, HTTPServer + +def _port(): + for i, a in enumerate(sys.argv): + if a == "--port": + return int(sys.argv[i + 1]) + return 4096 + +class H(BaseHTTPRequestHandler): + def log_message(self, *a): + pass + def do_GET(self): + if self.path == "/global/health": + body = json.dumps({"version": "0.0.0-fake"}).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + else: + self.send_response(404) + self.end_headers() + +HTTPServer(("127.0.0.1", _port()), H).serve_forever() +""" + +_FAKE_ARGV = [sys.executable, "-c", _FAKE_OPENCODE_SRC] + +# The real opencode_healthy, forced to treat OpenCode's literal default +# address as dead regardless of the actual machine's state -- a real, +# unrelated `opencode serve` left running on the default port (easy to +# accumulate: nothing in this codebase kills one automatically once +# started, confirmed live more than once during development) would +# otherwise make _ensure_reachable's "already healthy, leave it alone" +# branch fire for real here, which is correct behaviour but defeats the +# point of THIS test -- it wants to force the "was dead, needed +# resolving" path deterministically. Every other address still gets a +# real health check. +import weightslab.opencode_process as _ocp # noqa: E402 + +_real_opencode_healthy = _ocp.opencode_healthy + + +def _healthy_except_bare_default(url: str, timeout: float = 1.5) -> bool: + if url.rstrip("/") == "http://127.0.0.1:4096": + return False + return _real_opencode_healthy(url, timeout=timeout) + + +class TestBackendAgentStartsFirst(unittest.TestCase): + """Order (a) from the user's own description: the backend SDK agent + needs a server before `weightslab start` ever calls /agent-server/start + for the same experiment directory.""" + + def setUp(self): + self.tmp = tempfile.mkdtemp() + + def test_ui_server_adopts_the_backend_agents_server_instead_of_spawning(self): + with patch("weightslab.opencode_process.opencode_healthy", side_effect=_healthy_except_bare_default): + chat = OpenCodeChat("http://127.0.0.1:4096", workspace_dir=self.tmp, url_is_explicit=False) + with patch("weightslab.opencode_process.resolve_opencode_argv", return_value=_FAKE_ARGV): + chat._ensure_reachable() # the backend agent spawns first + backend_url = chat.base_url + self.assertNotEqual(backend_url, "http://127.0.0.1:4096", "the dead default should have been replaced") + + session = ui_server._OpencodeSession() + try: + with patch.object(ui_server, "_resolve_opencode_argv") as argv_mock, \ + patch.object(ui_server, "_opencode_healthy", side_effect=_healthy_except_bare_default): + result = session.ensure(self.tmp, "http://localhost:5173") + finally: + session.shutdown() + + argv_mock.assert_not_called() # no second server spawned + self.assertTrue(result["ok"], result) + self.assertEqual(result["url"], backend_url) + self.assertEqual(result.get("adopted"), "lockfile") + + +class TestUiServerStartsFirst(unittest.TestCase): + """Order (b): `weightslab start` (the UI server) needs a server first, + and the backend SDK agent's first query comes along afterward.""" + + def setUp(self): + self.tmp = tempfile.mkdtemp() + + def test_backend_agent_adopts_the_ui_servers_server_instead_of_spawning(self): + with patch("weightslab.opencode_process.opencode_healthy", side_effect=_healthy_except_bare_default), \ + patch.object(ui_server, "_opencode_healthy", side_effect=_healthy_except_bare_default): + session = ui_server._OpencodeSession() + try: + with patch.object(ui_server, "_resolve_opencode_argv", return_value=_FAKE_ARGV): + started = session.ensure(self.tmp, "http://localhost:5173") + self.assertTrue(started["ok"], started) + + chat = OpenCodeChat("http://127.0.0.1:4096", workspace_dir=self.tmp, url_is_explicit=False) + with patch("weightslab.opencode_process.resolve_opencode_argv") as argv_mock: + chat._ensure_reachable() + finally: + session.shutdown() + + argv_mock.assert_not_called() # no second server spawned + self.assertEqual(chat.base_url, started["url"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/trainer/services/test_agent_live_prompt_evaluation.py b/tests/trainer/services/test_agent_live_prompt_evaluation.py index ceaedeaa..1139289d 100644 --- a/tests/trainer/services/test_agent_live_prompt_evaluation.py +++ b/tests/trainer/services/test_agent_live_prompt_evaluation.py @@ -2,30 +2,25 @@ Live-LLM evaluation of the Data Manipulation Agent against a batch of realistic user prompts. -This suite is OPT-IN: every test calls a REAL LLM through OpenRouter -(consuming API credits and real wall-clock time), so it is skipped entirely -unless an OpenRouter API key can be resolved. The key/model are resolved, in -priority order, from: - - 1. The dedicated ``UTEST_AGENT_PROMPT_EVALUATION`` / - ``UTEST_AGENT_PROMPT_EVALUATION_MODEL`` env vars (explicit opt-in). - 2. The standard ``OPENROUTER_API_KEY`` / ``OPENROUTER_MODEL`` env vars — - including any loaded from a repo ``.env`` file — so the same credentials - the running agent uses also drive this suite with no extra setup. - -So any of these work from the command line: - - # reuse your existing OpenRouter config (.env or exported env var) +This suite is OPT-IN: every test calls a REAL LLM through a local OpenCode +server (opencode.ai) (consuming real wall-clock time), so it is skipped +entirely unless explicitly turned on via the ``UTEST_AGENT_PROMPT_EVALUATION`` +env var (any non-empty value). Unlike the old OpenRouter-backed version of +this suite, there is no API key to resolve here -- OpenCode's credential +lives in its own config (``opencode auth login``), not in an env var -- so +opting in just means "run me", and ``OPENCODE_URL``/``OPENCODE_MODEL`` (the +SAME vars the agent itself reads) pick which server/model to run against. + +So this works from the command line, with a local OpenCode server already +running and authenticated: + + set UTEST_AGENT_PROMPT_EVALUATION=1 pytest weightslab/tests/trainer/services/test_agent_live_prompt_evaluation.py -v - # or pass explicitly for this run only (PowerShell) - $env:OPENROUTER_API_KEY="sk-or-..."; $env:OPENROUTER_MODEL="google/gemini-flash-latest"; pytest ... -v + # or target a non-default server/model for this run only (PowerShell) + $env:UTEST_AGENT_PROMPT_EVALUATION="1"; $env:OPENCODE_URL="http://127.0.0.1:4096"; $env:OPENCODE_MODEL="openrouter/anthropic/claude-opus-4.6"; pytest ... -v - # or the dedicated opt-in vars (cmd.exe) - set UTEST_AGENT_PROMPT_EVALUATION=sk-or-... - pytest ... -v - -The model defaults to the agent's own default OpenRouter model when unset. +The model defaults to the OpenCode server's own configured default when unset. Each test reproduces a specific, previously-reported bug/scenario and verifies the agent's plan, once executed against a realistic synthetic @@ -58,13 +53,16 @@ logger = logging.getLogger(__name__) -def _resolve_live_credentials() -> "tuple[str, str | None]": - """Resolve the OpenRouter (key, model) for the live suite. +def _resolve_live_opencode_config() -> "tuple[bool, str]": + """Resolve whether to run the live suite, and which OpenCode model to + request. - Mirrors the agent's own config loading: pull in any repo ``.env`` first, - then prefer the dedicated UTEST_* opt-in vars, falling back to the standard - OPENROUTER_* vars so the same credentials the agent runs on also drive this - suite without duplicating them. + OpenCode has no API-key concept (its credential lives in the OpenCode + server's own config, entered once via ``opencode auth login``), so opting + in is just a plain on/off switch -- ``UTEST_AGENT_PROMPT_EVALUATION`` set + to any non-empty value. ``OPENCODE_URL``/``OPENCODE_MODEL`` (the SAME env + vars the agent itself reads) pick which server/model to run against; any + repo ``.env`` is loaded first so they can live there too. """ if load_dotenv is not None: # weightslab/tests/trainer/services/ -> parents[4] = repo root, @@ -75,26 +73,19 @@ def _resolve_live_credentials() -> "tuple[str, str | None]": load_dotenv(dotenv_path=candidate, override=False) load_dotenv(override=False) - key = ( - os.environ.get("UTEST_AGENT_PROMPT_EVALUATION", "").strip() - or os.environ.get("OPENROUTER_API_KEY", "").strip() - ) - model = ( - os.environ.get("UTEST_AGENT_PROMPT_EVALUATION_MODEL", "").strip() - or os.environ.get("OPENROUTER_MODEL", "").strip() - or None - ) - return key, model + run_live = bool(os.environ.get("UTEST_AGENT_PROMPT_EVALUATION", "").strip()) + model = os.environ.get("OPENCODE_MODEL", "").strip() or None + return run_live, model -API_KEY, MODEL = _resolve_live_credentials() +RUN_LIVE, MODEL = _resolve_live_opencode_config() -if not API_KEY: +if not RUN_LIVE: logger.info( - "[test_agent_live_prompt_evaluation] No OpenRouter key found " - "(checked UTEST_AGENT_PROMPT_EVALUATION and OPENROUTER_API_KEY, incl. .env) -- " - "skipping live-LLM agent prompt evaluation tests. Set one of those (and optionally " - "OPENROUTER_MODEL) to run this suite against a real model." + "[test_agent_live_prompt_evaluation] UTEST_AGENT_PROMPT_EVALUATION not set -- " + "skipping live-LLM agent prompt evaluation tests. Set it to any value, with a " + "local OpenCode server running and authenticated (see OPENCODE_URL/OPENCODE_MODEL), " + "to run this suite against a real model." ) @@ -146,7 +137,7 @@ def _make_live_agent(df: pd.DataFrame, exp_ctx=None) -> DataManipulationAgent: # for the data-only tests (no model registered). ctx = SimpleNamespace(_all_datasets_df=df, _ctx=exp_ctx) agent = DataManipulationAgent(ctx) - ok, message = agent.initialize_with_cloud_key(API_KEY, "openrouter", MODEL or agent.openrouter_model) + ok, message = agent.initialize_with_cloud_key("", "opencode", MODEL) if not ok: raise RuntimeError(f"Failed to initialize live agent for testing: {message}") return agent @@ -246,7 +237,7 @@ def _run_ops(df: pd.DataFrame, ops: list, model_service=None) -> "tuple[pd.DataF return df, messages -@unittest.skipUnless(API_KEY, "UTEST_AGENT_PROMPT_EVALUATION not set; skipping live-LLM agent evaluation") +@unittest.skipUnless(RUN_LIVE, "UTEST_AGENT_PROMPT_EVALUATION not set; skipping live-LLM agent evaluation") class TestAgentLivePromptEvaluation(unittest.TestCase): """Runs a battery of realistic user prompts against a REAL LLM and verifies the resulting dataframe/message state. Each test is a @@ -403,7 +394,7 @@ def test_reset_view_is_recognized(self): self.assertTrue(any(op.get("params", {}).get("__agent_reset__") for op in ops), ops) -@unittest.skipUnless(API_KEY, "UTEST_AGENT_PROMPT_EVALUATION not set; skipping live-LLM agent evaluation") +@unittest.skipUnless(RUN_LIVE, "UTEST_AGENT_PROMPT_EVALUATION not set; skipping live-LLM agent evaluation") class TestAgentRstDocumentedPrompts(unittest.TestCase): """ One test per example prompt listed in docs/agent.rst's "Example prompts diff --git a/tests/trainer/services/test_agent_model_and_safety_unit.py b/tests/trainer/services/test_agent_model_and_safety_unit.py index 6f475cda..93e4d50d 100644 --- a/tests/trainer/services/test_agent_model_and_safety_unit.py +++ b/tests/trainer/services/test_agent_model_and_safety_unit.py @@ -16,13 +16,9 @@ def _install_agent_dependency_stubs(): stubs = { - "langchain_ollama": types.ModuleType("langchain_ollama"), - "langchain_openai": types.ModuleType("langchain_openai"), "langchain_core": types.ModuleType("langchain_core"), "langchain_core.prompts": types.ModuleType("langchain_core.prompts"), } - stubs["langchain_ollama"].ChatOllama = object - stubs["langchain_openai"].ChatOpenAI = object stubs["langchain_core.prompts"].ChatPromptTemplate = object return stubs @@ -37,9 +33,7 @@ def _make_agent(df=None): # `_ctx=None` means `_setup_model_schema` bails out early (no live model), # matching how a standalone agent behaves before any model is registered. ctx = SimpleNamespace(_all_datasets_df=df, _ctx=None) - - with mock.patch.object(agent_mod, "ChatOpenAI", None), mock.patch.object(agent_mod, "ChatOllama", None): - agent = agent_mod.DataManipulationAgent(ctx) + agent = agent_mod.DataManipulationAgent(ctx) return agent_mod, agent @@ -134,8 +128,6 @@ def fake_try_query_provider(provider, instruction, system_prompt): return None agent._try_query_provider = fake_try_query_provider - agent.preferred_provider = "openrouter" - agent.fallback_to_local = False agent.query("keep only validation or test samples") @@ -168,8 +160,6 @@ def fake_try_query_provider(provider, instruction, system_prompt): return next(it) agent._try_query_provider = fake_try_query_provider - agent.preferred_provider = "openrouter" - agent.fallback_to_local = False return agent, calls def test_history_starts_empty(self): @@ -231,64 +221,17 @@ def test_history_unaffected_by_a_failed_query(self): self.assertEqual(agent.history, []) -class TestStartupProviderVerification(unittest.TestCase): +class TestQueryAuthFailureHandling(unittest.TestCase): """Reported bug: CheckAgentHealth said the agent was available, but a - real query then failed with "401 Unauthorized". Root cause: - is_available() only checks that a provider CLIENT OBJECT exists, not - that its credentials were ever confirmed to work -- true for any key - loaded from agent_config.yaml/env vars, which (unlike the /init UI flow) - never runs a connectivity check. _verify_startup_providers() now probes - once at construction time, and a 401 detected during a real query - invalidates the cached connection so is_available() reflects reality - immediately instead of staying stale until the process restarts.""" - - def test_startup_verification_disables_a_bad_openrouter_key(self): - agent_mod, agent = _make_agent() - - class _FailingChatModel: - def __init__(self, *a, **kw): pass - def invoke(self, prompt): - raise RuntimeError("401 Unauthorized") - - agent.openrouter_api_key = "bad-key" - with mock.patch.object(agent_mod, "ChatOpenAI", _FailingChatModel): - agent._setup_providers() - agent._verify_startup_providers() - - self.assertIsNone(agent.chain_openrouter) - self.assertFalse(agent.is_available()) - - def test_startup_verification_keeps_a_good_openrouter_key(self): - agent_mod, agent = _make_agent() - - class _OkChatModel: - def __init__(self, *a, **kw): pass - def invoke(self, prompt): - return SimpleNamespace(content="OK") - - agent.openrouter_api_key = "good-key" - with mock.patch.object(agent_mod, "ChatOpenAI", _OkChatModel): - agent._setup_providers() - agent._verify_startup_providers() - - self.assertIsNotNone(agent.chain_openrouter) - self.assertTrue(agent.is_available()) - - def test_startup_verification_skips_probe_when_no_chain_was_built(self): - # No openrouter_api_key configured at all -> chain_openrouter stays - # None -> nothing to probe, must not raise. - _, agent = _make_agent() - agent.chain_openrouter = None - - agent._verify_startup_providers() # should be a no-op, not raise - - self.assertIsNone(agent.chain_openrouter) + real query then failed with "401 Unauthorized". A 401 detected during a + real query invalidates the cached OpenCode connection so is_available() + reflects reality immediately instead of staying stale until the process + restarts; a non-auth (e.g. transient/timeout) failure must NOT do that, + since the connection might still be perfectly valid.""" def test_401_during_query_invalidates_the_cached_connection(self): _, agent = _make_agent() - agent.chain_openrouter = object() # simulates an already-"available" cached client - agent.preferred_provider = "openrouter" - agent.fallback_to_local = False + agent.chain_opencode = object() # simulates an already-"available" cached client def fake_try_query_provider(provider, instruction, system_prompt): agent._last_query_error = RuntimeError("401 Unauthorized") @@ -299,7 +242,7 @@ def fake_try_query_provider(provider, instruction, system_prompt): self.assertTrue(agent.is_available()) # stale "available" before the query result = agent.query("do something") - self.assertIsNone(agent.chain_openrouter) + self.assertIsNone(agent.chain_opencode) self.assertFalse(agent.is_available()) self.assertIn("Agent not connected", result[0]["params"]["reason"]) @@ -307,9 +250,7 @@ def test_non_auth_failure_does_not_invalidate_the_connection(self): # A transient/non-auth failure (e.g. timeout) must NOT disable a # connection that might still be perfectly valid. _, agent = _make_agent() - agent.chain_openrouter = object() - agent.preferred_provider = "openrouter" - agent.fallback_to_local = False + agent.chain_opencode = object() def fake_try_query_provider(provider, instruction, system_prompt): agent._last_query_error = RuntimeError("Connection timed out") @@ -319,7 +260,7 @@ def fake_try_query_provider(provider, instruction, system_prompt): agent.query("do something") - self.assertIsNotNone(agent.chain_openrouter) + self.assertIsNotNone(agent.chain_opencode) self.assertTrue(agent.is_available()) diff --git a/tests/trainer/services/test_agent_opencode_provider.py b/tests/trainer/services/test_agent_opencode_provider.py new file mode 100644 index 00000000..4f33aeb8 --- /dev/null +++ b/tests/trainer/services/test_agent_opencode_provider.py @@ -0,0 +1,448 @@ +"""Tests for DataManipulationAgent's OpenCode provider: config loading, +_setup_providers wiring self.chain_opencode, initialize_with_cloud_key/ +change_model/get_available_models/reset_connection, and clear_history/ +compact_history. OpenCode is the only supported agent backend. + +Follows the exact same "_install_agent_dependency_stubs + _make_agent" pattern +already used in test_agent_model_and_safety_unit.py / test_agent_prompt_unit.py +(duplicated per-file by this repo's own convention, not imported across test +files). OpenCodeChat itself is mocked out here -- its own HTTP/SSE behavior is +covered by test_opencode_chat.py against a real fake server. +""" + +import importlib +import json +import sys +import types +import unittest +from types import SimpleNamespace +from unittest import mock +from unittest.mock import MagicMock + +import pandas as pd + +# `_make_agent()` (like test_agent_model_and_safety_unit.py's identical helper) +# wraps each import in `mock.patch.dict(sys.modules, stubs, clear=False)`, which +# restores sys.modules to its EXACT pre-`with` snapshot on exit -- including +# evicting every module (torch, numpy, and their transitive dependency tree) +# that wasn't already resident when the block started. Those C extensions +# cannot be safely re-initialized after eviction, and the failure is +# order-dependent: it only shows up on the SECOND-and-later `_make_agent()` +# call in a run where nothing had already pulled in agent.py's full transitive +# import tree first. The sibling test file (test_agent_model_and_safety_unit.py) +# avoids this via its own top-level `from weightslab.trainer.services.data_service +# import ...`, which happens to import that whole tree before any stubbing runs. +# Importing the same module here for the same reason, not because this file +# needs DataService itself. +from weightslab.trainer.services.data_service import DataService # noqa: F401 + + +def _install_agent_dependency_stubs(): + stubs = { + "langchain_core": types.ModuleType("langchain_core"), + "langchain_core.prompts": types.ModuleType("langchain_core.prompts"), + } + stubs["langchain_core.prompts"].ChatPromptTemplate = object + return stubs + + +def _make_agent(df=None): + with mock.patch.dict(sys.modules, _install_agent_dependency_stubs(), clear=False): + agent_mod = importlib.import_module("weightslab.trainer.services.agent.agent") + + if df is None: + df = pd.DataFrame({"loss": [0.1, 0.9], "discarded": [False, False]}) + + ctx = SimpleNamespace(_all_datasets_df=df, _ctx=None) + agent = agent_mod.DataManipulationAgent(ctx) + + return agent_mod, agent + + +class TestOpenCodeConfigLoading(unittest.TestCase): + def test_opencode_url_and_model_default(self): + with mock.patch.dict("os.environ", {}, clear=False): + for key in ("OPENCODE_URL", "OPENCODE_MODEL"): + import os + os.environ.pop(key, None) + _, agent = _make_agent() + self.assertEqual(agent.opencode_url, "http://127.0.0.1:4096") + self.assertEqual(agent.opencode_model, "") + + def test_opencode_url_and_model_read_from_env(self): + with mock.patch.dict("os.environ", { + "OPENCODE_URL": "http://127.0.0.1:5555", + "OPENCODE_MODEL": "openrouter/anthropic/claude-opus-4.6", + }, clear=False): + _, agent = _make_agent() + self.assertEqual(agent.opencode_url, "http://127.0.0.1:5555") + self.assertEqual(agent.opencode_model, "openrouter/anthropic/claude-opus-4.6") + + def test_default_url_is_not_marked_explicit(self): + # Nobody chose this address -- OpenCodeChat._ensure_reachable must be + # free to replace it via auto-discovery/spawn if it's ever dead. + import os + with mock.patch.dict("os.environ", {}, clear=False): + os.environ.pop("OPENCODE_URL", None) + _, agent = _make_agent() + self.assertFalse(agent._opencode_url_explicit) + + def test_env_provided_url_is_marked_explicit(self): + # The opposite: an operator who deliberately set OPENCODE_URL is + # opting OUT of auto-discovery, not asking for it. + with mock.patch.dict("os.environ", {"OPENCODE_URL": "http://127.0.0.1:5555"}, clear=False): + _, agent = _make_agent() + self.assertTrue(agent._opencode_url_explicit) + + def test_workspace_dir_follows_weightslab_root_log_dir(self): + # The same directory `weightslab start ` roots the browser + # landing-page agent at -- the shared key opencode_process.py's lock + # file is discovered/published under. + with mock.patch.dict("os.environ", {"WEIGHTSLAB_ROOT_LOG_DIR": "/tmp/some-experiment"}, clear=False): + _, agent = _make_agent() + self.assertEqual(agent.opencode_workspace_dir, "/tmp/some-experiment") + + def test_workspace_dir_falls_back_to_cwd_when_unset(self): + import os + with mock.patch.dict("os.environ", {}, clear=False): + os.environ.pop("WEIGHTSLAB_ROOT_LOG_DIR", None) + _, agent = _make_agent() + self.assertEqual(agent.opencode_workspace_dir, os.getcwd()) + + +class TestSetupProvidersOpenCode(unittest.TestCase): + def test_opencode_chain_is_built(self): + agent_mod, agent = _make_agent() + fake_runnable = MagicMock() + with mock.patch.object(agent_mod, "OpenCodeChat") as mock_cls: + mock_cls.return_value.as_runnable.return_value = fake_runnable + initialized = agent._setup_providers() + + mock_cls.assert_called_once_with( + agent.opencode_url, agent.opencode_model, + workspace_dir=agent.opencode_workspace_dir, + url_is_explicit=agent._opencode_url_explicit, + model_is_explicit=agent._opencode_model_explicit, + ) + self.assertTrue(initialized) + self.assertIs(agent.chain_opencode, fake_runnable) + + def test_no_api_key_gate(self): + """OpenCode has no API-key concept at all -- the credential lives in + OpenCode's own config.""" + agent_mod, agent = _make_agent() + with mock.patch.object(agent_mod, "OpenCodeChat") as mock_cls: + mock_cls.return_value.as_runnable.return_value = MagicMock() + initialized = agent._setup_providers() + self.assertTrue(initialized) # succeeded with no key configured anywhere + + def test_setup_error_is_caught_and_reported_as_not_initialized(self): + agent_mod, agent = _make_agent() + with mock.patch.object(agent_mod, "OpenCodeChat", side_effect=RuntimeError("boom")): + initialized = agent._setup_providers() + self.assertFalse(initialized) + self.assertIsNone(agent.chain_opencode) + + +class TestInitializeWithCloudKeyOpenCode(unittest.TestCase): + def test_accepts_opencode_and_ignores_empty_api_key(self): + agent_mod, agent = _make_agent() + with mock.patch.object(agent_mod, "OpenCodeChat") as mock_cls: + mock_cls.return_value.as_runnable.return_value = MagicMock() + success, message = agent.initialize_with_cloud_key("", "opencode", "openrouter/openai/gpt-5") + + self.assertTrue(success) + self.assertIn("OpenCode", message) + self.assertEqual(agent.preferred_provider, "opencode") + self.assertEqual(agent.opencode_model, "openrouter/openai/gpt-5") + + def test_reports_failure_when_opencode_unreachable(self): + agent_mod, agent = _make_agent() + with mock.patch.object(agent_mod, "OpenCodeChat", side_effect=RuntimeError("connection refused")): + success, message = agent.initialize_with_cloud_key("", "opencode", None) + self.assertFalse(success) + self.assertIn("OpenCode", message) + + def test_rejects_any_provider_other_than_opencode(self): + _, agent = _make_agent() + success, message = agent.initialize_with_cloud_key("key", "anthropic-direct", None) + self.assertFalse(success) + self.assertIn("Only OpenCode", message) + + def test_rejects_openrouter_now_that_it_is_removed(self): + _, agent = _make_agent() + success, message = agent.initialize_with_cloud_key("sk-or-test", "openrouter", "openai/gpt-5") + self.assertFalse(success) + self.assertIn("Only OpenCode", message) + + +class TestChangeModelOpenCode(unittest.TestCase): + def test_switches_opencode_model(self): + agent_mod, agent = _make_agent() + with mock.patch.object(agent_mod, "OpenCodeChat") as mock_cls: + mock_cls.return_value.as_runnable.return_value = MagicMock() + success, message = agent.change_model("openrouter/openai/gpt-5-mini") + self.assertTrue(success) + self.assertEqual(agent.opencode_model, "openrouter/openai/gpt-5-mini") + + def test_reports_failure_when_opencode_unreachable(self): + agent_mod, agent = _make_agent() + with mock.patch.object(agent_mod, "OpenCodeChat", side_effect=RuntimeError("down")): + success, message = agent.change_model("openrouter/openai/gpt-5-mini") + self.assertFalse(success) + + def test_empty_model_is_rejected(self): + _, agent = _make_agent() + success, message = agent.change_model(" ") + self.assertFalse(success) + self.assertIn("Model cannot be empty", message) + + +class TestGetAvailableModelsOpenCode(unittest.TestCase): + def test_flattens_providers_into_provider_slash_model_strings(self): + _, agent = _make_agent() + # Stub out the self-heal itself (covered on its own in + # TestOpencodeBaseUrlHelper below) so this test only exercises the + # response-flattening logic against a fixed URL. + agent._opencode_chat._ensure_reachable = MagicMock() + agent._opencode_chat.base_url = "http://127.0.0.1:4096" + + fake_payload = { + "providers": [ + {"id": "openrouter", "models": {"anthropic/claude-opus-4.6": {}, "openai/gpt-5": {}}}, + {"id": "ollama", "models": {"llama3.2:3b": {}}}, + ], + } + + class _FakeResp: + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + def read(self): + import json + return json.dumps(fake_payload).encode() + + with mock.patch("urllib.request.urlopen", return_value=_FakeResp()): + ok, models, message = agent.get_available_models() + + self.assertTrue(ok) + self.assertEqual( + models, + sorted([ + "openrouter/anthropic/claude-opus-4.6", + "openrouter/openai/gpt-5", + "ollama/llama3.2:3b", + ]), + ) + + def test_reports_failure_when_opencode_server_unreachable(self): + _, agent = _make_agent() + agent._opencode_chat._ensure_reachable = MagicMock() + with mock.patch("urllib.request.urlopen", side_effect=OSError("refused")): + ok, models, message = agent.get_available_models() + self.assertFalse(ok) + self.assertEqual(models, []) + self.assertIn("OpenCode", message) + + def test_self_heals_the_base_url_before_querying(self): + """The bug this covers: `agent models` used to hit the raw + OPENCODE_URL directly and fail outright ("connection refused") if + nothing had spawned/discovered an OpenCode server yet -- unlike a + chat turn, which self-heals via OpenCodeChat._ensure_reachable. Pins + down that get_available_models now goes through the same self-heal + (via _opencode_base_url) before querying /config/providers.""" + _, agent = _make_agent() + with mock.patch.object(agent._opencode_chat, "_ensure_reachable") as ensure_mock: + with mock.patch("urllib.request.urlopen", side_effect=OSError("refused")): + agent.get_available_models() + ensure_mock.assert_called_once() + + +class TestOpencodeBaseUrlHelper(unittest.TestCase): + def test_falls_back_to_opencode_url_when_chat_not_yet_built(self): + _, agent = _make_agent() + agent._opencode_chat = None + agent.opencode_url = "http://127.0.0.1:7777" + self.assertEqual(agent._opencode_base_url(), "http://127.0.0.1:7777") + + def test_delegates_to_the_chats_self_healed_base_url(self): + _, agent = _make_agent() + agent._opencode_chat = MagicMock(base_url="http://127.0.0.1:9999") + self.assertEqual(agent._opencode_base_url(), "http://127.0.0.1:9999") + agent._opencode_chat._ensure_reachable.assert_called_once() + + +class TestGetContextUsage(unittest.TestCase): + """DataManipulationAgent.get_context_usage() -- backs the /context + command. Combines OpenCodeChat.last_usage (mocked directly here; its own + population from real SSE events is covered by test_opencode_chat.py) with + a context-window lookup via /config/providers (mocked urllib, same + _FakeResp pattern as TestGetAvailableModelsOpenCode above).""" + + class _FakeResp: + def __init__(self, payload): + self._payload = payload + + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + def read(self): + return json.dumps(self._payload).encode() + + def test_not_configured_when_opencode_chat_is_none(self): + _, agent = _make_agent() + agent._opencode_chat = None + + ok, usage, message = agent.get_context_usage() + + self.assertFalse(ok) + self.assertEqual(usage, {}) + self.assertIn("/init", message) + + def test_reports_no_turns_yet_when_last_usage_is_none(self): + _, agent = _make_agent() + agent.opencode_model = "" + agent._opencode_chat = MagicMock(last_usage=None) + + ok, usage, message = agent.get_context_usage() + + self.assertTrue(ok) + self.assertEqual(usage["context_window"], 0) + self.assertIn("No agent turns yet", message) + + def test_combines_last_usage_with_the_models_context_window(self): + _, agent = _make_agent() + agent.opencode_url = "http://127.0.0.1:4096" + agent.opencode_model = "openrouter/anthropic/claude-opus-4.6" + agent._opencode_chat = MagicMock(base_url="http://127.0.0.1:4096", last_usage={ + "input": 100, "output": 20, "reasoning": 5, "cache_read": 60, "cache_write": 10, + }) + + payload = { + "providers": [ + {"id": "openrouter", "models": { + "anthropic/claude-opus-4.6": {"limit": {"context": 200000, "output": 8192}}, + }}, + ], + } + with mock.patch("urllib.request.urlopen", return_value=self._FakeResp(payload)): + ok, usage, message = agent.get_context_usage() + + self.assertTrue(ok) + self.assertEqual(message, "") + self.assertEqual(usage, { + "model": "openrouter/anthropic/claude-opus-4.6", + "context_window": 200000, + "input_tokens": 100, + "output_tokens": 20, + "reasoning_tokens": 5, + "cache_read_tokens": 60, + "cache_write_tokens": 10, + }) + + def test_context_window_defaults_to_zero_when_the_server_is_unreachable(self): + """A failed /config/providers lookup must not sink the whole command -- + usage numbers are still worth showing without a window/percentage.""" + _, agent = _make_agent() + agent.opencode_model = "openrouter/anthropic/claude-opus-4.6" + agent._opencode_chat = MagicMock(base_url="http://127.0.0.1:4096", last_usage={ + "input": 10, "output": 2, "reasoning": 0, "cache_read": 0, "cache_write": 0, + }) + + with mock.patch("urllib.request.urlopen", side_effect=OSError("refused")): + ok, usage, message = agent.get_context_usage() + + self.assertTrue(ok) + self.assertEqual(usage["context_window"], 0) + self.assertEqual(usage["input_tokens"], 10) + + +class TestResetConnection(unittest.TestCase): + def test_reset_clears_opencode_chain_and_model(self): + agent_mod, agent = _make_agent() + agent.chain_opencode = MagicMock() + agent.opencode_model = "openrouter/openai/gpt-5" + + success, message = agent.reset_connection() + + self.assertTrue(success) + self.assertIsNone(agent.chain_opencode) + self.assertEqual(agent.preferred_provider, "opencode") + + +class _FakePipedRunnable: + """Stands in for `(ChatPromptTemplate | chain)` -- skips actual prompt + formatting (irrelevant to what compact_history does with the result) and + just forwards straight to the underlying chain's `.invoke`, matching real + LangChain's RunnableSequence semantics for this purpose.""" + + def __init__(self, chain): + self._chain = chain + + def invoke(self, variables): + return self._chain.invoke(variables) + + +class _FakeChatPromptTemplate: + @classmethod + def from_messages(cls, messages): + return cls() + + def __or__(self, chain): + return _FakePipedRunnable(chain) + + +class TestClearAndCompactHistory(unittest.TestCase): + def test_clear_history_empties_and_reports_count(self): + _, agent = _make_agent() + agent.history = ["User: a", "Action: 1 ops executed", "User: b", "Action: 2 ops executed"] + + success, message = agent.clear_history() + + self.assertTrue(success) + self.assertEqual(agent.history, []) + self.assertIn("4", message) + + def test_compact_history_on_empty_history_is_a_no_op_success(self): + # Short-circuits before touching ChatPromptTemplate at all -- no patch needed. + _, agent = _make_agent() + agent.history = [] + success, message = agent.compact_history() + self.assertTrue(success) + self.assertEqual(agent.history, []) + + def test_compact_history_replaces_history_with_one_summary(self): + agent_mod, agent = _make_agent() + agent.history = ["User: discard bad samples", "Action: 3 ops executed"] + + fake_reply = SimpleNamespace(content="Discarded 3 low-quality samples per user request.") + agent.chain_opencode = MagicMock(invoke=MagicMock(return_value=fake_reply)) + + with mock.patch.object(agent_mod, "ChatPromptTemplate", _FakeChatPromptTemplate): + success, message = agent.compact_history() + + self.assertTrue(success) + self.assertEqual(len(agent.history), 1) + self.assertIn("Discarded 3 low-quality samples", agent.history[0]) + + def test_compact_history_fails_cleanly_when_no_provider_available(self): + agent_mod, agent = _make_agent() + agent.history = ["User: x"] + agent.chain_opencode = None + + with mock.patch.object(agent_mod, "ChatPromptTemplate", _FakeChatPromptTemplate): + success, message = agent.compact_history() + + self.assertFalse(success) + # History is left untouched on failure -- nothing was actually compacted. + self.assertEqual(agent.history, ["User: x"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/trainer/services/test_agent_prompt_unit.py b/tests/trainer/services/test_agent_prompt_unit.py index 116d7cdb..093d992a 100644 --- a/tests/trainer/services/test_agent_prompt_unit.py +++ b/tests/trainer/services/test_agent_prompt_unit.py @@ -10,13 +10,9 @@ def _install_agent_dependency_stubs(): stubs = { - "langchain_ollama": types.ModuleType("langchain_ollama"), - "langchain_openai": types.ModuleType("langchain_openai"), "langchain_core": types.ModuleType("langchain_core"), "langchain_core.prompts": types.ModuleType("langchain_core.prompts"), } - stubs["langchain_ollama"].ChatOllama = object - stubs["langchain_openai"].ChatOpenAI = object stubs["langchain_core.prompts"].ChatPromptTemplate = object return stubs @@ -37,18 +33,6 @@ def _rewrite_origin_literals(self, code): return code -class _FakeChatModel: - def __init__(self, *args, **kwargs): - self.args = args - self.kwargs = kwargs - - def with_structured_output(self, schema): - return self - - def invoke(self, prompt): - return SimpleNamespace(content="OK") - - class TestAgentPromptUnit(unittest.TestCase): def test_intent_prompt_contains_expected_placeholders(self): self.assertIn("{row_count}", INTENT_PROMPT) @@ -134,44 +118,7 @@ def test_agent_models_and_handlers(self): ) self.assertEqual(action["function"], "action.save") - def test_initialize_with_cloud_key_checks_chat_connectivity(self): - with unittest.mock.patch.dict(sys.modules, _install_agent_dependency_stubs(), clear=False): - agent_mod = importlib.import_module("weightslab.trainer.services.agent.agent") - - ctx = SimpleNamespace( - _all_datasets_df=agent_mod.pd.DataFrame({"metric": [1.0, 2.0]}), - ) - - with mock.patch.object(agent_mod, "ChatOpenAI", _FakeChatModel), mock.patch.object(agent_mod, "ChatOllama", _FakeChatModel): - agent = agent_mod.DataManipulationAgent(ctx) - ok, message = agent.initialize_with_cloud_key("test-key", "openrouter", "google/gemini-2.5-flash") - - self.assertTrue(ok) - self.assertIn("initialized successfully", message) - self.assertIsNotNone(agent.chain_openrouter) - self.assertEqual(agent.openrouter_model, "google/gemini-2.5-flash") - - def test_initialize_with_cloud_key_fails_when_probe_fails(self): - with unittest.mock.patch.dict(sys.modules, _install_agent_dependency_stubs(), clear=False): - agent_mod = importlib.import_module("weightslab.trainer.services.agent.agent") - - class _FailingChatModel(_FakeChatModel): - def invoke(self, prompt): - raise RuntimeError("401 Unauthorized") - - ctx = SimpleNamespace( - _all_datasets_df=agent_mod.pd.DataFrame({"metric": [1.0, 2.0]}), - ) - - with mock.patch.object(agent_mod, "ChatOpenAI", _FailingChatModel), mock.patch.object(agent_mod, "ChatOllama", _FakeChatModel): - agent = agent_mod.DataManipulationAgent(ctx) - ok, message = agent.initialize_with_cloud_key("bad-key", "openrouter", "~google/gemini-flash-latest") - - self.assertFalse(ok) - self.assertIn("connectivity check failed", message) - self.assertIsNone(agent.chain_openrouter) - - def test_initialize_with_cloud_key_rejects_non_openrouter_provider(self): + def test_initialize_with_cloud_key_rejects_non_opencode_provider(self): with unittest.mock.patch.dict(sys.modules, _install_agent_dependency_stubs(), clear=False): agent_mod = importlib.import_module("weightslab.trainer.services.agent.agent") @@ -179,12 +126,11 @@ def test_initialize_with_cloud_key_rejects_non_openrouter_provider(self): _all_datasets_df=agent_mod.pd.DataFrame({"metric": [1.0, 2.0]}), ) - with mock.patch.object(agent_mod, "ChatOpenAI", _FakeChatModel), mock.patch.object(agent_mod, "ChatOllama", _FakeChatModel): - agent = agent_mod.DataManipulationAgent(ctx) - ok, message = agent.initialize_with_cloud_key("test-key", "grok", "grok-3-mini") + agent = agent_mod.DataManipulationAgent(ctx) + ok, message = agent.initialize_with_cloud_key("test-key", "grok", "grok-3-mini") self.assertFalse(ok) - self.assertIn("Only OpenRouter", message) + self.assertIn("Only OpenCode", message) def test_build_python_mask_keeps_string_literals(self): with unittest.mock.patch.dict(sys.modules, _install_agent_dependency_stubs(), clear=False): @@ -202,8 +148,7 @@ def test_build_python_mask_keeps_string_literals(self): ), ) - with mock.patch.object(agent_mod, "ChatOpenAI", _FakeChatModel), mock.patch.object(agent_mod, "ChatOllama", _FakeChatModel): - agent = agent_mod.DataManipulationAgent(ctx) + agent = agent_mod.DataManipulationAgent(ctx) mask = agent._build_python_mask( [agent_mod.Condition(column="origin", op="==", value="train")] @@ -234,8 +179,7 @@ def test_compact_schema_for_prompt_separates_index_levels_from_columns(self): ), ) - with mock.patch.object(agent_mod, "ChatOpenAI", _FakeChatModel), mock.patch.object(agent_mod, "ChatOllama", _FakeChatModel): - agent = agent_mod.DataManipulationAgent(ctx) + agent = agent_mod.DataManipulationAgent(ctx) schema_text = agent._compact_schema_for_prompt() @@ -321,8 +265,7 @@ def test_filter_by_origin_and_loss(self): ), ) - with mock.patch.object(agent_mod, "ChatOpenAI", _FakeChatModel), mock.patch.object(agent_mod, "ChatOllama", _FakeChatModel): - agent = agent_mod.DataManipulationAgent(ctx) + agent = agent_mod.DataManipulationAgent(ctx) cond1 = agent_mod.Condition(column="origin", op="==", value="train") cond2 = agent_mod.Condition(column="loss", op="<", value=0.3) @@ -368,8 +311,7 @@ def test_tag_high_loss_samples(self): ), ) - with mock.patch.object(agent_mod, "ChatOpenAI", _FakeChatModel), mock.patch.object(agent_mod, "ChatOllama", _FakeChatModel): - agent = agent_mod.DataManipulationAgent(ctx) + agent = agent_mod.DataManipulationAgent(ctx) step = agent_mod.AtomicIntent( kind="transform", @@ -401,8 +343,7 @@ def test_tag_from_quantile_computation(self): ), ) - with mock.patch.object(agent_mod, "ChatOpenAI", _FakeChatModel), mock.patch.object(agent_mod, "ChatOllama", _FakeChatModel): - agent = agent_mod.DataManipulationAgent(ctx) + agent = agent_mod.DataManipulationAgent(ctx) step = agent_mod.AtomicIntent( kind="transform", @@ -496,8 +437,7 @@ def test_analysis_train_loss_stddev(self): ), ) - with mock.patch.object(agent_mod, "ChatOpenAI", _FakeChatModel), mock.patch.object(agent_mod, "ChatOllama", _FakeChatModel): - agent = agent_mod.DataManipulationAgent(ctx) + agent = agent_mod.DataManipulationAgent(ctx) step = agent_mod.AtomicIntent( kind="analysis", @@ -531,8 +471,7 @@ def test_tag_outliers_by_stddev(self): ), ) - with mock.patch.object(agent_mod, "ChatOpenAI", _FakeChatModel), mock.patch.object(agent_mod, "ChatOllama", _FakeChatModel): - agent = agent_mod.DataManipulationAgent(ctx) + agent = agent_mod.DataManipulationAgent(ctx) step = agent_mod.AtomicIntent( kind="transform", @@ -563,8 +502,7 @@ def test_tag_outliers_by_iqr(self): ), ) - with mock.patch.object(agent_mod, "ChatOpenAI", _FakeChatModel), mock.patch.object(agent_mod, "ChatOllama", _FakeChatModel): - agent = agent_mod.DataManipulationAgent(ctx) + agent = agent_mod.DataManipulationAgent(ctx) step = agent_mod.AtomicIntent( kind="transform", @@ -597,8 +535,7 @@ def test_filter_and_tag_combination(self): ), ) - with mock.patch.object(agent_mod, "ChatOpenAI", _FakeChatModel), mock.patch.object(agent_mod, "ChatOllama", _FakeChatModel): - agent_mod.DataManipulationAgent(ctx) + agent_mod.DataManipulationAgent(ctx) # First filter: keep only train filt_cond = agent_mod.Condition(column="origin", op="==", value="train") @@ -635,8 +572,7 @@ def test_untag_operation(self): ), ) - with mock.patch.object(agent_mod, "ChatOpenAI", _FakeChatModel), mock.patch.object(agent_mod, "ChatOllama", _FakeChatModel): - agent = agent_mod.DataManipulationAgent(ctx) + agent = agent_mod.DataManipulationAgent(ctx) step = agent_mod.AtomicIntent( kind="transform", @@ -667,8 +603,7 @@ def test_rename_tag_operation(self): ), ) - with mock.patch.object(agent_mod, "ChatOpenAI", _FakeChatModel), mock.patch.object(agent_mod, "ChatOllama", _FakeChatModel): - agent = agent_mod.DataManipulationAgent(ctx) + agent = agent_mod.DataManipulationAgent(ctx) step = agent_mod.AtomicIntent( kind="transform", diff --git a/tests/trainer/services/test_agent_service_unit.py b/tests/trainer/services/test_agent_service_unit.py index fb26bd6c..c3d394ed 100644 --- a/tests/trainer/services/test_agent_service_unit.py +++ b/tests/trainer/services/test_agent_service_unit.py @@ -24,9 +24,11 @@ def test_check_agent_health_reports_ready_when_available(self): self.assertTrue(response.available) self.assertIn('Ready to help you.', response.message) - def test_initialize_agent_delegates_to_agent_with_openrouter(self): + def test_initialize_agent_rejects_openrouter_now_that_it_is_removed(self): + """PROVIDER_OPENROUTER (0) is kept in the .proto only for wire + compatibility with older frontends -- requesting it must be rejected + cleanly rather than reaching the (now opencode-only) agent.""" agent = MagicMock() - agent.initialize_with_cloud_key.return_value = (True, 'ok') service, _ = self._make_service(agent=agent) response = service.InitializeAgent( @@ -38,13 +40,9 @@ def test_initialize_agent_delegates_to_agent_with_openrouter(self): None, ) - agent.initialize_with_cloud_key.assert_called_once_with( - 'sk-or-test', - 'openrouter', - '~google/gemini-flash-latest', - ) - self.assertTrue(response.success) - self.assertEqual(response.message, 'ok') + agent.initialize_with_cloud_key.assert_not_called() + self.assertFalse(response.success) + self.assertIn('Only OpenCode', response.message) def test_initialize_agent_rejects_unsupported_provider(self): agent = MagicMock() @@ -61,7 +59,7 @@ def test_initialize_agent_rejects_unsupported_provider(self): agent.initialize_with_cloud_key.assert_not_called() self.assertFalse(response.success) - self.assertIn('Only OpenRouter', response.message) + self.assertIn('Only OpenCode', response.message) def test_change_get_and_reset_agent_delegate_to_agent(self): agent = MagicMock() @@ -99,6 +97,89 @@ def test_methods_fail_cleanly_when_agent_backend_missing(self): self.assertEqual(list(list_response.models), []) self.assertIn('not running', init_response.message) + def test_initialize_agent_accepts_opencode_provider(self): + agent = MagicMock() + agent.initialize_with_cloud_key.return_value = (True, 'Agent initialized successfully via OpenCode. Ready to help you.') + service, _ = self._make_service(agent=agent) + + response = service.InitializeAgent( + pb2.InitializeAgentRequest( + api_key='', # ignored for opencode -- credential lives in OpenCode's own config + provider=pb2.PROVIDER_OPENCODE, + model='openrouter/anthropic/claude-opus-4.6', + ), + None, + ) + + agent.initialize_with_cloud_key.assert_called_once_with( + '', 'opencode', 'openrouter/anthropic/claude-opus-4.6', + ) + self.assertTrue(response.success) + + def test_clear_agent_history_delegates_to_agent(self): + agent = MagicMock() + agent.clear_history.return_value = (True, 'Cleared 4 history entries.') + service, _ = self._make_service(agent=agent) + + response = service.ClearAgentHistory(pb2.Empty(), None) + + agent.clear_history.assert_called_once_with() + self.assertTrue(response.success) + self.assertEqual(response.message, 'Cleared 4 history entries.') + + def test_compact_agent_history_delegates_to_agent(self): + agent = MagicMock() + agent.compact_history.return_value = (True, 'Compacted 4 entries into one summary.') + service, _ = self._make_service(agent=agent) + + response = service.CompactAgentHistory(pb2.Empty(), None) + + agent.compact_history.assert_called_once_with() + self.assertTrue(response.success) + + def test_clear_and_compact_history_fail_cleanly_when_agent_backend_missing(self): + service, _ = self._make_service(agent=None) + + clear_response = service.ClearAgentHistory(pb2.Empty(), None) + compact_response = service.CompactAgentHistory(pb2.Empty(), None) + + self.assertFalse(clear_response.success) + self.assertFalse(compact_response.success) + self.assertIn('not running', clear_response.message) + + def test_get_agent_context_usage_delegates_to_agent(self): + agent = MagicMock() + agent.get_context_usage.return_value = (True, { + 'model': 'openrouter/anthropic/claude-opus-4.6', + 'context_window': 200000, + 'input_tokens': 100, + 'output_tokens': 20, + 'reasoning_tokens': 5, + 'cache_read_tokens': 60, + 'cache_write_tokens': 10, + }, '') + service, _ = self._make_service(agent=agent) + + response = service.GetAgentContextUsage(pb2.Empty(), None) + + agent.get_context_usage.assert_called_once_with() + self.assertTrue(response.success) + self.assertEqual(response.model, 'openrouter/anthropic/claude-opus-4.6') + self.assertEqual(response.context_window, 200000) + self.assertEqual(response.input_tokens, 100) + self.assertEqual(response.output_tokens, 20) + self.assertEqual(response.reasoning_tokens, 5) + self.assertEqual(response.cache_read_tokens, 60) + self.assertEqual(response.cache_write_tokens, 10) + + def test_get_agent_context_usage_fails_cleanly_when_agent_backend_missing(self): + service, _ = self._make_service(agent=None) + + response = service.GetAgentContextUsage(pb2.Empty(), None) + + self.assertFalse(response.success) + self.assertIn('not running', response.message) + if __name__ == '__main__': unittest.main() diff --git a/tests/trainer/services/test_opencode_chat.py b/tests/trainer/services/test_opencode_chat.py new file mode 100644 index 00000000..847b15ef --- /dev/null +++ b/tests/trainer/services/test_opencode_chat.py @@ -0,0 +1,463 @@ +"""Tests for OpenCodeChat (weightslab/trainer/services/agent/opencode_chat.py). + +Exercises the module against a REAL minimal HTTP server implementing the +subset of OpenCode's protocol this class needs (POST /session, POST +/session/{id}/message, GET /event as text/event-stream) rather than mocking +urllib -- the class's correctness hinges on stream-first ordering and SSE +event parsing, which a mocked urlopen would not honestly exercise. Mirrors +the fake-server technique already used in tests/ui/test_server_agent.py for +the same reason. +""" + +import json +import threading +import time +import unittest +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from unittest import mock + +from weightslab.trainer.services.agent.opencode_chat import OpenCodeChat, OpenCodeError + + +class _FakeOpenCodeHandler(BaseHTTPRequestHandler): + """Implements just enough of OpenCode's HTTP+SSE surface to drive + OpenCodeChat through a full create -> send -> stream -> idle cycle. + Configured per-server-instance via class attributes the test sets before + starting it (see _FakeOpenCodeServer below).""" + + def log_message(self, *a): + pass + + def _send_json(self, obj): + body = json.dumps(obj).encode("utf-8") + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def do_POST(self): + length = int(self.headers.get("Content-Length", "0") or 0) + raw = self.rfile.read(length) if length else b"" + try: + body = json.loads(raw.decode("utf-8")) if raw else {} + except ValueError: + body = {} + + if self.path == "/session": + self.server.recorded_session_titles.append(body.get("title")) + self._send_json({"id": self.server.session_id}) + return + + if self.path == f"/session/{self.server.session_id}/message": + self.server.recorded_messages.append(body) + # A real server holds this open for the whole turn; this fake + # returns immediately -- OpenCodeChat must not rely on this + # response for content, only on the SSE stream (see its own + # docstring). Sleep briefly so the message genuinely arrives + # after the stream has had a chance to open, exercising the + # stream-first ordering rather than accidentally passing by luck. + time.sleep(0.05) + self._send_json({}) + return + + self.send_response(404) + self.end_headers() + + def do_GET(self): + if self.path != "/event": + self.send_response(404) + self.end_headers() + return + + self.send_response(200) + self.send_header("Content-Type", "text/event-stream") + self.end_headers() + + def emit(event: dict) -> None: + payload = f"data: {json.dumps(event)}\n\n".encode("utf-8") + self.wfile.write(payload) + self.wfile.flush() + + session_id = self.server.session_id + # Wait for the message POST to actually land before emitting anything + # -- proves OpenCodeChat is genuinely reading a live stream, not just + # replaying a canned response. + deadline = time.monotonic() + 5 + while not self.server.recorded_messages and time.monotonic() < deadline: + time.sleep(0.01) + + info = {"id": "msg_1", "role": "assistant", "sessionID": session_id} + if self.server.reply_tokens is not None: + info["tokens"] = self.server.reply_tokens + emit({"type": "message.updated", "properties": {"info": info}}) + for delta in self.server.reply_deltas: + emit({"type": "message.part.updated", "properties": {"part": { + "id": f"prt_{delta[:4]}", "type": "text", "text": delta, "messageID": "msg_1", "sessionID": session_id, + }}}) + if self.server.emit_error: + emit({"type": "session.error", "properties": {"sessionID": session_id}}) + else: + emit({"type": "session.idle", "properties": {"sessionID": session_id}}) + # Keep the connection open a moment so OpenCodeChat's break-on-idle has + # definitely already fired before we tear down. + time.sleep(0.05) + + +class _FakeOpenCodeServer(ThreadingHTTPServer): + daemon_threads = True + + def __init__(self, *args, reply_deltas, emit_error=False, reply_tokens=None, **kwargs): + super().__init__(*args, **kwargs) + self.session_id = "ses_test1" + self.reply_deltas = reply_deltas + self.emit_error = emit_error + self.reply_tokens = reply_tokens + self.recorded_messages = [] + self.recorded_session_titles = [] + + +class _ServerTestCase(unittest.TestCase): + def _start_server(self, reply_deltas, emit_error=False, reply_tokens=None): + self.httpd = _FakeOpenCodeServer( + ("127.0.0.1", 0), _FakeOpenCodeHandler, + reply_deltas=reply_deltas, emit_error=emit_error, reply_tokens=reply_tokens, + ) + self.port = self.httpd.server_address[1] + self.thread = threading.Thread(target=self.httpd.serve_forever, daemon=True) + self.thread.start() + + def tearDown(self): + if hasattr(self, "httpd"): + self.httpd.shutdown() + self.thread.join(timeout=5) + + +class TestOpenCodeChatCall(_ServerTestCase): + def test_collects_streamed_text_and_returns_an_ai_message(self): + self._start_server(reply_deltas=["Here is ", "the answer."]) + chat = OpenCodeChat(f"http://127.0.0.1:{self.port}", model="openrouter/openai/gpt-5", timeout=10) + + result = chat.as_runnable().invoke("do the thing") + + self.assertEqual(result.content, "Here is the answer.") + + def test_sends_the_model_ref_split_on_first_slash_only(self): + self._start_server(reply_deltas=["ok"]) + chat = OpenCodeChat(f"http://127.0.0.1:{self.port}", model="openrouter/anthropic/claude-opus-4.6", timeout=10) + + chat._call("hello") + + self.assertEqual(len(self.httpd.recorded_messages), 1) + self.assertEqual( + self.httpd.recorded_messages[0]["model"], + {"providerID": "openrouter", "modelID": "anthropic/claude-opus-4.6"}, + ) + + def test_disables_every_mutating_tool_on_the_outgoing_message(self): + """This wrapper backs the SDK agent's text/JSON call sites, which parse + the reply themselves -- unlike the Weights Studio landing chat, it must + never let OpenCode write files as a side effect.""" + self._start_server(reply_deltas=["ok"]) + chat = OpenCodeChat(f"http://127.0.0.1:{self.port}", model=None, timeout=10) + + chat._call("hello") + + tools = self.httpd.recorded_messages[0]["tools"] + self.assertEqual(tools, {"write": False, "edit": False, "patch": False, "bash": False}) + + def test_omits_model_field_when_none_configured(self): + self._start_server(reply_deltas=["ok"]) + chat = OpenCodeChat(f"http://127.0.0.1:{self.port}", model=None, timeout=10) + + chat._call("hello") + + self.assertNotIn("model", self.httpd.recorded_messages[0]) + + def test_creates_a_fresh_session_per_call(self): + self._start_server(reply_deltas=["a"]) + chat = OpenCodeChat(f"http://127.0.0.1:{self.port}", model=None, timeout=10) + + chat._call("first") + chat._call("second") + + # A fresh session per call (self.history on DataManipulationAgent + # already carries cross-call context) -- both calls hit /session, not + # a reused id. + self.assertEqual(len(self.httpd.recorded_session_titles), 2) + + def test_degrades_to_empty_string_on_session_error_rather_than_raising(self): + self._start_server(reply_deltas=["partial"], emit_error=True) + chat = OpenCodeChat(f"http://127.0.0.1:{self.port}", model=None, timeout=10) + + # session.error still ends the read loop cleanly; whatever text arrived + # before the error is still returned rather than raising. + result = chat._call("hello") + self.assertEqual(result.content, "partial") + + def test_raises_opencode_error_when_the_server_is_unreachable(self): + chat = OpenCodeChat("http://127.0.0.1:1", model=None, timeout=2) # nothing listens on port 1 + with self.assertRaises(OpenCodeError): + chat._call("hello") + + def test_populates_last_usage_from_the_assistant_messages_tokens(self): + self._start_server( + reply_deltas=["ok"], + reply_tokens={"input": 120, "output": 30, "reasoning": 5, "cache": {"read": 80, "write": 10}}, + ) + chat = OpenCodeChat(f"http://127.0.0.1:{self.port}", model=None, timeout=10) + + self.assertIsNone(chat.last_usage) + chat._call("hello") + + self.assertEqual( + chat.last_usage, + {"input": 120, "output": 30, "reasoning": 5, "cache_read": 80, "cache_write": 10}, + ) + + def test_last_usage_is_none_when_the_reply_carries_no_tokens_field(self): + self._start_server(reply_deltas=["ok"]) # reply_tokens defaults to None + chat = OpenCodeChat(f"http://127.0.0.1:{self.port}", model=None, timeout=10) + + chat._call("hello") + + self.assertIsNone(chat.last_usage) + + +class TestModelRefParsing(unittest.TestCase): + def test_splits_on_first_slash_only(self): + chat = OpenCodeChat("http://x", model="openrouter/anthropic/claude-opus-4.6") + self.assertEqual(chat._model_ref(), {"providerID": "openrouter", "modelID": "anthropic/claude-opus-4.6"}) + + def test_none_when_no_model_configured(self): + chat = OpenCodeChat("http://x", model=None) + self.assertIsNone(chat._model_ref()) + + def test_none_when_model_has_no_slash(self): + chat = OpenCodeChat("http://x", model="justamodel") + self.assertIsNone(chat._model_ref()) + + +class TestHandleEvent(unittest.TestCase): + """Unit-level checks on the event state machine, independent of the + network -- complements the end-to-end server tests above.""" + + def test_ignores_parts_from_a_message_not_yet_known_to_be_assistant(self): + text_parts = {} + assistant_ids = set() + payload = json.dumps({ + "type": "message.part.updated", + "properties": {"part": {"id": "p1", "type": "text", "text": "x", "messageID": "msg_unknown", "sessionID": "s1"}}, + }) + outcome = OpenCodeChat._handle_event(payload, "s1", assistant_ids, text_parts) + self.assertIsNone(outcome) + self.assertEqual(text_parts, {}) + + def test_ignores_events_for_a_different_session(self): + text_parts = {} + assistant_ids = {"msg_1"} + payload = json.dumps({ + "type": "message.part.updated", + "properties": {"part": {"id": "p1", "type": "text", "text": "x", "messageID": "msg_1", "sessionID": "OTHER"}}, + }) + outcome = OpenCodeChat._handle_event(payload, "s1", assistant_ids, text_parts) + self.assertIsNone(outcome) + self.assertEqual(text_parts, {}) + + def test_malformed_payload_is_ignored_not_raised(self): + outcome = OpenCodeChat._handle_event("not json", "s1", set(), {}) + self.assertIsNone(outcome) + + def test_session_idle_signals_completion(self): + outcome = OpenCodeChat._handle_event( + json.dumps({"type": "session.idle", "properties": {"sessionID": "s1"}}), "s1", set(), {}, + ) + self.assertEqual(outcome, "idle") + + def test_populates_the_usage_dict_when_the_assistant_message_carries_tokens(self): + usage = {} + payload = json.dumps({ + "type": "message.updated", + "properties": {"info": { + "id": "msg_1", "role": "assistant", "sessionID": "s1", + "tokens": {"input": 10, "output": 2, "reasoning": 0, "cache": {"read": 5, "write": 1}}, + }}, + }) + outcome = OpenCodeChat._handle_event(payload, "s1", set(), {}, usage) + self.assertIsNone(outcome) + self.assertEqual(usage, {"input": 10, "output": 2, "reasoning": 0, "cache_read": 5, "cache_write": 1}) + + def test_usage_param_is_optional_and_ignored_when_omitted(self): + # Existing call sites that predate the usage tracking must keep working. + payload = json.dumps({ + "type": "message.updated", + "properties": {"info": {"id": "msg_1", "role": "assistant", "sessionID": "s1", "tokens": {"input": 1}}}, + }) + outcome = OpenCodeChat._handle_event(payload, "s1", set(), {}) + self.assertIsNone(outcome) + + +class TestEnsureReachable(unittest.TestCase): + """The other half of the cross-process handoff opencode_process.py + implements: this side self-heals its own base_url via that module's + resolve_or_spawn_opencode instead of staying pointed at a dead address + forever. The module itself (env/lockfile/spawn precedence, real + subprocess spawn+poll) is covered by tests/test_opencode_process.py -- + these mock it out to pin down exactly when OpenCodeChat calls it.""" + + def test_explicit_url_is_never_auto_replaced_even_if_dead(self): + chat = OpenCodeChat("http://127.0.0.1:1", workspace_dir="/tmp/x", url_is_explicit=True) + with mock.patch("weightslab.opencode_process.opencode_healthy") as healthy, \ + mock.patch("weightslab.opencode_process.resolve_or_spawn_opencode") as resolve: + chat._ensure_reachable() + healthy.assert_not_called() + resolve.assert_not_called() + self.assertEqual(chat.base_url, "http://127.0.0.1:1") + + def test_already_healthy_non_explicit_url_is_left_alone(self): + chat = OpenCodeChat("http://127.0.0.1:4096", workspace_dir="/tmp/x", url_is_explicit=False) + with mock.patch("weightslab.opencode_process.opencode_healthy", return_value=True) as healthy, \ + mock.patch("weightslab.opencode_process.resolve_or_spawn_opencode") as resolve: + chat._ensure_reachable() + healthy.assert_called_once_with("http://127.0.0.1:4096") + resolve.assert_not_called() + self.assertEqual(chat.base_url, "http://127.0.0.1:4096") + + def test_dead_non_explicit_url_resolves_or_spawns_a_replacement(self): + chat = OpenCodeChat("http://127.0.0.1:4096", workspace_dir="/tmp/x", url_is_explicit=False) + with mock.patch("weightslab.opencode_process.opencode_healthy", return_value=False), \ + mock.patch("weightslab.opencode_process.resolve_or_spawn_opencode", + return_value={"ok": True, "url": "http://127.0.0.1:9999"}) as resolve: + chat._ensure_reachable() + resolve.assert_called_once_with("/tmp/x") + self.assertEqual(chat.base_url, "http://127.0.0.1:9999") + + def test_failed_resolve_leaves_the_dead_url_in_place(self): + # A visible connection error on the next real call is more honest + # than silently pretending nothing changed. + chat = OpenCodeChat("http://127.0.0.1:4096", workspace_dir="/tmp/x", url_is_explicit=False) + with mock.patch("weightslab.opencode_process.opencode_healthy", return_value=False), \ + mock.patch("weightslab.opencode_process.resolve_or_spawn_opencode", + return_value={"ok": False, "error": "no opencode"}): + chat._ensure_reachable() + self.assertEqual(chat.base_url, "http://127.0.0.1:4096") + + def test_missing_workspace_dir_falls_back_to_the_current_directory(self): + chat = OpenCodeChat("http://127.0.0.1:4096", url_is_explicit=False) # workspace_dir defaults to None + with mock.patch("weightslab.opencode_process.opencode_healthy", return_value=False), \ + mock.patch("weightslab.opencode_process.resolve_or_spawn_opencode", + return_value={"ok": True, "url": "http://127.0.0.1:9999"}) as resolve: + chat._ensure_reachable() + resolve.assert_called_once_with(".") + + def test_call_invokes_ensure_reachable_before_creating_a_session(self): + # _call is the single real entry point all three of + # DataManipulationAgent's call sites go through (see module + # docstring) -- this pins down that self-healing actually happens + # on the path real turns take, not just when called directly. + chat = OpenCodeChat("http://127.0.0.1:4096", url_is_explicit=False) + calls = [] + chat._ensure_reachable = lambda: calls.append("ensure_reachable") + chat._ensure_model_resolved = lambda: calls.append("ensure_model_resolved") + chat._create_session = lambda: (calls.append("create_session") or "ses_x") + chat._collect_reply = lambda session_id, text: (calls.append("collect_reply") or "ok") + chat._call("hello") + self.assertEqual( + calls, + ["ensure_reachable", "ensure_model_resolved", "create_session", "collect_reply"], + ) + + +class TestEnsureModelResolved(unittest.TestCase): + """Confirmed live: leaving `model` unset does NOT mean "OpenCode picks a + sensible default" -- it means OpenCode picks whatever's configured, + arbitrarily (an image-generation preview model, in the case that + surfaced this). _ensure_model_resolved is the fix: resolve a REAL + default (the user's last actual pick, or a provider's own configured + default) instead of leaving it to chance. Mirrors TestEnsureReachable's + own mocked-request style -- the wire format itself (GET /config, + GET /config/providers) is exercised for real in + TestGetAvailableModelsOpenCode (test_agent_opencode_provider.py).""" + + def _fake_response(self, payload): + class _Resp: + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + def read(self): + return json.dumps(payload).encode() + return _Resp() + + def test_explicit_model_is_never_auto_replaced(self): + chat = OpenCodeChat("http://127.0.0.1:4096", model="openrouter/anthropic/claude-opus-4.6", model_is_explicit=True) + with mock.patch.object(chat, "_request") as request_mock: + chat._ensure_model_resolved() + request_mock.assert_not_called() + self.assertEqual(chat.model, "openrouter/anthropic/claude-opus-4.6") + + def test_already_set_non_explicit_model_is_left_alone(self): + # Already resolved once (e.g. a prior call) -- don't re-resolve or + # re-request every single turn. + chat = OpenCodeChat("http://127.0.0.1:4096", model="openrouter/openai/gpt-5", model_is_explicit=False) + with mock.patch.object(chat, "_request") as request_mock: + chat._ensure_model_resolved() + request_mock.assert_not_called() + self.assertEqual(chat.model, "openrouter/openai/gpt-5") + + def test_unset_model_resolves_from_config_own_model_field(self): + chat = OpenCodeChat("http://127.0.0.1:4096", model=None, model_is_explicit=False) + with mock.patch.object(chat, "_request", return_value=self._fake_response({"model": "anthropic/claude-haiku-4.5"})) as request_mock: + chat._ensure_model_resolved() + request_mock.assert_called_once_with("/config") + self.assertEqual(chat.model, "anthropic/claude-haiku-4.5") + + def test_falls_back_to_provider_default_when_config_has_no_model(self): + chat = OpenCodeChat("http://127.0.0.1:4096", model=None, model_is_explicit=False) + providers_payload = { + "providers": [{"id": "openrouter"}, {"id": "ollama"}], + "default": {"openrouter": "openai/gpt-5-mini"}, + } + responses = [self._fake_response({}), self._fake_response(providers_payload)] + with mock.patch.object(chat, "_request", side_effect=lambda *a, **k: responses.pop(0)) as request_mock: + chat._ensure_model_resolved() + self.assertEqual(request_mock.call_args_list[0].args, ("/config",)) + self.assertEqual(request_mock.call_args_list[1].args, ("/config/providers",)) + self.assertEqual(chat.model, "openrouter/openai/gpt-5-mini") + + def test_no_resolvable_model_falls_back_to_the_hardcoded_default(self): + chat = OpenCodeChat("http://127.0.0.1:4096", model=None, model_is_explicit=False) + with mock.patch.object(chat, "_request", side_effect=OSError("refused")): + chat._ensure_model_resolved() # must not raise + self.assertEqual(chat.model, "opencode/deepseek-v4-flash-free") + + def test_falls_back_to_the_hardcoded_default_when_providers_have_no_default_either(self): + # /config and /config/providers both answer, but neither has anything + # usable (a fresh OpenCode install with no provider credentials at + # all) -- still must not leave `model` unset. + chat = OpenCodeChat("http://127.0.0.1:4096", model=None, model_is_explicit=False) + responses = [ + self._fake_response({}), + self._fake_response({"providers": [{"id": "openrouter"}], "default": {}}), + ] + with mock.patch.object(chat, "_request", side_effect=lambda *a, **k: responses.pop(0)): + chat._ensure_model_resolved() + self.assertEqual(chat.model, "opencode/deepseek-v4-flash-free") + + def test_config_field_that_is_not_provider_slash_model_falls_through(self): + # A malformed/unexpected `model` field (missing the "/") is treated + # the same as absent, not used as-is. + chat = OpenCodeChat("http://127.0.0.1:4096", model=None, model_is_explicit=False) + responses = [ + self._fake_response({"model": "not-a-provider-model-pair"}), + self._fake_response({"providers": [{"id": "openrouter"}], "default": {"openrouter": "openai/gpt-5"}}), + ] + with mock.patch.object(chat, "_request", side_effect=lambda *a, **k: responses.pop(0)): + chat._ensure_model_resolved() + self.assertEqual(chat.model, "openrouter/openai/gpt-5") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/trainer/services/test_trainer_services_unit.py b/tests/trainer/services/test_trainer_services_unit.py index ef7ecbb6..cc8c1588 100644 --- a/tests/trainer/services/test_trainer_services_unit.py +++ b/tests/trainer/services/test_trainer_services_unit.py @@ -53,6 +53,35 @@ def test_get_latest_logger_data_full_history_nested(self): self.assertEqual(response.points[0].model_age, 0) self.assertEqual(response.points[1].model_age, 1) + def test_get_latest_logger_data_full_history_downsample_keeps_last_point(self): + # Regression test: a fixed-stride slice (signal_history[::step]) starts at + # index 0 and drops the tail of a fixed-length run. With 10000 points and + # max_points=1000, step=10 lands the last kept index at 9990, silently + # dropping steps 9991-9999. _downsample_uniform must be used instead so the + # most recent point is always present. + signal_logger = MagicMock() + signal_logger.get_signal_history.return_value = [ + { + "metric_name": "train/loss_CE", + "model_age": age, + "metric_value": 1.0 / (age + 1), + "experiment_hash": "exp-hash", + } + for age in range(10000) + ] + ctx = _DummyCtx(components={"signal_logger": signal_logger}) + with patch("weightslab.trainer.services.experiment_service.DataService"): + service = ExperimentService(ctx) + + request = pb2.GetLatestLoggerDataRequest(request_full_history=True, max_points=1000, break_by_slices=False) + response = service.GetLatestLoggerData(request, None) + + model_ages = [p.model_age for p in response.points] + self.assertLessEqual(len(model_ages), 1000) + self.assertEqual(model_ages[0], 0) + self.assertEqual(max(model_ages), 9999) + self.assertEqual(model_ages[-1], 9999) + def test_get_latest_logger_data_queue_mode(self): signal_logger = MagicMock() signal_logger.get_and_clear_queue.return_value = [ diff --git a/tests/ui/test_server_agent.py b/tests/ui/test_server_agent.py new file mode 100644 index 00000000..4782726e --- /dev/null +++ b/tests/ui/test_server_agent.py @@ -0,0 +1,514 @@ +"""Tests for weightslab/ui/server.py's OpenCode-agent supervisor: + +- POST /agent-server/start -- spawns (or reuses) a local OpenCode server rooted + at the experiment directory, so the browser never has to run `opencode serve` + by hand. Mirrors /local-notebook's "the browser can't spawn a process, so it + asks us to" shape. +- GET /agent-server/status -- none / running / killed, for the composer's status + line to poll. + +CI has no real `opencode` binary, so every test replaces +`ui_server._resolve_opencode_argv` with a tiny stand-in HTTP server (started via +`python -c`) that serves /global/health the same way the real OpenCode server +does. That is the only thing `_OpencodeSession` depends on to decide the child +started successfully, so it exercises the real spawn/health-poll/status code +path without depending on Node or the OpenCode package being installed. +""" + +import json +import os +import sys +import tempfile +import threading +import time +import unittest +import urllib.error +import urllib.request +from unittest.mock import patch + +from weightslab.ui import server as ui_server + +# A minimal stand-in for `opencode serve`: binds the --port it was given and +# answers /global/health like the real server does. Reads --port out of +# sys.argv positionally rather than assuming argv[0] is anything in particular, +# since "python -c serve --hostname H --port N" hands the child an argv +# whose exact shape depends on the platform's python launcher. +_FAKE_OPENCODE_SRC = r""" +import sys, json +from http.server import BaseHTTPRequestHandler, HTTPServer + +def _port(): + for i, a in enumerate(sys.argv): + if a == "--port": + return int(sys.argv[i + 1]) + return 4096 + +class H(BaseHTTPRequestHandler): + def log_message(self, *a): + pass + def do_GET(self): + if self.path == "/global/health": + body = json.dumps({"version": "0.0.0-fake"}).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + else: + self.send_response(404) + self.end_headers() + +HTTPServer(("127.0.0.1", _port()), H).serve_forever() +""" + +_FAKE_ARGV = [sys.executable, "-c", _FAKE_OPENCODE_SRC] + +# A stand-in that never becomes healthy -- models a child that starts but never +# binds (a bad flag, a crash loop) so ensure()'s timeout path is exercised. +_HANGING_ARGV = [sys.executable, "-c", "import time; time.sleep(60)"] + + +class TestOpencodeSessionUnit(unittest.TestCase): + """Exercises _OpencodeSession directly -- no HTTP layer, no real OpenCode.""" + + def setUp(self): + self.tmp = tempfile.mkdtemp() + self.session = ui_server._OpencodeSession() + # Keep polling fast so a genuine failure test doesn't sit for 45s. + self._timeout_patch = patch.object(ui_server, "_OPENCODE_START_TIMEOUT", 3.0) + self._timeout_patch.start() + + def tearDown(self): + self.session.shutdown() + self._timeout_patch.stop() + + def test_starts_and_reports_running(self): + with patch.object(ui_server, "_resolve_opencode_argv", return_value=_FAKE_ARGV): + result = self.session.ensure(self.tmp, "http://localhost:5173") + self.assertTrue(result["ok"], result) + self.assertFalse(result["reused"]) + self.assertEqual(result["workspace"], self.tmp) + self.assertTrue(result["url"].startswith("http://127.0.0.1:")) + + status = self.session.status() + self.assertEqual(status["state"], "running") + self.assertEqual(status["workspace"], self.tmp) + + def test_second_call_reuses_the_same_process(self): + with patch.object(ui_server, "_resolve_opencode_argv", return_value=_FAKE_ARGV): + first = self.session.ensure(self.tmp, "http://localhost:5173") + second = self.session.ensure(self.tmp, "http://localhost:5173") + self.assertFalse(first["reused"]) + self.assertTrue(second["reused"]) + self.assertEqual(first["url"], second["url"]) + + def test_missing_binary_and_missing_npx_reports_a_clear_error(self): + with patch.object(ui_server, "_resolve_opencode_argv", return_value=None): + result = self.session.ensure(self.tmp, "http://localhost:5173") + self.assertFalse(result["ok"]) + self.assertIn("opencode", result["error"].lower()) + + status = self.session.status() + self.assertEqual(status["state"], "none") + + def test_child_that_never_becomes_healthy_times_out_and_is_killed(self): + with patch.object(ui_server, "_resolve_opencode_argv", return_value=_HANGING_ARGV): + result = self.session.ensure(self.tmp, "http://localhost:5173") + self.assertFalse(result["ok"]) + self.assertIn("did not come up", result["error"]) + + # The timed-out child must actually be killed, not leaked. + status = self.session.status() + self.assertEqual(status["state"], "killed") + + def test_shutdown_stops_a_running_process(self): + with patch.object(ui_server, "_resolve_opencode_argv", return_value=_FAKE_ARGV): + result = self.session.ensure(self.tmp, "http://localhost:5173") + self.assertTrue(result["ok"]) + process = self.session._process + self.session.shutdown() + self.assertIsNotNone(process.poll()) # exited + + def test_ensure_drops_agents_md_into_a_fresh_workspace(self): + # This test process runs from an actual repo checkout, so AGENTS.md + # resolves for real via _repo_doc_path -- no mocking needed to prove + # ensure() actually reaches the workspace with it. + with patch.object(ui_server, "_resolve_opencode_argv", return_value=_FAKE_ARGV): + result = self.session.ensure(self.tmp, "http://localhost:5173") + self.assertTrue(result["ok"], result) + copied = os.path.join(self.tmp, "AGENTS.md") + self.assertTrue(os.path.isfile(copied)) + with open(copied, encoding="utf-8") as fh: + self.assertTrue(fh.read().strip()) + + def test_ensure_never_overwrites_a_workspace_own_agents_md(self): + # A workspace's own AGENTS.md might be the USER's project + # instructions -- ensure() must never clobber it with ours. + own_path = os.path.join(self.tmp, "AGENTS.md") + with open(own_path, "w", encoding="utf-8") as fh: + fh.write("this workspace's own instructions") + with patch.object(ui_server, "_resolve_opencode_argv", return_value=_FAKE_ARGV): + result = self.session.ensure(self.tmp, "http://localhost:5173") + self.assertTrue(result["ok"], result) + with open(own_path, encoding="utf-8") as fh: + self.assertEqual(fh.read(), "this workspace's own instructions") + + def test_ensure_copy_is_best_effort_when_no_source_agents_md_exists(self): + # No source to copy from (e.g. a stripped-down install) must not + # fail ensure() outright -- the agent server should still start. + with patch.object(ui_server, "_repo_doc_path", return_value=None): + with patch.object(ui_server, "_resolve_opencode_argv", return_value=_FAKE_ARGV): + result = self.session.ensure(self.tmp, "http://localhost:5173") + self.assertTrue(result["ok"], result) + self.assertFalse(os.path.isfile(os.path.join(self.tmp, "AGENTS.md"))) + + def test_successful_spawn_writes_a_lock_file_for_this_workspace(self): + # The other half of the cross-process handoff: the backend SDK agent + # (agent.py's OpenCodeChat) discovers THIS server via the same file + # -- see test_opencode_process.py's cross-process test for the full + # round trip. + with patch.object(ui_server, "_resolve_opencode_argv", return_value=_FAKE_ARGV): + result = self.session.ensure(self.tmp, "http://localhost:5173") + self.assertTrue(result["ok"], result) + lock = ui_server.opencode_process.read_lock(self.tmp) + self.assertEqual(lock["url"], result["url"]) + + def test_adopts_a_healthy_lockfile_instead_of_spawning_a_second_server(self): + # Simulates order (a) from the tabbed-agent-window plan: the backend + # SDK agent already published a server for this workspace before + # `weightslab start` (this session) ever called ensure(). + other = ui_server._OpencodeSession() + try: + with patch.object(ui_server, "_resolve_opencode_argv", return_value=_FAKE_ARGV): + published = other.ensure(self.tmp, None) + self.assertTrue(published["ok"], published) + + with patch.object(ui_server, "_resolve_opencode_argv") as argv_mock: + result = self.session.ensure(self.tmp, "http://localhost:5173") + finally: + other.shutdown() + + argv_mock.assert_not_called() + self.assertTrue(result["ok"], result) + self.assertEqual(result["url"], published["url"]) + self.assertEqual(result.get("adopted"), "lockfile") + # Adopted, not spawned -- nothing of this session's own to kill. + self.assertIsNone(self.session._process) + + def test_stale_lockfile_is_ignored_and_a_fresh_server_is_spawned(self): + # The process a stale lock file names is long gone -- must fall + # through to a normal spawn rather than failing or hanging. + ui_server.opencode_process.write_lock(self.tmp, "http://127.0.0.1:1", pid=999999) + with patch.object(ui_server, "_resolve_opencode_argv", return_value=_FAKE_ARGV): + result = self.session.ensure(self.tmp, "http://localhost:5173") + self.assertTrue(result["ok"], result) + self.assertNotEqual(result["url"], "http://127.0.0.1:1") + # The stale entry was overwritten with the newly-spawned server. + self.assertEqual(ui_server.opencode_process.read_lock(self.tmp)["url"], result["url"]) + + +class TestEnsureWorkspaceAgentsMd(unittest.TestCase): + """_ensure_workspace_agents_md directly -- no OpenCode process involved.""" + + def setUp(self): + self.tmp = tempfile.mkdtemp() + + def test_copies_from_the_installed_package_location(self): + ui_server._ensure_workspace_agents_md(self.tmp) + target = os.path.join(self.tmp, "AGENTS.md") + self.assertTrue(os.path.isfile(target)) + with open(target, encoding="utf-8") as fh: + content = fh.read() + self.assertEqual(content, ui_server._read_repo_doc("AGENTS.md")) + + def test_is_a_no_op_when_the_workspace_already_has_one(self): + target = os.path.join(self.tmp, "AGENTS.md") + with open(target, "w", encoding="utf-8") as fh: + fh.write("mine") + ui_server._ensure_workspace_agents_md(self.tmp) + with open(target, encoding="utf-8") as fh: + self.assertEqual(fh.read(), "mine") + + def test_does_nothing_when_no_source_is_found(self): + with patch.object(ui_server, "_repo_doc_path", return_value=None): + ui_server._ensure_workspace_agents_md(self.tmp) + self.assertFalse(os.path.isfile(os.path.join(self.tmp, "AGENTS.md"))) + + +class TestCorsOriginVariants(unittest.TestCase): + """The localhost <-> 127.0.0.1 expansion is the #1 way this feature goes + silently wrong -- a mismatch here makes every request look like the agent + server is simply not there.""" + + def test_expands_localhost_to_127_0_0_1(self): + variants = ui_server._cors_origin_variants("http://localhost:5173") + self.assertIn("http://localhost:5173", variants) + self.assertIn("http://127.0.0.1:5173", variants) + + def test_expands_127_0_0_1_to_localhost(self): + variants = ui_server._cors_origin_variants("http://127.0.0.1:8080") + self.assertIn("http://127.0.0.1:8080", variants) + self.assertIn("http://localhost:8080", variants) + + def test_leaves_a_non_loopback_origin_alone(self): + variants = ui_server._cors_origin_variants("https://weightslab.example.com") + self.assertEqual(variants, ["https://weightslab.example.com"]) + + def test_none_origin_yields_no_variants(self): + self.assertEqual(ui_server._cors_origin_variants(None), []) + + +class _ServerTestCase(unittest.TestCase): + """Spins up a real serve_ui() on 127.0.0.1: per test, rooted at a + fresh temp dir passed as experiment_dir -- same shape as + test_server_experiment_reports.py.""" + + def setUp(self): + self.tmp = tempfile.mkdtemp() + self.httpd = ui_server.serve_ui( + ui_host="127.0.0.1", ui_port=0, + backend_host="localhost", backend_port=50051, + open_browser=False, block=False, + experiment_dir=self.tmp, + ) + self.port = self.httpd.server_address[1] + self.thread = threading.Thread(target=self.httpd.serve_forever, daemon=True) + self.thread.start() + time.sleep(0.1) + + # The module-level singleton is shared across the whole test process; + # swap in a fresh one per test so a leftover child from another test + # can't make this one see {reused: true} unexpectedly. + self._orig_session = ui_server._opencode_session + ui_server._opencode_session = ui_server._OpencodeSession() + self._timeout_patch = patch.object(ui_server, "_OPENCODE_START_TIMEOUT", 3.0) + self._timeout_patch.start() + + def tearDown(self): + self._timeout_patch.stop() + ui_server._opencode_session.shutdown() + ui_server._opencode_session = self._orig_session + self.httpd.shutdown() + self.thread.join(timeout=5) + + def _get(self, path): + return urllib.request.urlopen(f"http://127.0.0.1:{self.port}{path}", timeout=5) + + def _post(self, path, origin=None): + req = urllib.request.Request( + f"http://127.0.0.1:{self.port}{path}", method="POST", data=b"", + ) + if origin: + req.add_header("Origin", origin) + return urllib.request.urlopen(req, timeout=10) + + +class TestAgentServerEndpoint(_ServerTestCase): + + def test_status_is_none_before_anything_starts(self): + with self._get("/agent-server/status") as r: + data = json.loads(r.read().decode()) + self.assertEqual(data["state"], "none") + + def test_start_spawns_rooted_at_the_experiment_dir(self): + with patch.object(ui_server, "_resolve_opencode_argv", return_value=_FAKE_ARGV): + with self._post("/agent-server/start", origin="http://localhost:5173") as r: + data = json.loads(r.read().decode()) + self.assertTrue(data["ok"], data) + self.assertEqual(data["workspace"], self.tmp) + + with self._get("/agent-server/status") as r: + status = json.loads(r.read().decode()) + self.assertEqual(status["state"], "running") + + def test_missing_opencode_and_npx_returns_a_clear_error_not_a_500_crash(self): + with patch.object(ui_server, "_resolve_opencode_argv", return_value=None): + with self.assertRaises(urllib.error.HTTPError) as ctx: + self._post("/agent-server/start", origin="http://localhost:5173") + self.assertEqual(ctx.exception.code, 500) + data = json.loads(ctx.exception.read().decode()) + self.assertFalse(data["ok"]) + self.assertIn("opencode", data["error"].lower()) + + def test_falls_back_to_reconstructing_origin_from_host_header(self): + # No Origin header at all (e.g. a same-origin fetch some browsers omit + # it for) -- must not crash, and must still start the server. + with patch.object(ui_server, "_resolve_opencode_argv", return_value=_FAKE_ARGV): + with self._post("/agent-server/start") as r: + data = json.loads(r.read().decode()) + self.assertTrue(data["ok"], data) + + +class TestAgentDocsEndpoint(_ServerTestCase): + """GET /agent-server/docs[?example=] -- AGENTS.md, plus + optionally one matching PyTorch usecase example, for the landing chat's + preset prompts to attach (see agentChat.ts's PRESET_PROMPTS). README.md + was dropped from this endpoint on purpose -- AGENTS.md alone carries the + weightslab integration pattern the presets need.""" + + def test_returns_agents_md_when_present_in_a_repo_checkout(self): + # This test process runs from an actual repo checkout, so it should + # resolve via _read_repo_doc. + with self._get("/agent-server/docs") as r: + data = json.loads(r.read().decode()) + names = {f["name"] for f in data["files"]} + self.assertEqual(names, {"AGENTS.md"}) + for f in data["files"]: + self.assertTrue(f["content"].strip()) + + def test_omits_the_doc_when_it_cannot_be_found_instead_of_erroring(self): + with patch.object(ui_server, "_read_repo_doc", return_value=None): + with self._get("/agent-server/docs") as r: + data = json.loads(r.read().decode()) + self.assertEqual(data["files"], []) + + def test_example_query_param_additionally_attaches_that_usecases_main_py(self): + with self._get("/agent-server/docs?example=wl-classification") as r: + data = json.loads(r.read().decode()) + names = {f["name"] for f in data["files"]} + self.assertEqual(names, {"AGENTS.md", "examples/PyTorch/wl-classification/main.py"}) + + def test_example_query_param_is_repeatable_for_multiple_usecases(self): + with self._get("/agent-server/docs?example=wl-detection&example=wl-segmentation") as r: + data = json.loads(r.read().decode()) + names = {f["name"] for f in data["files"]} + self.assertEqual(names, { + "AGENTS.md", + "examples/PyTorch/wl-detection/main.py", + "examples/PyTorch/wl-segmentation/main.py", + }) + + def test_example_query_param_is_ignored_when_not_a_known_usecase(self): + # Client-supplied -- must be checked against the allowlist, never + # trusted as a path component (e.g. "../../../etc/passwd"). + with self._get("/agent-server/docs?example=../../../etc/passwd") as r: + data = json.loads(r.read().decode()) + names = {f["name"] for f in data["files"]} + self.assertEqual(names, {"AGENTS.md"}) + + def test_no_example_query_param_means_no_example_file(self): + with self._get("/agent-server/docs") as r: + data = json.loads(r.read().decode()) + names = {f["name"] for f in data["files"]} + self.assertNotIn("examples/PyTorch/wl-classification/main.py", names) + + +class TestLoopRegistryUnit(unittest.TestCase): + """_LoopRegistry directly -- no HTTP server, no real OpenCode process. + Mocks the module-level _opencode_json_request/_opencode_send_and_collect/ + _opencode_get_messages functions the registry itself calls, so these + exercise its own eager-session-creation/locking/preamble-once logic in + isolation. A fresh registry per test, not the module-level singleton.""" + + def setUp(self): + self.registry = ui_server._LoopRegistry() + + def tearDown(self): + for job in self.registry._jobs.values(): + if job.timer is not None: + job.timer.cancel() + + def test_start_creates_the_session_eagerly_not_on_first_tick(self): + with patch.object(ui_server, "_opencode_session") as mock_session, \ + patch.object(ui_server, "_opencode_json_request", return_value={"id": "sess-1"}) as mock_req, \ + patch.object(ui_server, "_opencode_send_and_collect", return_value="ok"): + mock_session.ensure.return_value = {"ok": True, "url": "http://fake"} + result = self.registry.start("monitor training", 60.0, "/tmp/ws", "http://localhost:5173") + + self.assertTrue(result["ok"], result) + job = self.registry._jobs[result["id"]] + # Eager: already set by start() itself, not left for _fire's first + # tick (which runs in a background thread started right after). + self.assertEqual(job.session_id, "sess-1") + mock_req.assert_any_call("http://fake", "/session", method="POST", body=unittest.mock.ANY) + + def test_rejected_after_session_creation_deletes_the_orphaned_session(self): + def _ensure_and_fill_concurrently(*_args, **_kwargs): + # Simulates 3 OTHER starts winning the race while this one's + # ensure() call was in flight -- by the time start() re-checks + # under the lock, the cap has already been hit by them. + for i in range(3): + self.registry._jobs[str(i)] = ui_server._LoopJob(str(i), "p", 60.0, "/tmp") + return {"ok": True, "url": "http://fake"} + + with patch.object(ui_server, "_opencode_session") as mock_session, \ + patch.object(ui_server, "_opencode_json_request", return_value={"id": "sess-orphan"}) as mock_req: + mock_session.ensure.side_effect = _ensure_and_fill_concurrently + result = self.registry.start("monitor training", 60.0, "/tmp/ws", "http://localhost:5173") + + self.assertFalse(result["ok"]) + self.assertIn("already running", result["error"]) + mock_req.assert_any_call("http://fake", "/session/sess-orphan", method="DELETE") + + def test_fire_sends_the_preamble_once_then_plain_prompt_on_later_ticks(self): + job = ui_server._LoopJob("1", "check the loss", 60.0, "/tmp") + job.session_id, job.base_url = "sess-1", "http://fake" + self.registry._jobs["1"] = job + + sent_texts = [] + + def _fake_send(_base_url, _session_id, text, _model=None, timeout=600.0): # noqa: ARG001 + sent_texts.append(text) + return "tick result", None + + with patch.object(ui_server, "_opencode_send_and_collect", side_effect=_fake_send): + self.registry._fire("1", "http://fake") + self.assertTrue(job.preamble_sent) + self.assertIn("recurring monitoring agent", sent_texts[0]) + self.assertIn("check the loss", sent_texts[0]) + job.timer.cancel() + + self.registry._fire("1", "http://fake") + self.assertEqual(sent_texts[1], "check the loss") + job.timer.cancel() + + def test_get_messages_success(self): + job = ui_server._LoopJob("1", "p", 60.0, "/tmp") + job.session_id, job.base_url = "sess-1", "http://fake" + self.registry._jobs["1"] = job + canned = [{"info": {"role": "user"}, "parts": [{"type": "text", "text": "hi"}]}] + with patch.object(ui_server, "_opencode_get_messages", return_value=canned): + result = self.registry.get_messages("1") + self.assertTrue(result["ok"], result) + self.assertEqual(result["messages"], canned) + + def test_get_messages_unknown_job(self): + result = self.registry.get_messages("nope") + self.assertFalse(result["ok"]) + + +class TestLoopMessagesEndpoint(_ServerTestCase): + """GET /agent-server/loop//messages -- a loop tab's read-only + transcript, proxied through the module-level _loop_registry singleton to + a job's own OpenCode session (mocked here; no real OpenCode process). + Loopback-gating itself mirrors the existing loop routes (_stop_loop et + al.), which have no dedicated test for the negative case either -- a + real test client is always loopback.""" + + def _seed_job(self): + job = ui_server._LoopJob("1", "check the loss", 60.0, self.tmp) + job.session_id, job.base_url = "sess-1", "http://fake" + ui_server._loop_registry._jobs["1"] = job + return job + + def tearDown(self): + ui_server._loop_registry._jobs.clear() + super().tearDown() + + def test_get_messages_returns_the_sessions_history(self): + self._seed_job() + canned = [{"info": {"role": "assistant"}, "parts": [{"type": "text", "text": "hi"}]}] + with patch.object(ui_server, "_opencode_get_messages", return_value=canned): + with self._get("/agent-server/loop/1/messages") as r: + data = json.loads(r.read().decode()) + self.assertTrue(data["ok"], data) + self.assertEqual(data["messages"], canned) + + def test_get_messages_404s_for_an_unknown_loop(self): + with self.assertRaises(urllib.error.HTTPError) as ctx: + self._get("/agent-server/loop/999/messages") + self.assertEqual(ctx.exception.code, 404) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/ui/test_server_data_query.py b/tests/ui/test_server_data_query.py new file mode 100644 index 00000000..3dacdc74 --- /dev/null +++ b/tests/ui/test_server_data_query.py @@ -0,0 +1,211 @@ +"""Tests for weightslab/ui/server.py's POST /agent-server/data-query -- +lets the landing-page agent chat perform dataset/model actions (discard, +tag, sort, filter, analyze, compute stats, ...) itself, the same way the +now-retired "Backend Agent" tab's query bar always did: by calling +ExperimentService.ApplyDataQuery over the SAME upstream gRPC channel +_proxy_grpc_web already proxies everything else through. + +Spins up a REAL grpc.server() implementing just ApplyDataQuery (not a mock) +alongside a real serve_ui() instance pointed at it -- same "real subprocess/ +real network, not mocked" philosophy as test_server_agent.py's fake-OpenCode +HTTP server -- so this exercises the actual request-building/response- +translation code, not just its shape. +""" + +import json +import tempfile +import threading +import time +import unittest +import urllib.error +import urllib.request +from concurrent import futures + +import grpc + +import weightslab.proto.experiment_service_pb2 as pb2 +import weightslab.proto.experiment_service_pb2_grpc as pb2_grpc +from weightslab.ui import server as ui_server + + +class _FakeExperimentService(pb2_grpc.ExperimentServiceServicer): + """Records every request it receives and returns whatever this test set + as `.next_response` (or raises `.next_error` instead, if set).""" + + def __init__(self): + self.received = [] + self.next_response = pb2.DataQueryResponse(success=True, message="ok") + self.next_error = None + + def ApplyDataQuery(self, request, context): + self.received.append(request) + if self.next_error is not None: + context.abort(self.next_error[0], self.next_error[1]) + return self.next_response + + +class _ServerTestCase(unittest.TestCase): + """Real serve_ui() + a real (fake) ExperimentService gRPC server behind + it, both on ephemeral ports.""" + + def setUp(self): + self.tmp = tempfile.mkdtemp() + + self.servicer = _FakeExperimentService() + self.grpc_server = grpc.server(futures.ThreadPoolExecutor(max_workers=2)) + pb2_grpc.add_ExperimentServiceServicer_to_server(self.servicer, self.grpc_server) + backend_port = self.grpc_server.add_insecure_port("127.0.0.1:0") + self.grpc_server.start() + + self.httpd = ui_server.serve_ui( + ui_host="127.0.0.1", ui_port=0, + backend_host="127.0.0.1", backend_port=backend_port, + open_browser=False, block=False, + experiment_dir=self.tmp, + ) + self.port = self.httpd.server_address[1] + self.thread = threading.Thread(target=self.httpd.serve_forever, daemon=True) + self.thread.start() + time.sleep(0.1) + + def tearDown(self): + self.httpd.shutdown() + self.thread.join(timeout=5) + self.grpc_server.stop(grace=None) + + def _post_json(self, path, body): + data = json.dumps(body).encode("utf-8") + req = urllib.request.Request( + f"http://127.0.0.1:{self.port}{path}", method="POST", data=data, + headers={"Content-Type": "application/json"}, + ) + return urllib.request.urlopen(req, timeout=10) + + def _get_json(self, path): + with urllib.request.urlopen(f"http://127.0.0.1:{self.port}{path}", timeout=10) as r: + return json.loads(r.read().decode()) + + +class TestDataQueryEndpoint(_ServerTestCase): + def test_builds_a_real_dataqueryrequest_and_translates_the_response(self): + self.servicer.next_response = pb2.DataQueryResponse( + success=True, message="Discarded 3 samples.", + number_of_all_samples=100, number_of_samples_in_the_loop=97, + number_of_discarded_samples=3, unique_tags=["reviewed"], + analysis_result="", + ) + + with self._post_json("/agent-server/data-query", {"query": "discard samples where loss > 5"}) as r: + data = json.loads(r.read().decode()) + + self.assertTrue(data["ok"]) + self.assertEqual(data["message"], "Discarded 3 samples.") + self.assertEqual(data["numberOfAllSamples"], 100) + self.assertEqual(data["numberOfSamplesInTheLoop"], 97) + self.assertEqual(data["numberOfDiscardedSamples"], 3) + self.assertEqual(data["uniqueTags"], ["reviewed"]) + + self.assertEqual(len(self.servicer.received), 1) + sent = self.servicer.received[0] + self.assertEqual(sent.query, "discard samples where loss > 5") + self.assertFalse(sent.accumulate) + self.assertTrue(sent.is_natural_language) + + def test_accumulate_flag_is_forwarded(self): + with self._post_json("/agent-server/data-query", {"query": "sort by loss", "accumulate": True}) as r: + r.read() + self.assertTrue(self.servicer.received[0].accumulate) + + def test_backend_reported_failure_is_not_an_http_error(self): + # The backend understood the request fine and answered -- it just + # couldn't do what was asked (ambiguous, out of scope, ...). That's + # a normal 200 with success=false, not a transport-level failure. + self.servicer.next_response = pb2.DataQueryResponse( + success=False, message="I don't understand which column you mean.", + ) + + with self._post_json("/agent-server/data-query", {"query": "do the thing"}) as r: + data = json.loads(r.read().decode()) + + self.assertFalse(data["ok"]) + self.assertIn("don't understand", data["message"]) + self.assertNotIn("error", data) + + def test_empty_query_is_rejected_without_reaching_the_backend(self): + try: + self._post_json("/agent-server/data-query", {"query": " "}) + self.fail("expected an HTTPError") + except urllib.error.HTTPError as exc: + self.assertEqual(exc.code, 400) + data = json.loads(exc.read().decode()) + self.assertFalse(data["ok"]) + self.assertIn("required", data["error"]) + self.assertEqual(self.servicer.received, []) + + def test_grpc_failure_reports_a_clear_error_not_a_stack_trace(self): + self.servicer.next_error = (grpc.StatusCode.INTERNAL, "dataframe not loaded yet") + + try: + self._post_json("/agent-server/data-query", {"query": "sort by loss"}) + self.fail("expected an HTTPError") + except urllib.error.HTTPError as exc: + self.assertEqual(exc.code, 500) + data = json.loads(exc.read().decode()) + self.assertFalse(data["ok"]) + self.assertIn("dataframe not loaded yet", data["error"]) + + +class TestLatestDataQueryEndpoint(_ServerTestCase): + """The agent's own bash/curl call to /agent-server/data-query never + touches the browser's JS -- agentChat.ts instead polls + GET /agent-server/data-query/latest once per finished turn to find out + a query ran and replay the grid-refresh/subview-banner reaction. See + _LatestDataQuery's docstring in server.py. + + _latest_data_query is a module-level singleton (shared by every + serve_ui() instance in this process), so it's swapped out per-test the + same way test_server_tracked_processes.py does for _tracked_processes -- + otherwise seq/records would leak across tests.""" + + def setUp(self): + super().setUp() + self._orig_latest_data_query = ui_server._latest_data_query + ui_server._latest_data_query = ui_server._LatestDataQuery() + + def tearDown(self): + ui_server._latest_data_query = self._orig_latest_data_query + super().tearDown() + + def test_reports_seq_zero_when_nothing_has_run_yet(self): + data = self._get_json("/agent-server/data-query/latest") + self.assertEqual(data, {"seq": 0}) + + def test_reflects_the_most_recent_data_query_call(self): + self.servicer.next_response = pb2.DataQueryResponse( + success=True, message="Filtered to label 7.", + number_of_all_samples=100, number_of_samples_in_the_loop=12, + ) + with self._post_json("/agent-server/data-query", {"query": "show only label 7"}) as r: + r.read() + + data = self._get_json("/agent-server/data-query/latest") + self.assertEqual(data["seq"], 1) + self.assertEqual(data["query"], "show only label 7") + self.assertTrue(data["ok"]) + self.assertEqual(data["message"], "Filtered to label 7.") + self.assertEqual(data["numberOfAllSamples"], 100) + self.assertEqual(data["numberOfSamplesInTheLoop"], 12) + + def test_seq_increments_on_each_new_call_so_the_frontend_can_dedupe(self): + with self._post_json("/agent-server/data-query", {"query": "first"}) as r: + r.read() + with self._post_json("/agent-server/data-query", {"query": "second"}) as r: + r.read() + + data = self._get_json("/agent-server/data-query/latest") + self.assertEqual(data["seq"], 2) + self.assertEqual(data["query"], "second") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/ui/test_server_loop.py b/tests/ui/test_server_loop.py new file mode 100644 index 00000000..7f602df6 --- /dev/null +++ b/tests/ui/test_server_loop.py @@ -0,0 +1,487 @@ +"""Tests for weightslab/ui/server.py's /loop feature -- recurring +OpenCode-backed monitoring jobs (`_LoopRegistry` + the +/agent-server/loop/{start,list,stop} endpoints). + +_LoopRegistry.start() delegates session/message plumbing to the module-level +_opencode_json_request/_opencode_send_and_collect helpers (already covered at +the wire-protocol level by tests/trainer/services/test_opencode_chat.py's +fake-SSE-server tests for the sibling Python client). Here those two helpers +are mocked so the tests exercise _LoopRegistry's OWN logic -- validation, +session reuse across ticks, error bookkeeping, stop/list, and the HTTP +wiring -- the same split test_server_agent.py uses (direct _OpencodeSession +unit tests, then a separate endpoint-level test class). +""" + +import json +import tempfile +import threading +import time +import unittest +import urllib.request +from unittest.mock import patch + +from weightslab.ui import server as ui_server + + +class TestLoopRegistryUnit(unittest.TestCase): + """Exercises _LoopRegistry directly, with the opencode session assumed + already up (_opencode_session.ensure mocked) and the session-create/ + send-and-collect wire calls mocked -- those are covered elsewhere (see + module docstring).""" + + def setUp(self): + self.registry = ui_server._LoopRegistry() + self._ensure_patch = patch.object( + ui_server._opencode_session, "ensure", + return_value={"ok": True, "url": "http://127.0.0.1:1", "workspace": "/tmp", "reused": False}, + ) + self._ensure_patch.start() + + def tearDown(self): + self.registry.shutdown() + self._ensure_patch.stop() + + def _wait_for_first_tick(self, job_id, timeout=2.0): + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + jobs = {j["id"]: j for j in self.registry.list()} + job = jobs.get(job_id) + if job and (job["lastResult"] is not None or job["lastError"] is not None): + return job + time.sleep(0.01) + self.fail("loop job never completed its first tick") + + def test_model_resolution_prefers_the_caller_then_config_then_provider_defaults(self): + """The three sources, in order -- see _opencode_resolve_model. The + point of the chain is that a loop is configured from the SAME place + the chat is, rather than needing to be told by whichever surface + happened to start it.""" + explicit = {"providerID": "openrouter", "modelID": "anthropic/claude-haiku-4.5"} + + # 1. explicit wins outright -- no lookups at all. + with patch.object(ui_server, "_opencode_json_request") as req: + self.assertEqual(ui_server._opencode_resolve_model("http://fake", explicit), explicit) + req.assert_not_called() + + # 2. opencode.json's own default (what the chat's picker writes back). + def _config_has_model(_base, path, **_kw): + if path == "/config": + return {"model": "openrouter/anthropic/claude-haiku-4.5"} + raise AssertionError(f"should not have reached {path}") + + with patch.object(ui_server, "_opencode_json_request", side_effect=_config_has_model): + self.assertEqual(ui_server._opencode_resolve_model("http://fake", None), explicit) + + # 3. provider defaults, when the config names no model of its own. + def _only_provider_defaults(_base, path, **_kw): + if path == "/config": + return {} + return {"providers": [{"id": "openrouter"}], "default": {"openrouter": "openai/gpt-5"}} + + with patch.object(ui_server, "_opencode_json_request", side_effect=_only_provider_defaults): + self.assertEqual( + ui_server._opencode_resolve_model("http://fake", None), + {"providerID": "openrouter", "modelID": "openai/gpt-5"}, + ) + + # Nothing reachable -- no model, and OpenCode decides per check-in. + with patch.object(ui_server, "_opencode_json_request", side_effect=OSError("down")): + self.assertIsNone(ui_server._opencode_resolve_model("http://fake", None)) + + def test_rejects_an_empty_prompt(self): + result = self.registry.start(" ", 120, "/tmp", None) + self.assertFalse(result["ok"]) + self.assertIn("prompt", result["error"]) + self.assertEqual(self.registry.list(), []) + + def test_rejects_an_interval_below_the_minimum(self): + result = self.registry.start("watch training", 10, "/tmp", None) + self.assertFalse(result["ok"]) + self.assertIn("Minimum", result["error"]) + + def test_surfaces_an_ensure_failure_without_starting_a_job(self): + self._ensure_patch.stop() + try: + with patch.object(ui_server._opencode_session, "ensure", + return_value={"ok": False, "error": "no opencode binary"}): + result = self.registry.start("watch training", 120, "/tmp", None) + finally: + self._ensure_patch.start() + + self.assertFalse(result["ok"]) + self.assertEqual(result["error"], "no opencode binary") + self.assertEqual(self.registry.list(), []) + + def test_first_tick_creates_a_session_seeded_with_the_system_preamble(self): + # Model resolution goes through _opencode_json_request too (it reads + # OpenCode's config) -- stubbed out so create_mock below is only ever + # the session-creation call this test is actually about. + with patch.object(ui_server, "_opencode_json_request", return_value={"id": "ses_abc"}) as create_mock, \ + patch.object(ui_server, "_opencode_resolve_model", return_value=None), \ + patch.object(ui_server, "_opencode_send_and_collect", return_value=("all good", None)) as send_mock: + result = self.registry.start("watch the loss", 120, "/tmp", "http://localhost:5173") + self.assertTrue(result["ok"], result) + job = self._wait_for_first_tick(result["id"]) + + self.assertEqual(job["lastResult"], "all good") + self.assertIsNone(job["lastError"]) + create_mock.assert_called_once() + self.assertEqual(create_mock.call_args.args[1], "/session") + sent_text = send_mock.call_args.args[2] + self.assertIn("watch the loss", sent_text) + self.assertIn("pause / resume", sent_text) # system preamble documents the CLI verbs + # A detached relaunch has no OS-level tie to this workspace, so the + # preamble must tell the model to register it -- with THIS job's own + # origin, not a placeholder -- for Ctrl+C-on-the-workspace cleanup. + self.assertIn("http://localhost:5173/agent-server/track-process", sent_text) + self.assertIn("-PassThru", sent_text) + + def test_second_tick_reuses_the_session_and_sends_the_bare_prompt(self): + with patch.object(ui_server, "_LOOP_MIN_INTERVAL_SECONDS", 0.01), \ + patch.object(ui_server, "_opencode_json_request", return_value={"id": "ses_abc"}) as create_mock, \ + patch.object(ui_server, "_opencode_resolve_model", return_value=None), \ + patch.object(ui_server, "_opencode_send_and_collect", return_value=("ok", None)) as send_mock: + result = self.registry.start("watch the loss", 0.02, "/tmp", None) + self._wait_for_first_tick(result["id"]) + + deadline = time.monotonic() + 2 + while send_mock.call_count < 2 and time.monotonic() < deadline: + time.sleep(0.01) + self.registry.stop(result["id"]) + + self.assertGreaterEqual(send_mock.call_count, 2) + create_mock.assert_called_once() # session created once, reused on tick 2 + second_call_text = send_mock.call_args_list[1].args[2] + self.assertEqual(second_call_text, "watch the loss") # no preamble wrapper on repeat ticks + + def test_records_last_error_and_does_not_set_last_result_on_a_failed_tick(self): + with patch.object(ui_server, "_opencode_json_request", return_value={"id": "ses_abc"}), \ + patch.object(ui_server, "_opencode_send_and_collect", side_effect=RuntimeError("boom")): + result = self.registry.start("watch training", 120, "/tmp", None) + job = self._wait_for_first_tick(result["id"]) + + self.assertEqual(job["lastError"], "boom") + self.assertIsNone(job["lastResult"]) + + def test_records_the_reported_error_when_a_tick_ends_without_raising(self): + """A turn can end via session.error having produced no text at all -- + a provider rejecting the request outright (e.g. a model with no + tool-use endpoints, against a prompt that hands the agent a full + toolset). Nothing raises, so this used to leave lastError None and + lastResult "": the loop's tab showed the check-in prompt with silence + under it, every interval, with nothing anywhere saying why.""" + reported = 'No endpoints found that support tool use.' + with patch.object(ui_server, "_opencode_json_request", return_value={"id": "ses_abc"}), \ + patch.object(ui_server, "_opencode_send_and_collect", return_value=("", reported)): + result = self.registry.start("watch training", 120, "/tmp", None) + job = self._wait_for_first_tick(result["id"]) + + self.assertEqual(job["lastError"], reported) + + def test_explains_an_empty_reply_that_reported_no_error_at_all(self): + with patch.object(ui_server, "_opencode_json_request", return_value={"id": "ses_abc"}), \ + patch.object(ui_server, "_opencode_send_and_collect", return_value=(" ", None)): + result = self.registry.start("watch training", 120, "/tmp", None) + job = self._wait_for_first_tick(result["id"]) + + self.assertIsNotNone(job["lastError"]) + self.assertIn("no reply", job["lastError"]) + + def test_passes_the_requested_model_through_to_every_check_in(self): + model = {"providerID": "openrouter", "modelID": "anthropic/claude-haiku-4.5"} + with patch.object(ui_server, "_opencode_json_request", return_value={"id": "ses_abc"}), \ + patch.object(ui_server, "_opencode_resolve_model", side_effect=lambda _u, m: m), \ + patch.object(ui_server, "_opencode_send_and_collect", return_value=("ok", None)) as send_mock: + result = self.registry.start("watch training", 120, "/tmp", None, model) + job = self._wait_for_first_tick(result["id"]) + + self.assertEqual(send_mock.call_args.args[3], model) + self.assertEqual(job["model"], "openrouter/anthropic/claude-haiku-4.5") + + def test_resolves_a_model_from_opencode_config_when_the_caller_offers_none(self): + """A loop started from a surface with no model picker of its own (or + with none chosen yet) still gets a concrete model: whatever the chat's + picker last wrote into opencode.json. Resolved ONCE at start and + pinned, so a model changed in the chat afterwards doesn't silently + change what an already-running job has been reporting.""" + resolved = {"providerID": "openrouter", "modelID": "anthropic/claude-haiku-4.5"} + with patch.object(ui_server, "_opencode_json_request", return_value={"id": "ses_abc"}), \ + patch.object(ui_server, "_opencode_resolve_model", return_value=resolved) as resolve_mock, \ + patch.object(ui_server, "_opencode_send_and_collect", return_value=("ok", None)) as send_mock: + result = self.registry.start("watch training", 120, "/tmp", None, None) + job = self._wait_for_first_tick(result["id"]) + + resolve_mock.assert_called_once() + self.assertIsNone(resolve_mock.call_args.args[1]) # nothing explicit to prefer + self.assertEqual(send_mock.call_args.args[3], resolved) + self.assertEqual(job["model"], "openrouter/anthropic/claude-haiku-4.5") + + def test_a_job_reports_running_only_while_a_check_in_is_in_flight(self): + """The tab has no other way to tell "the agent is working on this + right now" from "nothing is happening": next_run_at deliberately only + moves once a run FINISHES, so during one it still holds the previous + run's already-elapsed value.""" + started = threading.Event() + release = threading.Event() + + def _slow(*_args, **_kwargs): + started.set() + release.wait(timeout=5) + return ("done", None) + + with patch.object(ui_server, "_opencode_json_request", return_value={"id": "ses_abc"}), \ + patch.object(ui_server, "_opencode_send_and_collect", side_effect=_slow): + result = self.registry.start("watch training", 120, "/tmp", None) + job_id = result["id"] + self.assertTrue(started.wait(timeout=5)) + + in_flight = next(j for j in self.registry.list() if j["id"] == job_id) + self.assertTrue(in_flight["running"]) + # ...and the countdown has not been moved yet, which is exactly + # why `running` has to be reported separately. + self.assertIsNone(in_flight["nextRunAt"]) + + release.set() + job = self._wait_for_first_tick(job_id) + + self.assertFalse(job["running"]) + self.assertIsNotNone(job["nextRunAt"]) + + def test_stop_removes_the_job_and_prevents_a_further_tick(self): + with patch.object(ui_server, "_opencode_json_request", return_value={"id": "ses_abc"}), \ + patch.object(ui_server, "_opencode_send_and_collect", return_value=("ok", None)) as send_mock: + result = self.registry.start("watch training", 120, "/tmp", None) + self._wait_for_first_tick(result["id"]) + + stop_result = self.registry.stop(result["id"]) + self.assertTrue(stop_result["ok"]) + self.assertEqual(self.registry.list(), []) + + calls_at_stop = send_mock.call_count + time.sleep(0.1) + self.assertEqual(send_mock.call_count, calls_at_stop) # no tick after stop + + def test_stopping_an_unknown_job_id_reports_not_found(self): + result = self.registry.stop("does-not-exist") + self.assertFalse(result["ok"]) + + def test_list_reflects_multiple_concurrent_jobs(self): + with patch.object(ui_server, "_opencode_json_request", return_value={"id": "ses_abc"}), \ + patch.object(ui_server, "_opencode_send_and_collect", return_value=("ok", None)): + first = self.registry.start("watch a", 120, "/tmp", None) + second = self.registry.start("watch b", 120, "/tmp", None) + self._wait_for_first_tick(first["id"]) + self._wait_for_first_tick(second["id"]) + + ids = {j["id"] for j in self.registry.list()} + self.assertEqual(ids, {first["id"], second["id"]}) + + def test_shutdown_stops_every_job(self): + with patch.object(ui_server, "_opencode_json_request", return_value={"id": "ses_abc"}), \ + patch.object(ui_server, "_opencode_send_and_collect", return_value=("ok", None)): + result = self.registry.start("watch training", 120, "/tmp", None) + self._wait_for_first_tick(result["id"]) + + self.registry.shutdown() + self.assertEqual(self.registry.list(), []) + + def test_rejects_a_fourth_concurrent_job(self): + with patch.object(ui_server, "_opencode_json_request", return_value={"id": "ses_abc"}), \ + patch.object(ui_server, "_opencode_send_and_collect", return_value=("ok", None)): + for i in range(ui_server._LOOP_MAX_CONCURRENT): + result = self.registry.start(f"watch {i}", 120, "/tmp", None) + self.assertTrue(result["ok"], result) + + fourth = self.registry.start("one too many", 120, "/tmp", None) + + self.assertFalse(fourth["ok"]) + self.assertIn("already running", fourth["error"]) + self.assertEqual(len(self.registry.list()), ui_server._LOOP_MAX_CONCURRENT) + + def test_update_changes_the_prompt_without_touching_the_interval(self): + with patch.object(ui_server, "_opencode_json_request", return_value={"id": "ses_abc"}), \ + patch.object(ui_server, "_opencode_send_and_collect", return_value=("ok", None)): + result = self.registry.start("watch training", 120, "/tmp", None) + self._wait_for_first_tick(result["id"]) + + update_result = self.registry.update(result["id"], prompt="watch training more closely") + + self.assertTrue(update_result["ok"], update_result) + job = {j["id"]: j for j in self.registry.list()}[result["id"]] + self.assertEqual(job["prompt"], "watch training more closely") + self.assertEqual(job["intervalSeconds"], 120) + + def test_update_changes_the_interval_and_reschedules_immediately(self): + with patch.object(ui_server, "_opencode_json_request", return_value={"id": "ses_abc"}), \ + patch.object(ui_server, "_opencode_send_and_collect", return_value=("ok", None)) as send_mock: + result = self.registry.start("watch training", 120, "/tmp", None) + self._wait_for_first_tick(result["id"]) + + with patch.object(ui_server, "_LOOP_MIN_INTERVAL_SECONDS", 0.01): + update_result = self.registry.update(result["id"], interval_seconds=0.02) + self.assertTrue(update_result["ok"], update_result) + self.assertEqual(update_result["intervalSeconds"], 0.02) + + # If the reschedule took effect immediately, a second tick lands + # almost at once rather than after the original 120s interval. + deadline = time.monotonic() + 2 + while send_mock.call_count < 2 and time.monotonic() < deadline: + time.sleep(0.01) + self.registry.stop(result["id"]) + + self.assertGreaterEqual(send_mock.call_count, 2) + + def test_update_on_an_unknown_job_reports_not_found(self): + result = self.registry.update("does-not-exist", prompt="anything") + self.assertFalse(result["ok"]) + self.assertIn("No loop job", result["error"]) + + def test_update_rejects_an_empty_prompt(self): + with patch.object(ui_server, "_opencode_json_request", return_value={"id": "ses_abc"}), \ + patch.object(ui_server, "_opencode_send_and_collect", return_value=("ok", None)): + result = self.registry.start("watch training", 120, "/tmp", None) + self._wait_for_first_tick(result["id"]) + + update_result = self.registry.update(result["id"], prompt=" ") + + self.assertFalse(update_result["ok"]) + self.assertIn("prompt", update_result["error"]) + + def test_update_rejects_an_interval_below_the_minimum(self): + with patch.object(ui_server, "_opencode_json_request", return_value={"id": "ses_abc"}), \ + patch.object(ui_server, "_opencode_send_and_collect", return_value=("ok", None)): + result = self.registry.start("watch training", 120, "/tmp", None) + self._wait_for_first_tick(result["id"]) + + update_result = self.registry.update(result["id"], interval_seconds=10) + + self.assertFalse(update_result["ok"]) + self.assertIn("Minimum", update_result["error"]) + + +class _LoopServerTestCase(unittest.TestCase): + """Spins up a real serve_ui() on 127.0.0.1:, same shape as + test_server_agent.py's _ServerTestCase -- reused rather than imported + since that one is a module-private helper of its own file.""" + + def setUp(self): + self.tmp = tempfile.mkdtemp() + self.httpd = ui_server.serve_ui( + ui_host="127.0.0.1", ui_port=0, + backend_host="localhost", backend_port=50051, + open_browser=False, block=False, + experiment_dir=self.tmp, + ) + self.port = self.httpd.server_address[1] + self.thread = threading.Thread(target=self.httpd.serve_forever, daemon=True) + self.thread.start() + time.sleep(0.1) + + self._orig_registry = ui_server._loop_registry + ui_server._loop_registry = ui_server._LoopRegistry() + + def tearDown(self): + ui_server._loop_registry.shutdown() + ui_server._loop_registry = self._orig_registry + self.httpd.shutdown() + self.thread.join(timeout=5) + + def _get(self, path): + return urllib.request.urlopen(f"http://127.0.0.1:{self.port}{path}", timeout=5) + + def _post_json(self, path, body): + data = json.dumps(body).encode("utf-8") + req = urllib.request.Request( + f"http://127.0.0.1:{self.port}{path}", method="POST", data=data, + headers={"Content-Type": "application/json"}, + ) + return urllib.request.urlopen(req, timeout=10) + + def _post(self, path): + req = urllib.request.Request(f"http://127.0.0.1:{self.port}{path}", method="POST", data=b"") + return urllib.request.urlopen(req, timeout=10) + + +class TestLoopEndpoints(_LoopServerTestCase): + + def test_list_is_empty_before_anything_starts(self): + with self._get("/agent-server/loop/list") as r: + data = json.loads(r.read().decode()) + self.assertEqual(data["loops"], []) + + def test_start_delegates_to_the_registry_and_the_new_job_shows_up_in_list(self): + with patch.object(ui_server._loop_registry, "start", + return_value={"ok": True, "id": "1", "intervalSeconds": 1800.0}) as start_mock: + with self._post_json("/agent-server/loop/start", {"prompt": "watch the loss", "intervalMinutes": 30}) as r: + data = json.loads(r.read().decode()) + self.assertTrue(data["ok"], data) + self.assertEqual(data["id"], "1") + start_mock.assert_called_once() + args = start_mock.call_args.args + self.assertEqual(args[0], "watch the loss") + self.assertEqual(args[1], 30 * 60.0) + self.assertEqual(args[2], self.tmp) # rooted at the experiment dir + + def test_start_with_an_empty_prompt_returns_a_400(self): + import urllib.error + with self.assertRaises(urllib.error.HTTPError) as ctx: + self._post_json("/agent-server/loop/start", {"prompt": "", "intervalMinutes": 30}) + self.assertEqual(ctx.exception.code, 400) + data = json.loads(ctx.exception.read().decode()) + self.assertFalse(data["ok"]) + + def test_stop_delegates_to_the_registry_with_the_id_parsed_out_of_the_path(self): + with patch.object(ui_server._loop_registry, "stop", return_value={"ok": True}) as stop_mock: + with self._post("/agent-server/loop/42/stop") as r: + data = json.loads(r.read().decode()) + self.assertTrue(data["ok"]) + stop_mock.assert_called_once_with("42") + + def test_stopping_an_unknown_job_returns_a_404(self): + import urllib.error + with self.assertRaises(urllib.error.HTTPError) as ctx: + self._post("/agent-server/loop/does-not-exist/stop") + self.assertEqual(ctx.exception.code, 404) + + def test_update_delegates_to_the_registry_with_the_id_parsed_out_of_the_path(self): + with patch.object(ui_server._loop_registry, "update", + return_value={"ok": True, "id": "42", "prompt": "new prompt", "intervalSeconds": 300.0}) as update_mock: + with self._post_json("/agent-server/loop/42/update", {"prompt": "new prompt", "intervalMinutes": 5}) as r: + data = json.loads(r.read().decode()) + self.assertTrue(data["ok"], data) + update_mock.assert_called_once_with("42", prompt="new prompt", interval_seconds=5 * 60.0) + + def test_updating_an_unknown_job_returns_a_400(self): + import urllib.error + with self.assertRaises(urllib.error.HTTPError) as ctx: + self._post_json("/agent-server/loop/does-not-exist/update", {"prompt": "x"}) + self.assertEqual(ctx.exception.code, 400) + data = json.loads(ctx.exception.read().decode()) + self.assertFalse(data["ok"]) + + def test_list_reflects_a_real_start_end_to_end_through_the_registry(self): + with patch.object(ui_server._opencode_session, "ensure", + return_value={"ok": True, "url": "http://127.0.0.1:1", "workspace": self.tmp, "reused": False}), \ + patch.object(ui_server, "_opencode_json_request", return_value={"id": "ses_abc"}), \ + patch.object(ui_server, "_opencode_send_and_collect", return_value=("all quiet", None)): + with self._post_json("/agent-server/loop/start", {"prompt": "watch training", "intervalMinutes": 30}) as r: + started = json.loads(r.read().decode()) + self.assertTrue(started["ok"], started) + + deadline = time.monotonic() + 2 + job = None + while time.monotonic() < deadline: + with self._get("/agent-server/loop/list") as r: + loops = json.loads(r.read().decode())["loops"] + job = next((j for j in loops if j["id"] == started["id"]), None) + if job and job["lastResult"]: + break + time.sleep(0.02) + + self.assertIsNotNone(job) + self.assertEqual(job["lastResult"], "all quiet") + self.assertEqual(job["prompt"], "watch training") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/ui/test_server_shutdown_signals.py b/tests/ui/test_server_shutdown_signals.py new file mode 100644 index 00000000..348d410e --- /dev/null +++ b/tests/ui/test_server_shutdown_signals.py @@ -0,0 +1,151 @@ +"""Tests for weightslab/ui/server.py's termination-handler coverage. + +Problem: `weightslab start` only ever caught Ctrl+C (SIGINT, which Python +turns into a catchable KeyboardInterrupt by default). Closing the terminal +window or a bare `kill ` delivers a DIFFERENT signal/event +(SIGTERM/SIGHUP on POSIX, CTRL_CLOSE_EVENT on Windows) that Python does NOT +convert into a Python-level exception on its own -- so neither of those +ever ran this server's cleanup (_run_shutdown_cleanup: stopping tracked +detached processes, its own OpenCode/Jupyter children, /loop jobs), +confirmed live against real platform behavior. _install_termination_handlers +closes that gap. +""" + +import os +import signal +import threading +import time +import unittest +from unittest.mock import patch + +from weightslab.ui import server as ui_server + + +class TestRunShutdownCleanup(unittest.TestCase): + def test_calls_all_four_shutdown_methods(self): + with patch.object(ui_server._tracked_processes, "shutdown") as tp, \ + patch.object(ui_server._opencode_session, "shutdown") as oc, \ + patch.object(ui_server._loop_registry, "shutdown") as lr, \ + patch.object(ui_server._jupyter_session, "shutdown") as js: + ui_server._run_shutdown_cleanup() + tp.assert_called_once() + oc.assert_called_once() + lr.assert_called_once() + js.assert_called_once() + + +class TestRaiseKeyboardInterrupt(unittest.TestCase): + def test_raises_keyboardinterrupt(self): + # Exactly how Python's own signal machinery would invoke this: a + # (signum, frame) callback. Called directly, not via a real signal, + # so this passes identically on every platform. + with self.assertRaises(KeyboardInterrupt): + ui_server._raise_keyboard_interrupt(signal.SIGTERM, None) + + +class TestOnWindowsCtrlEvent(unittest.TestCase): + """Pure logic, no ctypes/real Windows API involved -- safe to run on + any platform, unlike _install_windows_console_handler itself below.""" + + def test_terminating_events_run_cleanup_and_report_handled(self): + for ctrl_type in (2, 5, 6): # CLOSE, LOGOFF, SHUTDOWN + with patch.object(ui_server, "_run_shutdown_cleanup") as cleanup_mock: + result = ui_server._on_windows_ctrl_event(ctrl_type) + cleanup_mock.assert_called_once() + self.assertTrue(result) + + def test_ctrl_c_and_ctrl_break_are_left_alone(self): + # Python's own signal module already turns these into SIGINT/ + # SIGBREAK -- this handler must not double-handle them. + for ctrl_type in (0, 1): + with patch.object(ui_server, "_run_shutdown_cleanup") as cleanup_mock: + result = ui_server._on_windows_ctrl_event(ctrl_type) + cleanup_mock.assert_not_called() + self.assertFalse(result) + + def test_unknown_event_is_left_alone(self): + with patch.object(ui_server, "_run_shutdown_cleanup") as cleanup_mock: + result = ui_server._on_windows_ctrl_event(99) + cleanup_mock.assert_not_called() + self.assertFalse(result) + + +@unittest.skipUnless(os.name == "nt", "SetConsoleCtrlHandler only exists on Windows") +class TestInstallWindowsConsoleHandler(unittest.TestCase): + def test_registers_a_real_handler_without_raising(self): + ui_server._install_windows_console_handler() + self.assertIsNotNone(ui_server._console_ctrl_handler_ref) + + +class TestInstallTerminationHandlers(unittest.TestCase): + def setUp(self): + self._orig_sigterm = signal.getsignal(signal.SIGTERM) + self._orig_sighup = signal.getsignal(signal.SIGHUP) if hasattr(signal, "SIGHUP") else None + + def tearDown(self): + signal.signal(signal.SIGTERM, self._orig_sigterm) + if hasattr(signal, "SIGHUP"): + signal.signal(signal.SIGHUP, self._orig_sighup) + + def test_registers_sigterm_to_raise_keyboardinterrupt(self): + ui_server._install_termination_handlers() + self.assertIs(signal.getsignal(signal.SIGTERM), ui_server._raise_keyboard_interrupt) + + @unittest.skipUnless(hasattr(signal, "SIGHUP"), "SIGHUP does not exist on this platform") + def test_registers_sighup_to_raise_keyboardinterrupt(self): + ui_server._install_termination_handlers() + self.assertIs(signal.getsignal(signal.SIGHUP), ui_server._raise_keyboard_interrupt) + + def test_installs_the_windows_console_handler_only_on_windows(self): + with patch.object(ui_server, "_install_windows_console_handler") as install_mock, \ + patch.object(ui_server.os, "name", "nt"): + ui_server._install_termination_handlers() + install_mock.assert_called_once() + + def test_skips_the_windows_console_handler_on_posix(self): + with patch.object(ui_server, "_install_windows_console_handler") as install_mock, \ + patch.object(ui_server.os, "name", "posix"): + ui_server._install_termination_handlers() + install_mock.assert_not_called() + + +@unittest.skipIf( + os.name == "nt", + "os.kill(pid, SIGTERM) maps to TerminateProcess on Windows (a hard kill " + "that bypasses any registered handler), so a self-SIGTERM isn't a safe " + "way to test this there -- the Windows-specific path is covered by " + "TestOnWindowsCtrlEvent/TestInstallWindowsConsoleHandler instead.", +) +class TestServeUiRealSigtermEndToEnd(unittest.TestCase): + """The real thing: an actual SIGTERM delivered to this process while + serve_ui(block=True) is blocking on the MAIN thread (signal handlers + only ever run on the main thread, so this only proves anything when + serve_forever() itself is there too -- exactly how `weightslab start` + really calls it).""" + + def test_sigterm_interrupts_serve_forever_and_runs_cleanup(self): + import tempfile + + tmp = tempfile.mkdtemp() + cleanup_called = threading.Event() + + def _send_sigterm_shortly(): + time.sleep(0.4) + os.kill(os.getpid(), signal.SIGTERM) + + with patch.object(ui_server, "_run_shutdown_cleanup", side_effect=cleanup_called.set): + sender = threading.Thread(target=_send_sigterm_shortly, daemon=True) + sender.start() + ui_server.serve_ui( + ui_host="127.0.0.1", ui_port=0, + backend_host="localhost", backend_port=50051, + open_browser=False, block=True, + experiment_dir=tmp, + ) + sender.join(timeout=5) + + self.assertTrue(cleanup_called.is_set()) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/ui/test_server_tracked_processes.py b/tests/ui/test_server_tracked_processes.py new file mode 100644 index 00000000..9c9149ff --- /dev/null +++ b/tests/ui/test_server_tracked_processes.py @@ -0,0 +1,159 @@ +"""Tests for weightslab/ui/server.py's POST /agent-server/track-process and +the _TrackedProcesses registry behind it. + +Problem this exists for: the agent is told to launch anything long-running +(training, a relaunched crashed run) DETACHED (Start-Process/setsid) so it +never blocks the chat turn. A detached process has no OS-level parent-child +relationship this server's own process-tree kill (_kill_process_tree) can +walk -- confirmed live: a detached launcher's own immediate shell exits +almost immediately after spawning it, and Windows keeps no record of an +exited process for `taskkill /T` to trace a grandchild through. Registering +the PID directly sidesteps that: this server kills it explicitly, by PID, +with no chain to walk at all. + +Uses a REAL child process (python -c "time.sleep(...)"), not a mock, so the +actual kill call is exercised end to end. +""" + +import json +import subprocess +import sys +import tempfile +import threading +import time +import unittest +import urllib.error +import urllib.request + +from weightslab.ui import server as ui_server + + +def _is_alive(pid: int) -> bool: + if ui_server.os.name == "nt": + out = subprocess.run( + ["tasklist", "/FI", f"PID eq {pid}"], + capture_output=True, text=True, + ).stdout + return str(pid) in out + try: + ui_server.os.kill(pid, 0) + return True + except OSError: + return False + + +class TestTrackedProcessesUnit(unittest.TestCase): + def setUp(self): + self.registry = ui_server._TrackedProcesses() + self.proc = subprocess.Popen( + [sys.executable, "-c", "import time; time.sleep(120)"], + start_new_session=True, + ) + + def tearDown(self): + if self.proc.poll() is None: + self.proc.kill() + self.proc.wait(timeout=5) + + def test_tracked_pid_is_killed_on_shutdown(self): + self.assertTrue(_is_alive(self.proc.pid)) + self.registry.track(self.proc.pid) + + self.registry.shutdown() + + deadline = time.monotonic() + 5 + while time.monotonic() < deadline and _is_alive(self.proc.pid): + time.sleep(0.1) + self.assertFalse(_is_alive(self.proc.pid)) + + def test_untracked_pid_is_left_alone(self): + self.registry.shutdown() # nothing tracked + self.assertTrue(_is_alive(self.proc.pid)) + + def test_shutdown_clears_the_registry_so_a_second_call_is_a_no_op(self): + self.registry.track(self.proc.pid) + self.registry.shutdown() + # A second shutdown() must not error just because the pid is already + # gone (e.g. serve_ui's explicit call racing its own atexit hook). + self.registry.shutdown() + + def test_killing_an_already_dead_pid_does_not_raise(self): + self.proc.kill() + self.proc.wait(timeout=5) + self.registry.track(self.proc.pid) + self.registry.shutdown() # must not raise + + +class _ServerTestCase(unittest.TestCase): + """Real serve_ui() on an ephemeral port -- same shape as + test_server_agent.py's own _ServerTestCase.""" + + def setUp(self): + self.tmp = tempfile.mkdtemp() + self.httpd = ui_server.serve_ui( + ui_host="127.0.0.1", ui_port=0, + backend_host="localhost", backend_port=50051, + open_browser=False, block=False, + experiment_dir=self.tmp, + ) + self.port = self.httpd.server_address[1] + self.thread = threading.Thread(target=self.httpd.serve_forever, daemon=True) + self.thread.start() + time.sleep(0.1) + + self._orig_tracked = ui_server._tracked_processes + ui_server._tracked_processes = ui_server._TrackedProcesses() + + def tearDown(self): + ui_server._tracked_processes = self._orig_tracked + self.httpd.shutdown() + self.thread.join(timeout=5) + + def _post_json(self, path, body): + data = json.dumps(body).encode("utf-8") + req = urllib.request.Request( + f"http://127.0.0.1:{self.port}{path}", method="POST", data=data, + headers={"Content-Type": "application/json"}, + ) + return urllib.request.urlopen(req, timeout=10) + + +class TestTrackProcessEndpoint(_ServerTestCase): + def test_registers_the_pid_end_to_end(self): + proc = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(120)"], start_new_session=True) + try: + with self._post_json("/agent-server/track-process", {"pid": proc.pid}) as r: + data = json.loads(r.read().decode()) + self.assertTrue(data["ok"]) + + ui_server._tracked_processes.shutdown() + + deadline = time.monotonic() + 5 + while time.monotonic() < deadline and _is_alive(proc.pid): + time.sleep(0.1) + self.assertFalse(_is_alive(proc.pid)) + finally: + if proc.poll() is None: + proc.kill() + proc.wait(timeout=5) + + def test_non_integer_pid_is_rejected(self): + try: + self._post_json("/agent-server/track-process", {"pid": "not-a-number"}) + self.fail("expected an HTTPError") + except urllib.error.HTTPError as exc: + self.assertEqual(exc.code, 400) + data = json.loads(exc.read().decode()) + self.assertFalse(data["ok"]) + self.assertIn("integer", data["error"]) + + def test_missing_pid_is_rejected(self): + try: + self._post_json("/agent-server/track-process", {}) + self.fail("expected an HTTPError") + except urllib.error.HTTPError as exc: + self.assertEqual(exc.code, 400) + + +if __name__ == "__main__": + unittest.main() diff --git a/weightslab/AGENTS.md b/weightslab/AGENTS.md new file mode 100644 index 00000000..57ed52cf --- /dev/null +++ b/weightslab/AGENTS.md @@ -0,0 +1,362 @@ +# WeightsLab — agent context for users & debugging + +Portable context for AI coding agents (and their humans) to **install, run, +integrate, and debug WeightsLab / Weights Studio** without reverse-engineering +the system first. Covers two repos: **weightslab** (Python backend — training +instrumentation, data ledger, gRPC service, shared proto) and **weights_studio** +(browser frontend that inspects/edits a *running* experiment). + +> File/line refs drift — verify against current source. Env var names/defaults +> are stable; authoritative reference is `weightslab/docs/configuration.rst`. + +--- + +## 0. Loading this guide into Claude Code + +- **Repo checkout:** committed as `AGENTS.md`; a gitignored `CLAUDE.md` copy at + the root gets auto-loaded every session. Nothing to do. +- **`pip install weightslab` only** (no checkout): absolute paths are fragile + across venvs/OS. Use a skill that locates the installed copy at runtime — + `~/.claude/skills/weightslab/SKILL.md`: + + ```yaml + --- + name: weightslab + description: Load the WeightsLab debugging & integration guide for weightslab/weights_studio problems (connection, TLS, env vars, training hangs, rendering, wl.* integration). + --- + !`python -c "import weightslab, os; print(open(os.path.join(os.path.dirname(weightslab.__file__), 'AGENTS.md')).read())"` + + Use the guide above to diagnose or implement the user's request. + ``` + + Requires the guide shipped as package data (`weightslab/weightslab/AGENTS.md` — see §7). +- **Quick-and-dirty:** copy this file to `~/.claude/WEIGHTSLAB.md`, `@`-import it + from `~/.claude/CLAUDE.md`. + +--- + +## 1. What it is, how the pieces connect + +A user wraps their PyTorch training script with WeightsLab so a running +experiment becomes inspectable/editable; Weights Studio is the UI. + +``` +Browser → weightslab start :8080 (grpc-web → grpc proxy) → Python gRPC servicer → training loop +``` + +- `weightslab start`: pure-Python HTTP server, serves the bundled SPA and + translates grpc-web↔gRPC. No Docker, no Envoy. Not running ⇒ no UI to load. +- gRPC servicer and training loop share **one process, different threads**, + coordinated by locks in `weightslab/weightslab/components/global_monitoring.py`. +- Proto is the single source of truth: `weightslab/weightslab/proto/experiment_service.proto`. + +--- + +## 2. Install & run + +```bash +pip install weightslab +``` + +```python +import weightslab as wl +# wrap objects so the studio can see/edit them (§3), then: +wl.serve(serving_grpc=True, serving_cli=True) +# ... training loop ... +wl.keep_serving() # keep process alive for the UI +``` + +```bash +weightslab start # http://localhost:8080 by default +``` + +For a new script, pick the closest match in +`weightslab/weightslab/examples/{PyTorch,Lightning,Ultralytics,Usecases}//main.py` +via the decision table in §3.9, then copy its `wl.*` calls — §3 documents that +whole API surface (reactive signals, group signals, the Ultralytics mixin, +etc. aren't in the `.rst` docs; the examples are the primary source). + +TLS/UI deploy details: `weightslab/docs/weights_studio.rst`. TLS is opt-in: +`weightslab se` once, then `weightslab start --certs`. + +--- + +## 3. The integration API (`import weightslab as wl`) + +How to wire a new training script correctly with no docs access — every verb +and kwarg here is real, taken from a shipping example under +`weightslab/weightslab/examples/` and checked against `weightslab/src.py`. + +### 3.1 Lifecycle + +```python +import weightslab as wl + +wl.watch_or_edit(..., flag=...) # register objects (§3.2) +wl.serve(serving_grpc=True, serving_cli=True) # background threads, same process +wl.start_training(timeout=3) # let UI/CLI attach before stepping +# ... training loop, guarded (§3.5) ... +wl.keep_serving() # block so the process/UI survives +``` + +- Register every object with `watch_or_edit` **before** `wl.serve`, using the parameter flag to define which object category it is. +- `timeout=0` skips the pre-start wait entirely. +- Skip `keep_serving()` for a script that should exit after writing a report + (`Usecases/*signals*` examples); include it otherwise. +- Tabular examples (`wl-fraud-detection`, `wl-ads-recommendation`) pass only + `serving_grpc=` to `serve` — no `serving_cli`. + +### 3.2 `wl.watch_or_edit(obj, flag=..., **kwargs)` + +Registers/wraps `obj` in the global ledger (`backend/ledgers.py`, +`GLOBAL_LEDGER`) and returns a live proxy. `flag` matches by substring +(case-insensitive). + +| flag | wraps | key kwargs | +|---|---|---| +| `"hyperparameters"` | plain `dict` | `defaults=parameters`, `poll_interval=1.0`, optional `name=` | +| `"model"` | `nn.Module` | `device=`; `compute_dependencies=False` (skip arch-op dependency graph when not editing architecture); `forced_model_wrapping=True` (Ultralytics only — load current object, not a checkpoint) | +| `"optimizer"` | `torch.optim.Optimizer` | none typically; build from the **watched** model's `.parameters()` | +| `"data"` | `Dataset` → tracked `DataLoader` | `loader_name=`, `batch_size=`, `shuffle=`, `is_training=`, `compute_hash=False`, `collate_fn=`, `preload_labels=`, `preload_metadata=` (both `True` for tabular — §3.4), `enable_h5_persistence=`, `num_workers=`; point-cloud/array data adds `array_autoload_arrays=`, `array_return_proxies=`, `array_use_cache=` | +| `"loss"` | `reduction="none"` criterion | `signal_name=`/`name=` (aliases), `log=True`, `per_sample=True` (one value/sample), `per_instance=True` (one value per `(sample_id, annotation_id)` — multi-box/mask samples); called as `criterion(preds_raw, targets, batch_ids=ids, preds=preds)` | +| `"metric"` | anything with `.compute()`/`.forward()` | same as `"loss"` | + +Registering `flag="loss"` auto-enrolls the signal for the background +loss-shape classifier (§3.6) unless overridden via `@wl.signal_classifier`. + +Objects need a `__name__` — set `obj.__name__ = "..."` manually if missing +(plain callables/custom loss modules). + +Hyperparameter proxy supports both `hp.get("lr")` and `hp["lr"]`, and stays +live (reflects later edits/re-registration). + +### 3.3 Per-sample / per-instance / grouped logging + +Watched loss/metric objects call `save_signals` internally on every +forward/compute — call these yourself only for derived values or anything not +from a watched object: + +- `wl.save_signals(batch_ids=ids, signals={...}, preds_raw=, targets=, preds=, log=True)` — + one value per sample id. `log=False` → stored as metadata, not a plotted signal. +- `wl.save_instance_signals(...)` — internal use by `per_instance=True`; rarely called directly. +- `wl.save_group_signals(signals={...}, group_ids=[...], origin="train_loader")` — + one row per group, for pairwise values (e.g. contrastive loss) that can't map + to a single sample. Needs a dataset that emits a `group_id` in its metadata + (`PyTorch/wl-generation`). +- `wl.trajectory_stats(values)` / `wl.classify_loss_shape(values)` — building + blocks behind the loss-shape tag (§3.6); call directly only for a custom classifier. + +### 3.4 `task_type` + +Set `self.task_type = "..."` on **both** model and dataset before +`watch_or_edit`. Confirmed values: + +| `task_type` | renders | set in | +|---|---|---| +| *(unset)* | classification (default) — also clustering, tabular, signal-tagging use cases | most examples | +| `"detection"` | 2D bounding boxes | `PyTorch/wl-detection/utils/{model,data}.py` | +| `"segmentation"` | instance/semantic masks | `PyTorch/wl-segmentation/utils/{model,data}.py` | +| `"detection_pointcloud"` | LiDAR point clouds, 2D or 3D (box column count disambiguates) | `Usecases/wl-{2d,3d}-lidar-detection/utils/{model,data}.py` | + +No `task_type="tabular"` exists — tabular rendering comes from the dataset +exposing feature values as sample **metadata** (`preload_labels=True, +preload_metadata=True`); mirror `PyTorch/wl-fraud-detection`, not the image +classification example. + +A dataset can implement `render_thumbnail_2d(...)` / `project_boxes_2d(...)` +for custom thumbnails — picked up automatically, no registration +(`Usecases/wl-3d-lidar-detection`, `CustomLidarDataset`). + +### 3.5 Guard contexts (required) + +```python +from weightslab import guard_training_context, guard_testing_context + +with guard_training_context: + ... # one training step +with guard_testing_context, torch.no_grad(): + ... # one eval step +``` + +Skip this and pause/resume and train/test stat separation break. Framework +variants: +- **Lightning:** wrap the body of `training_step`/`validation_step` — no manual loop. +- **Ultralytics mixin:** entered/exited manually (`guard.__enter__()`/`__exit__(None,None,None)`) + across `on_train_batch_start`/`_end` callback pairs (§3.9). + +Use `model.get_age()` (steps actually trained, survives checkpoints) for +step-based cadence, not a raw loop counter. + +### 3.6 Reactive signals & loss-shape classification + +```python +@wl.signal(name="sig/entropy", subscribe_to="loss_sample", batched=True) +def entropy(b): ... # b.logits, etc. → per-sample values + +@wl.signal(name="sig/hardness", inputs=["loss_sample", "sig/entropy"], batched=True) +def hardness(loss_vals, entropy_vals): ... +``` + +- `subscribe_to=` fires reactively when that signal saves (push); `inputs=[...]` + pulls named signals as args and can chain off other `@wl.signal` outputs. +- Define before `wl.serve()`/`wl.start_training()` — module scope or inline in `main()`. +- `@wl.signal_classifier(signal=)` overrides the built-in 6-way + loss-shape classifier (`monotonic/plateaued/Flat_high/high_variance/U_Shape/Spiked`) + for that signal, surfaced as categorical `tag:loss_shape` — no manual tagging + needed. Rebind at runtime: `wl.signal_classifier(signal=name)(fn)`. +- No custom classifier needed? Pass `loss_shape_signal=` to + `wl.write_dataframe(...)` (§3.7) to compute the built-in tag at dump time. + +### 3.7 Persisting & inspecting history + +- `wl.write_history()` / `wl.write_dataframe(path=, format="csv", columns=[...], loss_shape_signal=)` — + dump the ledger; `columns=` filters groups (e.g. `["signals","tags"]`). Call + periodically in long loops and once at the end. +- `wl.drain_signals()` — force-flush async signals before reading them back + (dataframe export, or a `GetDataSamples` call right after training). +- `wl.query_signal_history(...)` / `query_sample_history(...)` / `query_instance_history(...)` — + programmatic readback. + +### 3.8 Tagging & filtering samples + +`wl.tag_samples(...)`, `wl.register_categorical_tag(...)`/`set_categorical_tag(...)` +(multi-value, predefined categories — boolean tags are separate), +`wl.discard_samples(...)`, `wl.get_samples_by_tag(...)`, `wl.get_discarded_samples(...)`. +The automatic `tag:loss_shape` tag (§3.6) uses these same primitives. + +### 3.9 Which example to copy + +| Integrating... | Mirror | Notes | +|---|---|---| +| Plain PyTorch loop (classification) | `PyTorch/wl-classification` | Simplest pattern, manual loop in `main()`. | +| Detection (2D boxes) | `PyTorch/wl-detection` | `task_type="detection"`, `per_sample`/`per_instance`, custom `collate_fn`, decoded preds passed for overlays. | +| Segmentation | `PyTorch/wl-segmentation` | `task_type="segmentation"`, masks as list-of-tensors. | +| LiDAR detection (2D/3D) | `Usecases/wl-{2d,3d}-lidar-detection` | `task_type="detection_pointcloud"`; 3D adds `render_thumbnail_2d`. | +| Tabular / feature vectors | `PyTorch/wl-fraud-detection` | No `task_type`; `preload_labels=True, preload_metadata=True`; see §3.10 for a headless verification script. | +| Embedding / clustering | `PyTorch/wl-clustering` (+ `face/model.py`) | `watch_or_edit` calls live inside the model wrapper, not `main.py`; open-ended loop. | +| Paired/contrastive samples, group-level signals | `PyTorch/wl-generation` | `wl.save_group_signals`; dataset emits 2 rows per item via a `uids` metadata key. | +| Reactive signals / custom loss-shape tagging | `Usecases/wl-classification-signals_shape_classification`, `Usecases/ws-signals-mnist` | §3.6; the latter is the minimal variant with no custom classifier. | +| PyTorch Lightning | `Lightning/wl-classification` | Same `watch_or_edit` calls as plain PyTorch; guards wrap `training_step`/`validation_step` bodies; `Trainer(log_every_n_steps=0, enable_checkpointing=False, logger=False)`. | +| Ultralytics YOLO (detect/segment) | `Ultralytics/wl-detection` | Don't call `watch_or_edit` for model/optimizer/data/loss/metric — pass `trainer=WLAwareTrainer` (or `WLAwareSegmentationTrainer`) from `weightslab.integrations.ultralytics` to `YOLO(...).train(...)`. It wires everything via UL callbacks; you only watch the run config as `flag="hyperparameters"`. | + +### 3.10 Verifying an integration headlessly + +Watch model/optimizer/data/loss/metrics, `wl.serve(serving_grpc=True, grpc_port=...)`, +`wl.start_training()`, run real steps inside the guard contexts, +`wl.drain_signals()`, then assert on `wl.write_dataframe(..., format="csv")` +columns or a raw gRPC `GetDataSamples` call's `raw_data.type` (e.g. `"vector"` +for tabular). Plain script, not pytest — `python verify_integration.py` +(`PyTorch/wl-fraud-detection/verify_integration.py`). + +--- + +## 4. Configuration (environment variables) + +Authoritative reference: `weightslab/docs/configuration.rst`. High-signal ones: + +**Backend:** + +| Variable | Default | Why | +|---|---|---| +| `WEIGHTSLAB_LOG_LEVEL` | `INFO` | `DEBUG` for detail (`WATCHDOG` level sits between WARNING/ERROR). | +| `GRPC_BACKEND_HOST`/`PORT` | `0.0.0.0`/`50051` | Backend gRPC bind address. | +| `GRPC_TLS_ENABLED` | `0` | TLS on the gRPC socket; set with `weightslab start --certs`. | +| `GRPC_TLS_REQUIRE_CLIENT_AUTH` | `0` | mTLS; must match `--certs`. | +| `WEIGHTSLAB_CERTS_DIR` | `~/.weightslab-certs` | Cert lookup — single source of truth. | +| `GRPC_AUTH_TOKEN` | unset | Optional token auth on top of mTLS. | +| `GRPC_MAX_MESSAGE_BYTES` | `268435456` | Raise if large tensors/images fail to transfer. | +| `WEIGHTSLAB_DISABLE_WATCHDOGS` | `0` | Set `1` when breakpoint-debugging (§5). | +| `GRPC_WATCHDOG_STUCK_SECONDS` | `60` | Lock/RPC stuck threshold + lock-acquire timeout. | + +**Frontend — runtime `window.*` globals (injected at `weightslab start` time; restart+reload to apply):** + +| Variable | Default | Why | +|---|---|---| +| `WS_SERVER_HOST`/`PORT`/`PROTOCOL` | `localhost`/`8080`/`http` | How the browser reaches the server — #1 connection knob. | +| `WS_HISTOGRAM_MAX_BINS` | `512` | Metadata histogram bar cap. | +| `BB_THUMB_RENDER` | `10` | Max boxes per thumbnail, per overlay (GT/PRED independent). | +| `BB_MODAL_RENDER` | `100` | Max boxes per modal image, per overlay. | +| `ENABLE_PLOTS` | `1` | `0` removes plots board + Signals card. | +| `ENABLE_DATA_EXPLORATION` | `1` | `0` removes data grid + metadata panel. | +| `ENABLE_HYPERPARAMETERS_OPTIMIZATION` | `1` | `0` makes HP inputs read-only, stops HP poll. | +| `ENABLE_AGENT` | `1` | `0` removes agent chat bar. | +| `ENABLE_NOTEBOOK` | `1` | `0` removes the notebook (shared in-process kernel against the live experiment; persisted as `notebook.ipynb`; `>`-prefixed cells ask the agent for code). | + +`VITE_*` vars are build-time (need a frontend rebuild); `WS_*`/`BB_*`/`ENABLE_*` +are runtime (need only restart + reload). `ENABLE_*` default on; `0`/`false`/`no`/`off` disables. + +--- + +## 5. Troubleshooting + +**Sample grid empty / "failed to fetch" / gRPC errors.** Check in order: (1) +backend serving on `0.0.0.0:50051`; (2) `weightslab start` running, browser +reaches `:8080`; (3) TLS mismatch if using `--certs` — run `weightslab se` +first, export `WEIGHTSLAB_CERTS_DIR` (or drop TLS: omit `--certs`, `GRPC_TLS_ENABLED=0`). + +**Env var change not taking effect.** `VITE_*` → rebuild frontend. +`WS_*`/`BB_*`/`ENABLE_*` → restart `weightslab start` + reload tab. + +**Grid flashes empty on auto-refresh.** Refreshes now skip while a +`GetDataSamples` fetch is in flight (`isFetchInProgress()` in +`weights_studio/src/grid_data/gridDataManager.ts`) — confirm your build has this guard. + +**Detection overlays slow/cluttered.** Cap with `BB_THUMB_RENDER` / +`BB_MODAL_RENDER` (GT and PRED capped independently; render-only, no data dropped). + +**Training hangs; `RESOURCE_EXHAUSTED`; server "restarts".** A watchdog flags +locks/RPCs held past `GRPC_WATCHDOG_STUCK_SECONDS` (60s) and restarts the gRPC +server after repeated unhealthy polls. Debugging with breakpoints that +intentionally exceed this? Set `WEIGHTSLAB_DISABLE_WATCHDOGS=1`. +`RESOURCE_EXHAUSTED` = a handler couldn't get the lock in time — find what's holding it. + +**Pause/resume broken, or train/test stats mixed up.** Train/eval step isn't +wrapped in `guard_training_context`/`guard_testing_context` — see §3.5. + +**Large weights/images fail to transfer.** Raise `GRPC_MAX_MESSAGE_BYTES`. + +**Agent bar says unconfigured.** Backed by a local OpenCode server +(`OPENCODE_URL`, default `http://127.0.0.1:4096`), auto-started on first use. +`/init` from the UI (then `/model`, `/reset`). See `docs/agent.rst`, `docs/weights_studio.rst`. + +**Agent says it "cannot run code."** It has bash/read/write/edit tools rooted +at this workspace directory (the frontend sends no `tools` restriction) — a +model claiming otherwise is declining to call a tool it actually has, not +reporting a real limitation. Ask it directly and concretely: "use your bash +tool to run `python