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/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..4fd6e570 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,48 @@ 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.
+
+Both agent surfaces (see above) converge on the **same** OpenCode server via
+one shared environment variable:
-- **OpenRouter** — cloud-hosted models (recommended; interactive onboarding in
- the UI).
-- **Ollama** — local inference, available immediately at backend startup when
- configured in ``agent_config.yaml``.
+.. code-block:: bash
-You can initialize it three ways.
+ 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), and otherwise the free-tier
+``opencode/deepseek-v4-flash-free`` automatically — a provider's own
+reported default used to be tried in between, but that could itself be an
+arbitrary, non-text-reasoning model whenever any provider had credentials
+configured, so it no longer overrides this.
+
+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 +347,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 +375,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 +395,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 +720,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 a9c70387..3aa327a5 100644
--- a/docs/configuration.rst
+++ b/docs/configuration.rst
@@ -718,10 +718,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
@@ -730,19 +735,24 @@ 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 to
+ OpenCode's own last-picked model, and otherwise the free-tier
+ ``opencode/deepseek-v4-flash-free`` automatically.
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
@@ -779,7 +789,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
^^^^^^^^^^^^^^^^^^^
@@ -791,73 +801,34 @@ 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 to OpenCode's own last-picked model, and otherwise the
+ free-tier ``opencode/deepseek-v4-flash-free`` automatically.
-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
^^^^^^^^^^^^^^^^^^^^^^^
@@ -865,11 +836,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.
@@ -877,13 +847,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/experiment_reports.rst b/docs/experiment_reports.rst
index 85c25c10..83fdc472 100644
--- a/docs/experiment_reports.rst
+++ b/docs/experiment_reports.rst
@@ -30,6 +30,7 @@ they produce the same artifact:
wl.ai_report_generation(signals=["train_loss"]) # specific signals
wl.ai_report_generation(use_agent=False) # skip the LLM call
wl.ai_report_generation(output_path="reports/run.html") # choose the file
+ wl.ai_report_generation(distributions=["train_loss"]) # + a histogram section
It returns the path written.
@@ -42,6 +43,7 @@ they produce the same artifact:
report train_loss val_loss
report --output /tmp/run_42.html
report --no-agent
+ report --distributions train_loss,val_loss
The reply gives the path, the number of signals included, and whether the
written analysis made it in.
@@ -62,6 +64,56 @@ they produce the same artifact:
Generate an experiment report on train_loss and val_loss.
+ To add a value-distribution histogram for a specific column (see
+ `Distributions`_ below), including as a follow-up on a report you already
+ generated:
+
+ .. code-block:: text
+
+ Generate an experiment report and include a histogram of train_loss.
+ Add a distribution of val_loss to the report.
+
+ This always goes through the SAME single backend action — the agent must
+ never break "generate a report" into several separate analysis questions
+ and hand-write its own summary; that would skip the plots/styling below
+ entirely.
+
+Updating a report vs. generating a new one
+---------------------------------------------
+
+Every path above always writes a *fresh*, separately timestamped file by
+default. When you ask through chat, though, wording matters:
+
+.. code-block:: text
+
+ Generate a report. # always a NEW file
+ Update the report with a histogram of val_loss. # overwrites the last one
+ Add a histogram of val_loss to the report. # overwrites the last one
+ Also include the confidence signal in it. # overwrites the last one
+
+"Generate"/"create"/"how is this going" (no reference to one already made)
+always produces a new file. Wording that refers to an *existing* report
+("update", "add X to **the** report", "also include Y in **it**") overwrites
+the most recently generated report for this experiment instead — the agent's
+reply says which happened ("updated"/"generated ... experiment report") and
+still names the file. Asking to "update" when nothing has been generated yet
+isn't an error: it just creates the first one, same as a plain "generate"
+would.
+
+A follow-up "add" is intentionally cumulative — asking to add a histogram of
+``val_loss`` after already having one for ``train_loss`` keeps both in the
+updated file, not just the newest one, as long as the request stays in the
+same conversation. There's no server-side memory of a report's contents
+behind this — the agent reasons about what to keep from what you (and it)
+said earlier in the chat, so it works within one back-and-forth but doesn't
+persist across separate sessions.
+
+Python/CLI callers that want the same overwrite-in-place behavior can pass
+the previous run's own path back in as ``output_path``
+(:func:`ai_report_generation`) / ``--output`` (the CLI's ``report`` command)
+— they already have direct control over the file, so there's no separate
+"update" flag for them.
+
- **Weights Studio button**: the bar-chart icon immediately left of the
notebook button in the connected app's header. Left-click generates a
report (checking agent availability first — see below); right-click opens
@@ -107,6 +159,14 @@ What's in the report
history swung the most (``max - min``). Both are ranked *inside DuckDB*
(``LoggerQueue.top_k_samples_by_reduce``) and only the top few ever leave
the database — see `Why per-sample data doesn't blow up the report`_.
+- **Distributions** *(optional — only when asked for)* — a value-distribution
+ histogram plus n/mean/std/range for each column named via ``distributions``
+ (see `Generating a report`_ above). Unlike a Signals card, this reads the
+ *current* per-sample dataframe, not the aggregated training curve — so it
+ answers "how spread out is train_loss across samples right now", not "how
+ did it move over training". A name that doesn't resolve to a column, or
+ resolves to one with no numeric values, still gets a card saying so rather
+ than being silently dropped. Not present at all when nobody asked for one.
- **Loss-Shape Classification** — if per-sample loss-shape classification has
already been computed for this experiment (:doc:`logger`'s
``wl.write_loss_shapes`` / the background auto-tagger), a count of samples
@@ -116,6 +176,18 @@ What's in the report
- **Dataset** — total sample count, discard count/rate, per-split counts
(the ``origin`` column), and a breakdown of any ``tag:*`` columns present.
+Light / dark mode
+--------------------
+
+The report follows the browser's ``prefers-color-scheme`` automatically, and
+also has its own toggle button (top-right of the banner) for overriding that
+— the choice is remembered (via ``localStorage``, scoped to that report file)
+so reopening the same report keeps the theme you picked. Signal/distribution
+plots are rendered once by matplotlib on a fixed white canvas, so they sit in
+a small always-light thumbnail card in either theme — this keeps their own
+text and gridlines legible instead of rendering (and shipping) two copies of
+every plot.
+
Why per-sample data doesn't blow up the report
--------------------------------------------------
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 7dd9dd29..0eb7a009 100644
--- a/docs/user_commands.rst
+++ b/docs/user_commands.rst
@@ -532,9 +532,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 6f7936a5..ea344329 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",
@@ -172,6 +174,15 @@ 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 opencode_process's
+ # WORKSPACE_SEED_FILES) -- lives inside the package, not just at the
+ # repo root, so a `pip install weightslab` ships it too.
+ "AGENTS.md",
+ # Seeded next to AGENTS.md in every experiment workspace. OpenCode reads a
+ # project config from the directory it is started in, which is that same
+ # workspace; this one points `instructions` at AGENTS.md.
+ "opencode.json",
]
[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/test_report_entrypoints.py b/tests/test_report_entrypoints.py
index fc4a1c89..c0dd61b1 100644
--- a/tests/test_report_entrypoints.py
+++ b/tests/test_report_entrypoints.py
@@ -146,6 +146,19 @@ def test_no_logger_raises(self):
wl_src.ai_report_generation()
self.assertIn("no experiment logger", str(ctx.exception))
+ def test_distributions_param_adds_the_section(self):
+ import pandas as pd
+
+ df = pd.DataFrame({"sample_id": [1, 2, 3], "train_loss": [0.5, 0.9, 0.1]}).set_index("sample_id")
+ manager = MagicMock()
+ manager.get_combined_df.return_value = df
+
+ with mock.patch.object(wl_src, "get_dataframe", return_value=manager):
+ path = wl_src.ai_report_generation(distributions=["train_loss"], use_agent=False)
+
+ html = open(path, encoding="utf-8").read()
+ self.assertIn("Distributions", html)
+
def test_exported_on_the_package(self):
import weightslab as wl
@@ -187,6 +200,20 @@ def test_no_agent_flag_skips_the_analysis(self):
self.assertIn("no written analysis", result["message"])
self.agent.generate_report_narrative.assert_not_called()
+ def test_distributions_flag_adds_the_section(self):
+ import pandas as pd
+
+ df = pd.DataFrame({"sample_id": [1, 2, 3], "train_loss": [0.5, 0.9, 0.1]}).set_index("sample_id")
+ manager = MagicMock()
+ manager.get_combined_df.return_value = df
+
+ with mock.patch.object(wl_src, "get_dataframe", return_value=manager):
+ result = _handle_command("report --distributions train_loss")
+
+ self.assertTrue(result["ok"])
+ html = open(result["path"], encoding="utf-8").read()
+ self.assertIn("Distributions", html)
+
def test_output_flag_selects_the_file(self):
out = os.path.join(self.tmp.name, "cli_report.html")
result = _handle_command(f"report --output {out}")
@@ -219,17 +246,20 @@ class TestAgentActionSharesTheSamePath(unittest.TestCase):
the Python and CLI entry points — it is the same ``generate_report`` call,
only with the live service's own logger/dataframe/agent."""
- def _run_action(self, root_log_dir, logger_q, narrative_fn):
+ def _run_action(self, root_log_dir, logger_q, narrative_fn, df_manager=None,
+ distributions=None, update_existing=False):
from weightslab.trainer.services.data_service import DataService
service = types.SimpleNamespace(
_resolve_checkpoint_manager=lambda: None,
_root_log_dir=root_log_dir,
- _df_manager=None,
+ _df_manager=df_manager,
_agent=types.SimpleNamespace(generate_report_narrative=narrative_fn),
)
with mock.patch("weightslab.backend.ledgers.get_logger", return_value=logger_q):
- return DataService._agent_generate_experiment_report(service)
+ return DataService._agent_generate_experiment_report(
+ service, distributions=distributions, update_existing=update_existing,
+ )
def test_action_writes_the_report_and_names_it_in_the_reply(self):
logger_q = _lg_with("train_loss")
@@ -257,6 +287,43 @@ def test_action_reports_failures_as_text(self):
message = self._run_action(tmp, None, lambda _s: "unused")
self.assertIn("no experiment logger available", message)
+ def test_update_existing_overwrites_the_prior_report_and_says_updated(self):
+ logger_q = _lg_with("train_loss")
+ with tempfile.TemporaryDirectory() as tmp:
+ first_message = self._run_action(tmp, logger_q, lambda _s: "Healthy run.")
+ self.assertTrue(first_message.startswith("Action: generated experiment report"))
+ first_path = first_message.split(" at ", 1)[1]
+
+ second_message = self._run_action(
+ tmp, logger_q, lambda _s: "Still healthy.", update_existing=True,
+ )
+ self.assertTrue(second_message.startswith("Action: updated experiment report"))
+ second_path = second_message.split(" at ", 1)[1]
+ self.assertEqual(first_path, second_path)
+ self.assertIn("Still healthy.", open(second_path, encoding="utf-8").read())
+
+ def test_update_existing_without_a_prior_report_still_generates_one(self):
+ logger_q = _lg_with("train_loss")
+ with tempfile.TemporaryDirectory() as tmp:
+ message = self._run_action(tmp, logger_q, lambda _s: "Healthy run.", update_existing=True)
+ self.assertTrue(message.startswith("Action: generated experiment report"))
+
+ def test_distributions_param_is_forwarded_to_the_report(self):
+ import pandas as pd
+
+ logger_q = _lg_with("train_loss")
+ df = pd.DataFrame({"sample_id": [1, 2, 3], "train_loss": [0.5, 0.9, 0.1]}).set_index("sample_id")
+ df_manager = MagicMock()
+ df_manager.get_combined_df.return_value = df
+
+ with tempfile.TemporaryDirectory() as tmp:
+ message = self._run_action(
+ tmp, logger_q, lambda _s: "Healthy run.",
+ df_manager=df_manager, distributions=["train_loss"],
+ )
+ path = message.split(" at ", 1)[1]
+ self.assertIn("Distributions", open(path, encoding="utf-8").read())
+
if __name__ == "__main__":
unittest.main()
diff --git a/tests/test_reporting.py b/tests/test_reporting.py
index be1f2c8f..cbcc458d 100644
--- a/tests/test_reporting.py
+++ b/tests/test_reporting.py
@@ -11,6 +11,8 @@
- summarize_context_for_llm (plot-free, bounded LLM payload)
- generate_report (the shared collect -> narrate -> render path behind the
Studio button, wl.ai_report_generation and the CLI `report` command)
+- light/dark mode (theme toggle + CSS variables always present in the output)
+- distributions (opt-in histogram section: column resolution, stats, HTML)
"""
import base64
@@ -264,6 +266,128 @@ def test_no_narrative_shows_fallback_text(self):
html = open(path, encoding="utf-8").read()
self.assertIn("No narrative was generated", html)
+ def test_report_supports_light_and_dark_mode(self):
+ lg = _lg()
+ with tempfile.TemporaryDirectory() as tmp:
+ ctx = reporting.collect_report_context(tmp, lg, None)
+ path = reporting.render_report(ctx, reporting.default_report_path(tmp))
+ html = open(path, encoding="utf-8").read()
+ # A manual toggle button, persisted via localStorage...
+ self.assertIn('id="wl-theme-toggle"', html)
+ self.assertIn("wl-report-theme", html)
+ # ...and a system-preference fallback for readers who never touch it.
+ self.assertIn("prefers-color-scheme: dark", html)
+ self.assertIn("--wl-bg", html)
+
+
+# ---------------------------------------------------------------------------
+# Distributions (opt-in histogram section)
+# ---------------------------------------------------------------------------
+
+class TestResolveDistributionColumn(unittest.TestCase):
+
+ def test_exact_match(self):
+ df = pd.DataFrame({"train_loss": [1.0]})
+ self.assertEqual(reporting._resolve_distribution_column(df, "train_loss"), "train_loss")
+
+ def test_nested_suffix_match(self):
+ df = pd.DataFrame({"signals//train_loss": [1.0]})
+ self.assertEqual(
+ reporting._resolve_distribution_column(df, "train_loss"), "signals//train_loss",
+ )
+
+ def test_substring_fallback_match(self):
+ df = pd.DataFrame({"my_train_loss_v2": [1.0]})
+ self.assertEqual(
+ reporting._resolve_distribution_column(df, "train_loss"), "my_train_loss_v2",
+ )
+
+ def test_no_match_returns_none(self):
+ df = pd.DataFrame({"accuracy": [1.0]})
+ self.assertIsNone(reporting._resolve_distribution_column(df, "train_loss"))
+
+ def test_none_df_returns_none(self):
+ self.assertIsNone(reporting._resolve_distribution_column(None, "train_loss"))
+
+
+class TestComputeDistributionEntries(unittest.TestCase):
+
+ def test_empty_request_returns_empty_list(self):
+ df = pd.DataFrame({"train_loss": [1.0, 2.0]})
+ self.assertEqual(reporting.compute_distribution_entries(df, None), [])
+ self.assertEqual(reporting.compute_distribution_entries(df, []), [])
+
+ def test_none_df_returns_empty_list(self):
+ self.assertEqual(reporting.compute_distribution_entries(None, ["train_loss"]), [])
+
+ def test_unresolved_column_is_flagged_not_dropped(self):
+ df = pd.DataFrame({"accuracy": [1.0, 2.0]})
+ entries = reporting.compute_distribution_entries(df, ["train_loss"])
+ self.assertEqual(len(entries), 1)
+ self.assertFalse(entries[0]["resolved"])
+ self.assertEqual(entries[0]["name"], "train_loss")
+
+ def test_resolved_numeric_column_has_stats_and_no_plot_without_matplotlib(self):
+ df = pd.DataFrame({"train_loss": [1.0, 2.0, 3.0, 4.0]})
+ entries = reporting.compute_distribution_entries(df, ["train_loss"], plt=None)
+ self.assertEqual(len(entries), 1)
+ entry = entries[0]
+ self.assertTrue(entry["resolved"])
+ self.assertEqual(entry["n"], 4)
+ self.assertEqual(entry["mean"], 2.5)
+ self.assertIsNone(entry["plot_b64"])
+
+ def test_non_numeric_column_reports_zero_n(self):
+ df = pd.DataFrame({"origin": ["train", "val"]})
+ entries = reporting.compute_distribution_entries(df, ["origin"])
+ self.assertTrue(entries[0]["resolved"])
+ self.assertEqual(entries[0]["n"], 0)
+
+
+class TestDistributionsInReport(unittest.TestCase):
+
+ def test_omitted_when_not_requested(self):
+ lg = _lg()
+ df = pd.DataFrame({"sample_id": [1, 2], "train_loss": [0.5, 0.9]}).set_index("sample_id")
+ with tempfile.TemporaryDirectory() as tmp:
+ ctx = reporting.collect_report_context(tmp, lg, df)
+ self.assertEqual(ctx["distributions"], [])
+ path = reporting.render_report(ctx, reporting.default_report_path(tmp))
+ html = open(path, encoding="utf-8").read()
+ self.assertNotIn("Distributions", html)
+
+ def test_included_when_requested(self):
+ lg = _lg()
+ df = pd.DataFrame({"sample_id": [1, 2, 3], "train_loss": [0.5, 0.9, 0.1]}).set_index("sample_id")
+ with tempfile.TemporaryDirectory() as tmp:
+ ctx = reporting.collect_report_context(tmp, lg, df, distributions=["train_loss"])
+ self.assertEqual(len(ctx["distributions"]), 1)
+ path = reporting.render_report(ctx, reporting.default_report_path(tmp))
+ html = open(path, encoding="utf-8").read()
+ self.assertIn("Distributions", html)
+ self.assertIn("train_loss", html)
+
+ def test_generate_report_forwards_distributions(self):
+ lg = _lg()
+ _seed_signal(lg, "train_loss", [1.0, 0.8, 0.6])
+ df = pd.DataFrame({"sample_id": [1, 2], "train_loss": [0.5, 0.9]}).set_index("sample_id")
+ with tempfile.TemporaryDirectory() as tmp:
+ result = reporting.generate_report(tmp, lg, df, distributions=["train_loss"])
+ html = open(result["path"], encoding="utf-8").read()
+ self.assertIn("Distributions", html)
+
+ def test_summarize_context_for_llm_drops_distribution_plots(self):
+ context = {
+ "signals": [],
+ "distributions": [{"name": "train_loss", "n": 3, "mean": 0.5, "plot_b64": "AAAA" * 500}],
+ "loss_shape_tags": [],
+ "dataframe": {},
+ }
+ payload = reporting.summarize_context_for_llm(context)
+ self.assertNotIn("plot_b64", payload)
+ self.assertNotIn("AAAA", payload)
+ self.assertIn("train_loss", payload)
+
# ---------------------------------------------------------------------------
# summarize_context_for_llm / generate_report
@@ -366,5 +490,98 @@ def test_unreadable_output_path_raises_rather_than_reporting_success(self):
reporting.generate_report(tmp, lg, None, output_path=clash)
+# ---------------------------------------------------------------------------
+# list_reports / latest_report_path
+# ---------------------------------------------------------------------------
+
+class TestListReports(unittest.TestCase):
+
+ def test_no_reports_dir_returns_empty(self):
+ with tempfile.TemporaryDirectory() as tmp:
+ self.assertEqual(reporting.list_reports(tmp), [])
+ self.assertIsNone(reporting.latest_report_path(tmp))
+
+ def test_newest_first_by_mtime(self):
+ with tempfile.TemporaryDirectory() as tmp:
+ reports_dir = os.path.join(tmp, "reports")
+ os.makedirs(reports_dir)
+ older = os.path.join(reports_dir, "experiment_report_20260101_000000.html")
+ newer = os.path.join(reports_dir, "experiment_report_20260102_000000.html")
+ open(older, "w").close()
+ open(newer, "w").close()
+ os.utime(older, (1_000_000, 1_000_000))
+ os.utime(newer, (2_000_000, 2_000_000))
+
+ result = reporting.list_reports(tmp)
+ self.assertEqual([p.name for p in result], [os.path.basename(newer), os.path.basename(older)])
+ self.assertEqual(reporting.latest_report_path(tmp).name, os.path.basename(newer))
+
+ def test_non_html_files_are_ignored(self):
+ with tempfile.TemporaryDirectory() as tmp:
+ reports_dir = os.path.join(tmp, "reports")
+ os.makedirs(reports_dir)
+ open(os.path.join(reports_dir, "notes.txt"), "w").close()
+ self.assertEqual(reporting.list_reports(tmp), [])
+
+
+# ---------------------------------------------------------------------------
+# generate_report(update_existing=...)
+# ---------------------------------------------------------------------------
+
+class TestGenerateReportUpdateExisting(unittest.TestCase):
+
+ def _seeded_logger(self):
+ lg = _lg()
+ _seed_signal(lg, "train_loss", [1.0, 0.8, 0.6, 0.4])
+ return lg
+
+ def test_no_prior_report_falls_back_to_creating_one(self):
+ lg = self._seeded_logger()
+ with tempfile.TemporaryDirectory() as tmp:
+ result = reporting.generate_report(tmp, lg, None, update_existing=True)
+ self.assertFalse(result["updated_existing"])
+ self.assertTrue(os.path.isfile(result["path"]))
+
+ def test_overwrites_the_latest_report_in_place(self):
+ lg = self._seeded_logger()
+ df = pd.DataFrame({"sample_id": [1, 2], "train_loss": [0.5, 0.9]}).set_index("sample_id")
+ with tempfile.TemporaryDirectory() as tmp:
+ existing = os.path.join(tmp, "reports", "experiment_report_20260101_000000.html")
+ reporting.generate_report(tmp, lg, None, output_path=existing)
+
+ result = reporting.generate_report(tmp, lg, df, update_existing=True,
+ distributions=["train_loss"])
+
+ self.assertTrue(result["updated_existing"])
+ self.assertEqual(os.path.abspath(result["path"]), os.path.abspath(existing))
+ self.assertEqual(len(reporting.list_reports(tmp)), 1)
+ self.assertIn("Distributions", open(existing, encoding="utf-8").read())
+
+ def test_default_call_always_creates_a_separate_file(self):
+ lg = self._seeded_logger()
+ with tempfile.TemporaryDirectory() as tmp:
+ path_a = os.path.join(tmp, "reports", "a.html")
+ path_b = os.path.join(tmp, "reports", "b.html")
+ with mock.patch.object(reporting, "default_report_path", side_effect=[path_a, path_b]):
+ first = reporting.generate_report(tmp, lg, None)
+ second = reporting.generate_report(tmp, lg, None)
+
+ self.assertFalse(first["updated_existing"])
+ self.assertFalse(second["updated_existing"])
+ self.assertNotEqual(first["path"], second["path"])
+ self.assertEqual(len(reporting.list_reports(tmp)), 2)
+
+ def test_explicit_output_path_overrides_update_existing(self):
+ lg = self._seeded_logger()
+ with tempfile.TemporaryDirectory() as tmp:
+ reporting.generate_report(tmp, lg, None) # an existing report to (not) update
+ chosen = os.path.join(tmp, "explicit.html")
+
+ result = reporting.generate_report(tmp, lg, None, update_existing=True, output_path=chosen)
+
+ self.assertFalse(result["updated_existing"])
+ self.assertEqual(os.path.abspath(result["path"]), os.path.abspath(chosen))
+
+
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..81468e21
--- /dev/null
+++ b/tests/trainer/services/test_opencode_chat.py
@@ -0,0 +1,448 @@
+"""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_the_hardcoded_default_when_config_has_no_model(self):
+ # Confirmed live: a provider-reported default (/config/providers, no
+ # longer consulted at all -- see _ensure_model_resolved's own
+ # docstring) could itself be an arbitrary, non-text-reasoning model
+ # whenever ANY provider had credentials configured, silently
+ # pre-empting this default every time. A single /config miss now
+ # goes straight to the known-good free model, with no second request.
+ 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({})) as request_mock:
+ chat._ensure_model_resolved()
+ request_mock.assert_called_once_with("/config")
+ self.assertEqual(chat.model, "opencode/deepseek-v4-flash-free")
+
+ 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_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 -- falls to the hardcoded
+ # default, not a second request.
+ 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": "not-a-provider-model-pair"})) as request_mock:
+ chat._ensure_model_resolved()
+ request_mock.assert_called_once_with("/config")
+ self.assertEqual(chat.model, "opencode/deepseek-v4-flash-free")
+
+
+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 28d5ac0f..9cb8f17b 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/fake_opencode.py b/tests/ui/fake_opencode.py
new file mode 100644
index 00000000..35832891
--- /dev/null
+++ b/tests/ui/fake_opencode.py
@@ -0,0 +1,261 @@
+"""A fake OpenCode server, over real HTTP, for the backend agent tests.
+
+WHY. ``weightslab/ui/server.py`` talks to OpenCode with three stdlib helpers:
+``_opencode_json_request`` (plain JSON round trip), ``_opencode_get_messages``
+(a passthrough), and ``_opencode_send_and_collect`` -- which is where all the
+real complexity lives. That one opens the SSE event stream FIRST, sends the
+prompt from a second thread, then reads the stream until this session goes
+idle, assembling reply text per part id and distinguishing "the turn ended
+having said nothing" from "the turn failed".
+
+Every existing test patches those helpers out (see test_server_loop.py's own
+header), which is right for exercising ``_LoopRegistry``'s scheduling logic but
+means the wire layer itself -- stream-before-send ordering, SSE framing,
+keep-alive comments, which events count and which are ignored, error vs idle
+termination -- has never been executed against anything. This serves the real
+protocol on a real socket so it is.
+
+Deliberately stdlib-only, matching the module it tests (``ui/server.py`` is
+stdlib-only by design so ``weightslab start`` needs no extra dependency).
+
+USAGE::
+
+ with FakeOpencode() as server:
+ server.script = [
+ {"type": "message.updated",
+ "properties": {"info": {"id": "m1", "role": "assistant",
+ "sessionID": "ses_1"}}},
+ {"type": "message.part.updated",
+ "properties": {"part": {"id": "p1", "messageID": "m1",
+ "sessionID": "ses_1",
+ "type": "text", "text": "hello"}}},
+ {"type": "session.idle", "properties": {"sessionID": "ses_1"}},
+ ]
+ text, error = ui_server._opencode_send_and_collect(
+ server.base_url, "ses_1", "check in")
+
+The scripted events are emitted when the prompt POST arrives, which is what a
+real server does: the turn only starts once the message is sent.
+"""
+
+import json
+import queue
+import threading
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+
+
+class FakeOpencode:
+ """A scriptable OpenCode stand-in on 127.0.0.1, speaking real HTTP/SSE."""
+
+ def __init__(self):
+ #: Events emitted (in order) once the prompt POST arrives.
+ self.script: list = []
+ #: HTTP status the prompt POST answers with.
+ self.post_status: int = 200
+ #: Seconds the prompt POST blocks before answering, to exercise the
+ #: "the send itself is what times out" path.
+ self.post_delay: float = 0.0
+ #: Emit a `:` keep-alive comment ahead of the scripted events -- real
+ #: servers do, and the reader must skip them rather than treat one as
+ #: a truncated event.
+ self.send_keepalive: bool = True
+ #: Replies for the plain JSON routes, by "METHOD /path".
+ self.responses: dict = {}
+ #: What GET /session//message returns.
+ self.messages: list = []
+
+ #: Every request received, as (method, path, parsed-body-or-None).
+ self.requests: list = []
+ #: Bodies POSTed to /session//message specifically.
+ self.prompt_bodies: list = []
+
+ self._events: "queue.Queue" = queue.Queue()
+ self._server = ThreadingHTTPServer(("127.0.0.1", 0), _make_handler(self))
+ self._thread = threading.Thread(target=self._server.serve_forever, daemon=True)
+
+ # -- lifecycle ---------------------------------------------------------
+ @property
+ def base_url(self) -> str:
+ host, port = self._server.server_address[:2]
+ return f"http://{host}:{port}"
+
+ def start(self) -> "FakeOpencode":
+ self._thread.start()
+ return self
+
+ def stop(self) -> None:
+ # Unblock any streaming handler still parked on the queue, so
+ # shutdown() isn't waiting on a request that never finishes.
+ self._events.put(None)
+ self._server.shutdown()
+ self._server.server_close()
+
+ def __enter__(self) -> "FakeOpencode":
+ return self.start()
+
+ def __exit__(self, *_exc) -> None:
+ self.stop()
+
+ # -- scripting ---------------------------------------------------------
+ def emit(self, event: dict) -> None:
+ """Push one event onto the live stream, outside the scripted batch."""
+ self._events.put(event)
+
+ def end_stream(self) -> None:
+ """Close the event stream from the server side (EOF), as a restarting
+ or crashing server would."""
+ self._events.put(None)
+
+ def _release_script(self) -> None:
+ for event in self.script:
+ self._events.put(event)
+
+
+def _make_handler(state: FakeOpencode):
+ class Handler(BaseHTTPRequestHandler):
+ # HTTP/1.0: the response body is delimited by connection close, so a
+ # streamed SSE body needs no chunked framing and the client's
+ # line-by-line read still sees each event as it is written.
+ protocol_version = "HTTP/1.0"
+
+ def log_message(self, *_args) -> None:
+ pass # keep the test output clean
+
+ # -- helpers --
+ def _read_body(self):
+ length = int(self.headers.get("Content-Length") or 0)
+ if not length:
+ return None
+ raw = self.rfile.read(length)
+ try:
+ return json.loads(raw.decode("utf-8"))
+ except (ValueError, UnicodeDecodeError):
+ return raw
+
+ def _send_json(self, payload, status: int = 200) -> None:
+ body = json.dumps(payload).encode("utf-8")
+ self.send_response(status)
+ self.send_header("Content-Type", "application/json")
+ self.send_header("Content-Length", str(len(body)))
+ self.end_headers()
+ self.wfile.write(body)
+
+ def _stream_events(self) -> None:
+ """GET /event -- hold the connection open and write SSE frames as
+ they are queued, until a None sentinel or the client hangs up."""
+ self.send_response(200)
+ self.send_header("Content-Type", "text/event-stream")
+ self.send_header("Cache-Control", "no-cache")
+ self.end_headers()
+ try:
+ if state.send_keepalive:
+ self.wfile.write(b": keep-alive\n\n")
+ self.wfile.flush()
+ while True:
+ event = state._events.get()
+ if event is None:
+ return
+ frame = f"data: {json.dumps(event)}\n\n".encode("utf-8")
+ self.wfile.write(frame)
+ self.wfile.flush()
+ except (BrokenPipeError, ConnectionResetError, OSError):
+ # The reader breaks out on session.idle/session.error and
+ # closes its end -- an expected, ordinary hang-up here.
+ return
+
+ # -- verbs --
+ def do_GET(self) -> None:
+ state.requests.append(("GET", self.path, None))
+ if self.path == "/event":
+ self._stream_events()
+ return
+ if self.path.startswith("/session/") and self.path.endswith("/message"):
+ self._send_json(state.messages)
+ return
+ key = f"GET {self.path}"
+ if key in state.responses:
+ self._send_json(state.responses[key])
+ return
+ self._send_json({})
+
+ def do_POST(self) -> None:
+ body = self._read_body()
+ state.requests.append(("POST", self.path, body))
+
+ if self.path.startswith("/session/") and self.path.endswith("/message"):
+ state.prompt_bodies.append(body)
+ # The turn starts on send -- release the scripted events now,
+ # which is also what makes the stream-first ordering in
+ # _opencode_send_and_collect observable: had it sent first and
+ # subscribed after, it would miss everything below.
+ state._release_script()
+ if state.post_delay:
+ import time
+ time.sleep(state.post_delay)
+ if state.post_status != 200:
+ self._send_json({"error": "prompt rejected"}, status=state.post_status)
+ return
+ self._send_json({"ok": True})
+ return
+
+ key = f"POST {self.path}"
+ if key in state.responses:
+ self._send_json(state.responses[key])
+ return
+ self._send_json({})
+
+ def do_DELETE(self) -> None:
+ state.requests.append(("DELETE", self.path, None))
+ self._send_json({})
+
+ return Handler
+
+
+# --------------------------------------------------------------------------
+# Event builders -- the shapes ui/server.py's reader actually looks for, in one
+# place so a protocol change lands here rather than across every test.
+# --------------------------------------------------------------------------
+
+def assistant_message(message_id: str, session_id: str) -> dict:
+ return {
+ "type": "message.updated",
+ "properties": {"info": {"id": message_id, "role": "assistant", "sessionID": session_id}},
+ }
+
+
+def user_message(message_id: str, session_id: str) -> dict:
+ return {
+ "type": "message.updated",
+ "properties": {"info": {"id": message_id, "role": "user", "sessionID": session_id}},
+ }
+
+
+def text_part(part_id: str, message_id: str, session_id: str, text: str) -> dict:
+ return {
+ "type": "message.part.updated",
+ "properties": {"part": {
+ "id": part_id, "messageID": message_id, "sessionID": session_id,
+ "type": "text", "text": text,
+ }},
+ }
+
+
+def tool_part(part_id: str, message_id: str, session_id: str, tool: str = "bash") -> dict:
+ return {
+ "type": "message.part.updated",
+ "properties": {"part": {
+ "id": part_id, "messageID": message_id, "sessionID": session_id,
+ "type": "tool", "tool": tool, "state": {"status": "completed"},
+ }},
+ }
+
+
+def session_idle(session_id: str) -> dict:
+ return {"type": "session.idle", "properties": {"sessionID": session_id}}
+
+
+def session_error(session_id: str, name: str, message: str = None) -> dict:
+ error = {"name": name}
+ if message is not None:
+ error["data"] = {"message": message}
+ return {"type": "session.error", "properties": {"sessionID": session_id, "error": error}}
diff --git a/tests/ui/test_server_agent.py b/tests/ui/test_server_agent.py
new file mode 100644
index 00000000..0cf419ff
--- /dev/null
+++ b/tests/ui/test_server_agent.py
@@ -0,0 +1,538 @@
+"""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_seed_files_exist(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.opencode_process, "_packaged_file", 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)
+ for filename in ui_server.opencode_process.WORKSPACE_SEED_FILES:
+ self.assertFalse(os.path.isfile(os.path.join(self.tmp, filename)))
+
+ 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 TestEnsureWorkspaceAgentFiles(unittest.TestCase):
+ """ensure_workspace_agent_files directly -- no OpenCode process involved.
+
+ Both seeded files matter for the same reason: OpenCode reads a project
+ AGENTS.md and a project opencode.json out of the directory it is started
+ in, which is this workspace. A workspace missing either one silently
+ downgrades what the agent knows.
+ """
+
+ def setUp(self):
+ self.tmp = tempfile.mkdtemp()
+
+ def test_copies_both_seed_files_from_the_installed_package_location(self):
+ ui_server._ensure_workspace_agent_files(self.tmp)
+ for filename in ui_server.opencode_process.WORKSPACE_SEED_FILES:
+ target = os.path.join(self.tmp, filename)
+ self.assertTrue(os.path.isfile(target), f"{filename} was not seeded")
+ with open(target, encoding="utf-8") as fh:
+ self.assertEqual(fh.read(), ui_server._read_repo_doc(filename))
+
+ def test_the_seeded_config_points_opencode_at_agents_md(self):
+ """The whole reason opencode.json is seeded at all -- without the
+ `instructions` entry it would be an empty config doing nothing."""
+ ui_server._ensure_workspace_agent_files(self.tmp)
+ with open(os.path.join(self.tmp, "opencode.json"), encoding="utf-8") as fh:
+ config = json.load(fh)
+ self.assertIn("AGENTS.md", config["instructions"])
+ self.assertEqual(config["$schema"], "https://opencode.ai/config.json")
+
+ def test_is_a_no_op_when_the_workspace_already_has_one(self):
+ for filename in ui_server.opencode_process.WORKSPACE_SEED_FILES:
+ with open(os.path.join(self.tmp, filename), "w", encoding="utf-8") as fh:
+ fh.write("mine")
+ ui_server._ensure_workspace_agent_files(self.tmp)
+ for filename in ui_server.opencode_process.WORKSPACE_SEED_FILES:
+ with open(os.path.join(self.tmp, filename), encoding="utf-8") as fh:
+ self.assertEqual(fh.read(), "mine", f"{filename} was overwritten")
+
+ def test_seeds_the_other_file_when_one_is_already_present(self):
+ with open(os.path.join(self.tmp, "AGENTS.md"), "w", encoding="utf-8") as fh:
+ fh.write("mine")
+ ui_server._ensure_workspace_agent_files(self.tmp)
+ self.assertTrue(os.path.isfile(os.path.join(self.tmp, "opencode.json")))
+
+ def test_does_nothing_when_no_source_is_found(self):
+ with patch.object(ui_server.opencode_process, "_packaged_file", return_value=None):
+ ui_server._ensure_workspace_agent_files(self.tmp)
+ for filename in ui_server.opencode_process.WORKSPACE_SEED_FILES:
+ self.assertFalse(os.path.isfile(os.path.join(self.tmp, filename)))
+
+
+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_opencode_wire.py b/tests/ui/test_server_opencode_wire.py
new file mode 100644
index 00000000..060a1db1
--- /dev/null
+++ b/tests/ui/test_server_opencode_wire.py
@@ -0,0 +1,361 @@
+"""The OpenCode wire layer in ui/server.py, against a real fake server.
+
+Every other loop/agent test in this directory patches
+``_opencode_json_request``/``_opencode_send_and_collect`` out -- correct for
+testing ``_LoopRegistry``'s scheduling, but it means the protocol code itself
+(SSE framing, stream-before-send ordering, which events count, error vs idle
+termination, text assembly) has never actually run under test. This file runs
+it, over real HTTP, against ``fake_opencode.FakeOpencode``.
+
+No network, no credentials, no model calls: the fake binds 127.0.0.1:0 and
+every reply is scripted, so these are deterministic and fast.
+"""
+
+import threading
+import time
+import unittest
+from unittest import mock
+
+from weightslab.ui import server as ui_server
+
+from tests.ui.fake_opencode import (
+ FakeOpencode,
+ assistant_message,
+ session_error,
+ session_idle,
+ text_part,
+ tool_part,
+ user_message,
+)
+
+SES = "ses_1"
+
+
+class _FakeServerTestCase(unittest.TestCase):
+ def setUp(self):
+ self.server = FakeOpencode().start()
+ self.addCleanup(self.server.stop)
+
+ def collect(self, timeout=10.0, session=SES, text="check in", model=None):
+ return ui_server._opencode_send_and_collect(
+ self.server.base_url, session, text, model=model, timeout=timeout,
+ )
+
+
+class TestPlainJsonRoutes(_FakeServerTestCase):
+ def test_get_round_trip_parses_json(self):
+ self.server.responses["GET /config"] = {"model": "openrouter/anthropic/claude-opus-4.6"}
+ result = ui_server._opencode_json_request(self.server.base_url, "/config")
+ self.assertEqual(result["model"], "openrouter/anthropic/claude-opus-4.6")
+
+ def test_post_sends_a_json_body_and_content_type(self):
+ ui_server._opencode_json_request(
+ self.server.base_url, "/session", method="POST", body={"title": "loop"},
+ )
+ method, path, body = self.server.requests[-1]
+ self.assertEqual((method, path), ("POST", "/session"))
+ self.assertEqual(body, {"title": "loop"})
+
+ def test_a_non_2xx_status_raises_rather_than_returning_a_falsy_value(self):
+ # urllib raises HTTPError for 4xx/5xx; callers rely on that (the loop's
+ # own start() catches it to report "could not reach OpenCode").
+ self.server.post_status = 500
+ with self.assertRaises(Exception):
+ ui_server._opencode_json_request(
+ self.server.base_url, f"/session/{SES}/message", method="POST", body={"parts": []},
+ )
+
+ def test_get_messages_passes_the_list_straight_through(self):
+ self.server.messages = [{"info": {"id": "m1", "role": "assistant"}, "parts": []}]
+ result = ui_server._opencode_get_messages(self.server.base_url, SES)
+ self.assertEqual(result, self.server.messages)
+
+
+class TestSendAndCollectHappyPath(_FakeServerTestCase):
+ def test_assembles_the_reply_text_and_reports_no_error(self):
+ self.server.script = [
+ assistant_message("m1", SES),
+ text_part("p1", "m1", SES, "The run looks healthy."),
+ session_idle(SES),
+ ]
+ text, error = self.collect()
+ self.assertEqual(text, "The run looks healthy.")
+ self.assertIsNone(error)
+
+ def test_the_stream_is_opened_BEFORE_the_prompt_is_sent(self):
+ # The fake only releases its scripted events when the POST arrives, so
+ # collecting any text at all proves the subscription was already live.
+ # Sending first and subscribing after would lose the whole turn.
+ self.server.script = [
+ assistant_message("m1", SES),
+ text_part("p1", "m1", SES, "caught it"),
+ session_idle(SES),
+ ]
+ text, _ = self.collect()
+ self.assertEqual(text, "caught it")
+ self.assertEqual([r[1] for r in self.server.requests if r[0] == "GET"], ["/event"])
+
+ def test_a_later_delta_for_the_same_part_replaces_it_instead_of_appending(self):
+ self.server.script = [
+ assistant_message("m1", SES),
+ text_part("p1", "m1", SES, "Loss is "),
+ text_part("p1", "m1", SES, "Loss is 0.31 and falling."),
+ session_idle(SES),
+ ]
+ text, error = self.collect()
+ self.assertEqual(text, "Loss is 0.31 and falling.")
+ self.assertIsNone(error)
+
+ def test_several_parts_are_joined_in_arrival_order(self):
+ self.server.script = [
+ assistant_message("m1", SES),
+ text_part("p1", "m1", SES, "First. "),
+ text_part("p2", "m1", SES, "Second. "),
+ text_part("p3", "m1", SES, "Third."),
+ session_idle(SES),
+ ]
+ text, _ = self.collect()
+ self.assertEqual(text, "First. Second. Third.")
+
+ def test_keep_alive_comments_are_skipped_not_parsed(self):
+ self.server.send_keepalive = True
+ self.server.script = [
+ assistant_message("m1", SES),
+ text_part("p1", "m1", SES, "fine"),
+ session_idle(SES),
+ ]
+ text, error = self.collect()
+ self.assertEqual(text, "fine")
+ self.assertIsNone(error)
+
+ def test_the_selected_model_is_forwarded_on_the_prompt(self):
+ self.server.script = [session_idle(SES)]
+ model = {"providerID": "openrouter", "modelID": "anthropic/claude-opus-4.6"}
+ self.collect(model=model)
+ self.assertEqual(self.server.prompt_bodies[0]["model"], model)
+ self.assertEqual(self.server.prompt_bodies[0]["parts"][0]["text"], "check in")
+
+ def test_no_model_key_is_sent_when_none_was_resolved(self):
+ self.server.script = [session_idle(SES)]
+ self.collect(model=None)
+ self.assertNotIn("model", self.server.prompt_bodies[0])
+
+ def test_a_clean_turn_that_said_nothing_returns_empty_text_and_no_error(self):
+ # Distinct from a failure: the loop reports "no reply" for this, which
+ # is only correct because `error` is None here.
+ self.server.script = [assistant_message("m1", SES), session_idle(SES)]
+ text, error = self.collect()
+ self.assertEqual(text, "")
+ self.assertIsNone(error)
+
+
+class TestSendAndCollectFiltering(_FakeServerTestCase):
+ """Which events count toward the reply, and which must be ignored."""
+
+ def test_ignores_text_from_a_different_session(self):
+ self.server.script = [
+ assistant_message("m_other", "ses_stranger"),
+ text_part("p1", "m_other", "ses_stranger", "not for this loop"),
+ assistant_message("m1", SES),
+ text_part("p2", "m1", SES, "mine"),
+ session_idle(SES),
+ ]
+ text, _ = self.collect()
+ self.assertEqual(text, "mine")
+
+ def test_ignores_the_user_side_of_the_conversation(self):
+ self.server.script = [
+ user_message("m_user", SES),
+ text_part("p1", "m_user", SES, "the prompt echoed back"),
+ assistant_message("m1", SES),
+ text_part("p2", "m1", SES, "the answer"),
+ session_idle(SES),
+ ]
+ text, _ = self.collect()
+ self.assertEqual(text, "the answer")
+
+ def test_ignores_non_text_parts_such_as_tool_calls(self):
+ self.server.script = [
+ assistant_message("m1", SES),
+ tool_part("p1", "m1", SES, tool="bash"),
+ text_part("p2", "m1", SES, "ran it"),
+ session_idle(SES),
+ ]
+ text, _ = self.collect()
+ self.assertEqual(text, "ran it")
+
+ def test_a_part_arriving_BEFORE_its_message_announcement_is_dropped(self):
+ # Documents a real asymmetry with the browser client, which queues such
+ # parts and replays them once the message id is known (agentChat.ts's
+ # pendingParts). This reader has no queue: the id gate is checked once,
+ # so an out-of-order part is lost. Harmless in practice -- OpenCode
+ # announces the message first -- but worth pinning so the difference is
+ # a decision rather than a surprise.
+ self.server.script = [
+ text_part("p1", "m1", SES, "early"),
+ assistant_message("m1", SES),
+ text_part("p2", "m1", SES, "late"),
+ session_idle(SES),
+ ]
+ text, _ = self.collect()
+ self.assertEqual(text, "late")
+
+ def test_a_malformed_event_payload_does_not_abort_the_turn(self):
+ self.server.script = [
+ {"not": "a valid event"}, # no `type`
+ assistant_message("m1", SES),
+ text_part("p1", "m1", SES, "survived"),
+ session_idle(SES),
+ ]
+ text, error = self.collect()
+ self.assertEqual(text, "survived")
+ self.assertIsNone(error)
+
+ def test_an_idle_for_a_DIFFERENT_session_does_not_end_this_turn(self):
+ self.server.script = [
+ assistant_message("m1", SES),
+ session_idle("ses_stranger"),
+ text_part("p1", "m1", SES, "still streaming after the other idle"),
+ session_idle(SES),
+ ]
+ text, _ = self.collect()
+ self.assertEqual(text, "still streaming after the other idle")
+
+
+class TestSendAndCollectFailures(_FakeServerTestCase):
+ def test_a_session_error_with_no_text_is_reported_as_an_error_not_silence(self):
+ # The exact bug the (text, error) tuple exists for: a provider
+ # rejecting the request produced an empty string that looked
+ # identical to "the agent had nothing to say", so the job recorded no
+ # error and the tab showed silence under the prompt every interval.
+ self.server.script = [
+ session_error(SES, "ProviderAuthError", "OpenRouter rejected the API key"),
+ ]
+ text, error = self.collect()
+ self.assertEqual(text, "")
+ self.assertEqual(error, "OpenRouter rejected the API key")
+
+ def test_keeps_whatever_text_streamed_in_before_the_error(self):
+ self.server.script = [
+ assistant_message("m1", SES),
+ text_part("p1", "m1", SES, "Started looking, then "),
+ session_error(SES, "APIError", "upstream 502"),
+ ]
+ text, error = self.collect()
+ self.assertEqual(text, "Started looking, then ")
+ self.assertEqual(error, "upstream 502")
+
+ def test_falls_back_to_the_error_name_when_no_message_is_carried(self):
+ self.server.script = [session_error(SES, "ContextOverflowError")]
+ _text, error = self.collect()
+ self.assertEqual(error, "ContextOverflowError")
+
+ def test_an_error_for_a_different_session_is_ignored(self):
+ self.server.script = [
+ session_error("ses_stranger", "ProviderAuthError", "not mine"),
+ assistant_message("m1", SES),
+ text_part("p1", "m1", SES, "unaffected"),
+ session_idle(SES),
+ ]
+ text, error = self.collect()
+ self.assertEqual(text, "unaffected")
+ self.assertIsNone(error)
+
+ def test_a_rejected_prompt_with_no_text_collected_raises(self):
+ # The loop's _fire catches this and records it as job.last_error, which
+ # is what makes a failing check-in visible in its tab.
+ self.server.post_status = 500
+ self.server.script = []
+ with self.assertRaises(Exception):
+ self.collect(timeout=5.0)
+
+ def test_the_stream_closing_early_degrades_to_the_text_collected_so_far(self):
+ # A server restart mid-turn: EOF on the stream rather than an idle.
+ self.server.script = [
+ assistant_message("m1", SES),
+ text_part("p1", "m1", SES, "partial answer"),
+ ]
+ stop = threading.Timer(0.3, self.server.end_stream)
+ stop.start()
+ self.addCleanup(stop.cancel)
+ text, error = self.collect(timeout=5.0)
+ self.assertEqual(text, "partial answer")
+ self.assertIsNone(error)
+
+
+class TestLoopEndToEndOverTheWire(unittest.TestCase):
+ """_LoopRegistry driving a real (fake) server with NOTHING patched out --
+ session creation, the check-in, and the transcript read all go over HTTP."""
+
+ def setUp(self):
+ self.server = FakeOpencode().start()
+ self.addCleanup(self.server.stop)
+ # The registry normally spawns/reuses an `opencode serve` child; point
+ # it straight at the fake instead. This is the one thing still stubbed,
+ # and deliberately: process spawning is covered by test_opencode_process.py.
+ patcher = mock.patch.object(
+ ui_server._opencode_session, "ensure",
+ return_value={"ok": True, "url": self.server.base_url},
+ )
+ patcher.start()
+ self.addCleanup(patcher.stop)
+ self.registry = ui_server._LoopRegistry()
+ self.addCleanup(self.registry.shutdown)
+
+ def test_a_started_loop_creates_a_session_and_records_its_first_check_in(self):
+ self.server.responses["POST /session"] = {"id": "ses_loop"}
+ self.server.script = [
+ assistant_message("m1", "ses_loop"),
+ text_part("p1", "m1", "ses_loop", "Training is progressing; loss 0.42."),
+ session_idle("ses_loop"),
+ ]
+
+ result = self.registry.start("watch the loss", 60.0, "/tmp/ws", origin="http://localhost:8080")
+ self.assertTrue(result.get("ok"), result)
+ job_id = str(result["id"])
+
+ job = self.registry._jobs[job_id]
+ deadline = time.monotonic() + 10.0
+ while job.last_result is None and job.last_error is None and time.monotonic() < deadline:
+ time.sleep(0.05)
+
+ self.assertIsNone(job.last_error)
+ self.assertIn("loss 0.42", job.last_result or "")
+ # The prompt that went out carries the monitoring preamble AND the task.
+ sent = self.server.prompt_bodies[0]["parts"][0]["text"]
+ self.assertIn("watch the loss", sent)
+ self.assertIn("weightslab cli", sent)
+ self.assertIn("NEVER stop/kill a process you did not yourself start", sent)
+
+ def test_a_failing_check_in_is_recorded_as_the_jobs_error(self):
+ self.server.responses["POST /session"] = {"id": "ses_loop"}
+ self.server.script = [session_error("ses_loop", "ProviderAuthError", "no credentials")]
+
+ result = self.registry.start("watch it", 60.0, "/tmp/ws", None)
+ job = self.registry._jobs[str(result["id"])]
+ deadline = time.monotonic() + 10.0
+ while job.last_error is None and time.monotonic() < deadline:
+ time.sleep(0.05)
+
+ self.assertEqual(job.last_error, "no credentials")
+
+ def test_the_tabs_transcript_is_read_from_the_servers_own_message_list(self):
+ self.server.responses["POST /session"] = {"id": "ses_loop"}
+ self.server.script = [session_idle("ses_loop")]
+ self.server.messages = [
+ {"info": {"id": "m1", "role": "assistant"}, "parts": [{"type": "text", "text": "all good"}]},
+ ]
+ result = self.registry.start("watch it", 60.0, "/tmp/ws", None)
+ job_id = str(result["id"])
+ deadline = time.monotonic() + 10.0
+ job = self.registry._jobs[job_id]
+ while job.session_id is None and time.monotonic() < deadline:
+ time.sleep(0.05)
+
+ messages = self.registry.get_messages(job_id)
+ self.assertTrue(messages.get("ok"), messages)
+ self.assertEqual(messages["messages"], self.server.messages)
+
+
+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..fcf0fff1
--- /dev/null
+++ b/tests/ui/test_server_tracked_processes.py
@@ -0,0 +1,235 @@
+"""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 unittest.mock
+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)
+
+
+class TestSelfRegistration(_ServerTestCase):
+ """The other half of the endpoint above: a training process registering
+ ITSELF from inside ``wl.serve()``.
+
+ Why this exists: the endpoint alone only helps when someone remembers to
+ POST to it, and the only thing that ever did was the agent, by hand,
+ right after a detached launch. That step gets skipped -- the turn is
+ interrupted between launching and registering, the model forgets on a
+ relaunch, or the run was started by hand in a terminal no agent was
+ involved in -- and every miss leaves an orphaned python process holding
+ the GPU and the gRPC port after the UI is gone. serve() registering
+ itself makes it unskippable for anything that serves.
+ """
+
+ def test_serve_ui_publishes_its_own_origin_for_children_to_find(self):
+ # The whole chain hangs off this: the OpenCode server is spawned by
+ # this server, the agent's shell descends from that, and a training
+ # process it launches inherits the variable from there.
+ self.assertEqual(
+ ui_server.os.environ.get("WEIGHTSLAB_UI_ORIGIN"),
+ f"http://127.0.0.1:{self.port}",
+ )
+
+ def test_registers_this_process_against_the_live_server(self):
+ from weightslab import src as wl_src
+
+ with unittest.mock.patch.dict(
+ ui_server.os.environ,
+ {"WEIGHTSLAB_UI_ORIGIN": f"http://127.0.0.1:{self.port}"},
+ ):
+ wl_src._register_pid_with_ui_server()
+
+ # Off the calling thread on purpose (a stale origin must never
+ # stall a training run), so give the daemon thread a moment.
+ deadline = time.monotonic() + 5
+ own_pid = ui_server.os.getpid()
+ while time.monotonic() < deadline:
+ if own_pid in ui_server._tracked_processes._pids:
+ break
+ time.sleep(0.05)
+
+ # NOTE: deliberately never calls _tracked_processes.shutdown() here --
+ # the pid under test is this very test runner's. _ServerTestCase swaps
+ # in a throwaway registry and restores the real one in tearDown, so
+ # nothing else can act on it either.
+ self.assertIn(own_pid, ui_server._tracked_processes._pids)
+
+ def test_is_a_no_op_when_no_ui_server_owns_this_process(self):
+ from weightslab import src as wl_src
+
+ env = dict(ui_server.os.environ)
+ env.pop("WEIGHTSLAB_UI_ORIGIN", None)
+ with unittest.mock.patch.dict(ui_server.os.environ, env, clear=True):
+ wl_src._register_pid_with_ui_server()
+
+ time.sleep(0.2)
+ self.assertEqual(ui_server._tracked_processes._pids, set())
+
+ def test_an_unreachable_origin_never_raises(self):
+ """Best-effort is the whole contract: a stale WEIGHTSLAB_UI_ORIGIN
+ left over from a previous `weightslab start` must cost a training run
+ nothing at all, not even a raised exception on a background thread."""
+ from weightslab import src as wl_src
+
+ with unittest.mock.patch.dict(
+ ui_server.os.environ,
+ # Port 1 is reserved and never listening.
+ {"WEIGHTSLAB_UI_ORIGIN": "http://127.0.0.1:1"},
+ ):
+ wl_src._register_pid_with_ui_server()
+
+ time.sleep(0.3)
+ self.assertEqual(ui_server._tracked_processes._pids, set())
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/weightslab/AGENTS.md b/weightslab/AGENTS.md
new file mode 100644
index 00000000..283f31c6
--- /dev/null
+++ b/weightslab/AGENTS.md
@@ -0,0 +1,413 @@
+# 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).
+
+**Configuration discipline (REQUIRED).** Before writing any other code,
+collect every tunable value for the script — `batch_size`, `learning_rate`,
+`num_workers`, step budget, dataset paths, model dims, anything a user might
+reasonably want to change — into ONE plain dict (e.g. `CONFIG = {...}`, or a
+`parameters` dict loaded from YAML with `.setdefault(...)` calls, as in
+`PyTorch/wl-fraud-detection/main.py`). Wrap that dict **first**, via
+`hp = wl.watch_or_edit(CONFIG, flag="hyperparameters", defaults=CONFIG)`,
+**before** constructing the model, optimizer, or data loaders. Every
+downstream construction must then read its value from `hp`
+(`hp["batch_size"]`/`hp.get("batch_size", ...)`) — never as a hardcoded
+literal passed directly into a constructor, and never from a second, un-wrapped
+copy of the same value. A value that isn't sourced from the wrapped dict is
+invisible to the UI/agent's hyperparameter tuning (`set_hyperparam`, the HP
+panel) — it *looks* configurable but silently can't be changed, because
+nothing is reading back from the live proxy. Concretely: `DataLoader(...,
+batch_size=4)` with a literal `4` is wrong even if `CONFIG["batch_size"]`
+exists elsewhere in the file; it must be `DataLoader(..., batch_size=hp["batch_size"])`.
+
+**The dict must follow the WeightsLab config shape, not an invented one.**
+`set_hyperparam`/`show_config` resolve a handful of semantic names to fixed
+dotted paths (`_HP_ALIASES` in `trainer/services/data_service.py`) and try
+those paths **in order, using the first one that already exists** — they
+never invent a new key. A same-meaning value under a different name (e.g.
+`"epochs"` instead of `training_steps_to_do`, `"eval_every_epochs"` instead of
+`eval_full_to_train_steps_ratio`) silently falls outside that resolution and
+can't be tuned from the UI/agent at all, even though the dict has *a* key for
+it. Every generated or rewritten config — regardless of usecase — must use
+these exact top-level names and nesting, matching every shipped example
+(`PyTorch/wl-classification`, `PyTorch/wl-fraud-detection`, …):
+
+```python
+CONFIG = {
+ "experiment_name": "...",
+ "device": "auto", # resolved to cuda/cpu AFTER wrapping, never baked in as a literal
+ "root_log_dir": None, # or an explicit path; None -> a tempdir is created and logged
+ "training_steps_to_do": 1_000_000, # WL counts training STEPS (model.get_age()), never "epochs"/an epoch loop
+ "eval_full_to_train_steps_ratio": 100, # -> agent's "eval ratio" tuning; NOT "eval_every_epochs"/"eval_every_n"
+ "experiment_dump_to_train_steps_ratio": 100,# -> agent's "dump ratio" tuning; NOT "dump_every_epochs"/"checkpoint_every"
+ "optimizer": {"lr": 1e-3}, # nested under "optimizer" -> agent's "learning rate" tuning
+ "data": {
+ "train_loader": {"batch_size": 64}, # nested under "data." -> agent's "batch size" tuning
+ "test_loader": {"batch_size": 256},
+ },
+}
+```
+
+Anything with no semantic alias (`num_workers`, `grpc_port`, `start_timeout`,
+model width/depth, dataset paths, …) still belongs as a key in this same
+dict — never a bare literal in code — just without a required name; pick a
+short, descriptive one and stay consistent within the file.
+
+### 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
@@ -560,6 +794,7 @@ def _dataframe_section_html(stats: dict) -> str:
WeightsLabExperiment Report
+