ai_mouse/docs/superpowers/plans/2026-05-11-ai-mouse-library-refactor.md

ai_mouse Library Refactor Implementation Plan

For agentic workers: REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (- [ ]) syntax for tracking.

Goal: Refactor ai_mouse from a mixed application/library into a slim ONNX-Runtime-based inference SDK (src/ai_mouse/), with all training/server/eval code moved to repo-internal tools/ and pre-trained weights bundled inside the wheel.

Architecture: src/-layout package distributable via git URL; runtime depends only on numpy + onnxruntime. tools/ directory holds development-only code (torch, fastapi, etc.). Public API exposes MouseModel/ScrollModel classes plus cached generate/generate_scroll helper functions.

Tech Stack: Python 3.12+, NumPy, ONNX Runtime, hatchling (build backend), pytest, uv. Tools-only: PyTorch, FastAPI, scipy, matplotlib.

Spec: docs/superpowers/specs/2026-05-11-ai-mouse-library-design.md

---

File Structure (target end state)

ai_mouse/                         (repo root)
├── src/
│   └── ai_mouse/                 # wheel content
│       ├── __init__.py
│       ├── mouse.py
│       ├── scroll.py
│       ├── _coord.py
│       ├── _postprocess.py
│       ├── _assets.py
│       ├── errors.py
│       ├── py.typed
│       └── assets/
│           ├── flow_model.onnx
│           ├── scroll_decoder.onnx
│           ├── click_dist.json
│           ├── duration_dist.json
│           ├── train_config.json
│           └── scroll_config.json
├── tools/                        # dev-only, not in wheel
│   ├── __init__.py
│   ├── __main__.py
│   ├── train.py / serve.py / export_onnx.py
│   ├── trainer.py / models.py / collector.py / config.py
│   ├── server/ / eval/ / data_adapters/
│   └── scroll/{trainer,models,collector}.py
├── tests/{unit,tools}/
├── examples/quickstart.py
├── data/ / static/ / docs/        # unchanged
├── pyproject.toml / CHANGELOG.md / README.md / CLAUDE.md

---

Phase Map

Phase	Goal	Validation
0	Capture golden tests + train scroll model	golden npz files committed
1	Move dev-only code from ai_mouse/ to tools/	python -m tools train works; old from ai_mouse import generate still works
2	Switch to src/ layout + tighten pyproject	uv build produces clean wheel; runtime install has no torch
3	Write ONNX exporter + commit assets	tools/export_onnx.py produces .onnx files; parity check passes
4	Rewrite library in NumPy + ORT	Golden tests pass; import ai_mouse works without torch
5	Docs + cleanup	README, CHANGELOG, CLAUDE.md updated; examples runnable

---

Phase 0: Pre-flight

Task 0.1: Train scroll model

The repo has data/scroll_traces.jsonl but no trained scroll model. The current trainer in ai_mouse/scroll/trainer.py exists and works.

Files:

• Read: ai_mouse/scroll/trainer.py

• Output: data/scroll_models/{scroll_model.pt, scroll_config.json}

• Step 1: Locate the scroll training entry point

Run: uv run python -c "from ai_mouse.scroll.trainer import train; help(train)" Confirm there's a callable train(data_path, output_dir, ...) with default epochs around 100.

• Step 2: Train the scroll model

uv run python -c "
from pathlib import Path
from ai_mouse.scroll.trainer import train
train(
    data_path=Path('data/scroll_traces.jsonl'),
    output_dir=Path('data/scroll_models'),
)
"

Expected: runs ~100 epochs over ~3 minutes on CPU. Loss decreasing. Writes scroll_model.pt and scroll_config.json to data/scroll_models/.

• Step 3: Verify outputs exist

Run: ls data/scroll_models/ Expected: scroll_model.pt, scroll_config.json.

• Step 4: Smoke-test inference

uv run python -c "
from ai_mouse.scroll.generator import generate_scroll
events = generate_scroll(0, 1500, mode='target', model_dir='data/scroll_models')
print(f'Got {len(events)} events; sum deltaY = {sum(e[\"deltaY\"] for e in events)}')
"

Expected: prints something like Got 14 events; sum deltaY = 1480 (close to 1500).

• Step 5: Commit the model

git add data/scroll_models/scroll_model.pt data/scroll_models/scroll_config.json
git commit -m "chore(scroll): train initial scroll model from scroll_traces.jsonl"

---

Task 0.2: Build mouse golden npz

Capture deterministic output from the current torch-based generate() for use in regression tests later.

Files:

• Create: scripts/build_golden_mouse.py (temporary, will be deleted after Phase 4)

• Output: tests/unit/data/golden_mouse.npz

• Step 1: Ensure tests/unit/data/ exists

mkdir -p tests/unit/data

• Step 2: Create the build script

Create scripts/build_golden_mouse.py:

"""One-shot script to capture golden mouse trajectories from the current torch
implementation. Run BEFORE the migration so we can verify the numpy/ORT rewrite
in Phase 4 produces equivalent output.

Output: tests/unit/data/golden_mouse.npz
"""
from __future__ import annotations

import random
from pathlib import Path

import numpy as np
import torch

from ai_mouse import generate

CASES: list[tuple[tuple[int, int], tuple[int, int]]] = [
    ((100, 200), (900, 400)),    # horizontal 800px
    ((500, 500), (500, 100)),    # vertical 400px upward
    ((200, 600), (800, 200)),    # 720px diagonal
    ((100, 100), (130, 110)),    # very short 31px
    ((50,  50),  (1500, 900)),   # very long 1700px
    ((400, 300), (500, 300)),    # short horizontal 100px
    ((300, 300), (700, 700)),    # 45° diagonal
    ((600, 400), (200, 100)),    # reverse diagonal
]
SEEDS = (0, 1, 2, 3)


def main() -> None:
    out: dict[str, np.ndarray] = {}
    for case_idx, (start, end) in enumerate(CASES):
        for seed in SEEDS:
            random.seed(seed)
            np.random.seed(seed)
            torch.manual_seed(seed)
            pts = generate(start=start, end=end)
            out[f"case{case_idx}_seed{seed}"] = np.array(pts, dtype=np.int64)
    out_path = Path("tests/unit/data/golden_mouse.npz")
    np.savez_compressed(out_path, **out)
    print(f"Wrote {len(out)} golden traces to {out_path}")


if __name__ == "__main__":
    main()

• Step 3: Run the script

uv run python scripts/build_golden_mouse.py

Expected output:

Wrote 32 golden traces to tests/unit/data/golden_mouse.npz

• Step 4: Inspect the npz

uv run python -c "
import numpy as np
z = np.load('tests/unit/data/golden_mouse.npz')
print('keys:', list(z.keys())[:4], '...')
print('case0_seed0 shape:', z['case0_seed0'].shape)
print('case0_seed0 first 3 rows:', z['case0_seed0'][:3])
print('case0_seed0 last 2 rows (clicks):', z['case0_seed0'][-2:])
"

Expected: 32 keys, each shape (66, 3) — 64 moves + 2 click events. Last two rows share x,y; t increments.

• Step 5: Commit the golden file (not the script yet)

git add tests/unit/data/golden_mouse.npz scripts/build_golden_mouse.py
git commit -m "test: capture mouse generate() golden output (pre-migration)"

---

Task 0.3: Build scroll golden npz

Files:

• Create: scripts/build_golden_scroll.py (temporary)

• Output: tests/unit/data/golden_scroll.npz

• Step 1: Create the script

Create scripts/build_golden_scroll.py:

"""Capture golden scroll event sequences from current torch implementation."""
from __future__ import annotations

import random
from pathlib import Path

import numpy as np
import torch

from ai_mouse import generate_scroll

CASES: list[tuple[int, int, str]] = [
    (0,    1500, "target"),
    (0,    500,  "precise"),
    (0,    5000, "fast"),
    (2000, 0,    "target"),   # upward
    (0,    800,  "precise"),
    (0,    3500, "fast"),
    (1000, 1200, "precise"),  # tiny scroll
    (0,    10000, "fast"),     # very long
]
SEEDS = (0, 1, 2, 3)


def main() -> None:
    out: dict[str, np.ndarray] = {}
    for case_idx, (start_y, end_y, mode) in enumerate(CASES):
        for seed in SEEDS:
            random.seed(seed)
            np.random.seed(seed)
            torch.manual_seed(seed)
            events = generate_scroll(start_y, end_y, mode=mode)
            arr = np.array(
                [[e["deltaY"], e["deltaMode"], e["t"]] for e in events],
                dtype=np.int64,
            )
            out[f"case{case_idx}_seed{seed}"] = arr
    out_path = Path("tests/unit/data/golden_scroll.npz")
    np.savez_compressed(out_path, **out)
    print(f"Wrote {len(out)} scroll golden traces to {out_path}")


if __name__ == "__main__":
    main()

• Step 2: Run it

uv run python scripts/build_golden_scroll.py

Expected: Wrote 32 scroll golden traces to tests/unit/data/golden_scroll.npz

• Step 3: Commit

git add tests/unit/data/golden_scroll.npz scripts/build_golden_scroll.py
git commit -m "test: capture scroll generate() golden output (pre-migration)"

---

Phase 1: Move dev code out of the ai_mouse/ package

After Phase 1, ai_mouse/ package contains ONLY inference-related modules (still torch-based for now). All training/server/collector code lives under tools/. The library API from ai_mouse import generate still works because we haven't touched it yet.

Task 1.1: Scaffold tools/ directory

Files:

• Create: tools/__init__.py

• Step 1: Create tools/ and an empty init.py

mkdir -p tools/scroll
touch tools/__init__.py tools/scroll/__init__.py

• Step 2: Verify

ls tools/ tools/scroll/

Expected: __init__.py in both.

• Step 3: Commit

git add tools/__init__.py tools/scroll/__init__.py
git commit -m "chore: scaffold tools/ directory"

---

Task 1.2: Move trainer + models + utils + config to tools/

Move the torch-using mouse modules together so internal imports stay consistent within one commit.

Files:

• Move: ai_mouse/trainer.py → tools/trainer.py

• Move: ai_mouse/models.py → tools/models.py

• Move: ai_mouse/utils.py → tools/utils.py

• Move: ai_mouse/config.py → tools/config.py

• Modify: ai_mouse/generator.py (update imports)

• Step 1: git mv the files

git mv ai_mouse/trainer.py tools/trainer.py
git mv ai_mouse/models.py tools/models.py
git mv ai_mouse/utils.py tools/utils.py
git mv ai_mouse/config.py tools/config.py

• Step 2: Update imports inside moved files

In tools/trainer.py, replace:

• from ai_mouse.config import TrainConfig → from tools.config import TrainConfig
• from ai_mouse.coord import encode_trajectory → from ai_mouse.coord import encode_trajectory (unchanged — coord stays in package)
• from ai_mouse.models import TrajectoryFlowModel → from tools.models import TrajectoryFlowModel
• from ai_mouse.utils import resample_arc → from tools.utils import resample_arc

In tools/utils.py: no imports to change (pure numpy). In tools/models.py: no imports to change (torch only). In tools/config.py: no imports to change.

• Step 3: Update ai_mouse/generator.py to import torch model from tools

Find the imports near the top:

from ai_mouse.config import GenerateConfig
from ai_mouse.coord import decode_trajectory
from ai_mouse.models import TrajectoryFlowModel
from ai_mouse.utils import resample_arc

Replace with:

from ai_mouse.coord import decode_trajectory
from tools.config import GenerateConfig
from tools.models import TrajectoryFlowModel
from tools.utils import resample_arc

(Note: GenerateConfig also moved with config.py. We'll yank this cross-boundary import in Phase 4 when generator.py is replaced entirely.)

• Step 4: Verify package imports

uv run python -c "from ai_mouse import generate; print(generate.__module__)"

Expected: prints ai_mouse.generator with no ImportError.

• Step 5: Run existing tests

uv run pytest tests/test_generator.py tests/test_trainer.py tests/test_models.py -v

Expected: all pass (some test files may need import updates — see next step).

• Step 6: Update test imports if needed

In tests/test_trainer.py, tests/test_models.py, tests/conftest.py, replace:

• from ai_mouse.trainer import ... → from tools.trainer import ...

• from ai_mouse.models import ... → from tools.models import ...

• from ai_mouse.config import TrainConfig → from tools.config import TrainConfig

• Step 7: Re-run tests

uv run pytest tests/test_generator.py tests/test_trainer.py tests/test_models.py -v

Expected: all pass.

• Step 8: Commit

git add -A
git commit -m "refactor: move trainer/models/utils/config to tools/"

---

Task 1.3: Move scroll trainer / models / collector

Files:

• Move: ai_mouse/scroll/trainer.py → tools/scroll/trainer.py

• Move: ai_mouse/scroll/models.py → tools/scroll/models.py

• Move: ai_mouse/scroll/collector.py → tools/scroll/collector.py

• Modify: ai_mouse/scroll/__init__.py, ai_mouse/scroll/generator.py

• Step 1: git mv

git mv ai_mouse/scroll/trainer.py tools/scroll/trainer.py
git mv ai_mouse/scroll/models.py tools/scroll/models.py
git mv ai_mouse/scroll/collector.py tools/scroll/collector.py

• Step 2: Update imports inside moved files

In tools/scroll/trainer.py:

• from ai_mouse.scroll.models import ScrollCVAE → from tools.scroll.models import ScrollCVAE
• from ai_mouse.config import ScrollTrainConfig → from tools.config import ScrollTrainConfig

In tools/scroll/collector.py:

• from ai_mouse.config import SCROLL_MODES, ScrollModeConfig → from tools.config import SCROLL_MODES, ScrollModeConfig

• Step 3: Update ai_mouse/scroll/generator.py

Replace from ai_mouse.scroll.models import ScrollCVAE with from tools.scroll.models import ScrollCVAE.

• Step 4: Strip stale imports from ai_mouse/scroll/__init__.py

Read current content first:

cat ai_mouse/scroll/__init__.py

Edit it to only re-export generate_scroll (the only public surface that stays in the package):

"""Scroll wheel event generation (inference only)."""
from ai_mouse.scroll.generator import generate_scroll

__all__ = ["generate_scroll"]

• Step 5: Update test imports

In tests/test_scroll_trainer.py, tests/test_scroll_models.py, tests/test_scroll_collector.py:

• from ai_mouse.scroll.trainer import ... → from tools.scroll.trainer import ...
• from ai_mouse.scroll.models import ... → from tools.scroll.models import ...
• from ai_mouse.scroll.collector import ... → from tools.scroll.collector import ...

In tests/conftest.py:

• from ai_mouse.scroll.models import ScrollCVAE → from tools.scroll.models import ScrollCVAE

• Step 6: Run scroll tests

uv run pytest tests/test_scroll_*.py -v

Expected: all pass.

• Step 7: Commit

git add -A
git commit -m "refactor(scroll): move trainer/models/collector to tools/scroll/"

---

Task 1.4: Move mouse collector

Files:

• Move: ai_mouse/collector.py → tools/collector.py

• Step 1: git mv + import fix

git mv ai_mouse/collector.py tools/collector.py

In tools/collector.py, replace any from ai_mouse.config import ... with from tools.config import ....

• Step 2: Search for callers

grep -rn "from ai_mouse.collector" --include="*.py"
grep -rn "from ai_mouse import collector" --include="*.py"

Update each hit to from tools.collector import ....

• Step 3: Run tests touching collector

uv run pytest tests/ -k collector -v

Expected: pass.

• Step 4: Commit

git add -A
git commit -m "refactor: move collector to tools/"

---

Task 1.5: Move server/

Files:

• Move: ai_mouse/server/ → tools/server/

• Modify: tools/server/__init__.py (path resolution to static/)

• Step 1: git mv

git mv ai_mouse/server tools/server

• Step 2: Fix imports inside tools/server/

For each file in tools/server/ (__init__.py, deps.py, routes_collect.py, routes_train.py, routes_verify.py, routes_scroll.py), replace:

• from ai_mouse.collector import Collector → from tools.collector import Collector
• from ai_mouse.scroll.collector import ScrollCollector → from tools.scroll.collector import ScrollCollector
• from ai_mouse.scroll.trainer import train as train_scroll → from tools.scroll.trainer import train as train_scroll
• from ai_mouse.trainer import train → from tools.trainer import train
• from ai_mouse.config import ... → from tools.config import ...

Keep these unchanged (they're library API):

• from ai_mouse import generate, generate_scroll

• Step 3: Fix static path resolution in tools/server/__init__.py

The current code reads:

_HERE = Path(__file__).resolve().parent
_STATIC_DIR = _HERE.parent.parent / "static"

After moving, _HERE is tools/server/ so .parent.parent becomes the repo root — already correct. Verify by:

uv run python -c "from tools.server import create_app; app = create_app(); print('app routes:', len(app.routes))"

Expected: no error; prints route count.

• Step 4: Update test imports

In tests/test_server.py:

• from ai_mouse.server import create_app → from tools.server import create_app

• Any from ai_mouse.server.X import Y → from tools.server.X import Y

• import ai_mouse.server.deps as deps_module (if present) → import tools.server.deps as deps_module

• Step 5: Run server tests

uv run pytest tests/test_server.py -v

Expected: pass.

• Step 6: Commit

git add -A
git commit -m "refactor: move server/ to tools/server/"

---

Task 1.6: Move eval/ and data_adapters/

Files:

• Move: ai_mouse/eval/ → tools/eval/

• Move: ai_mouse/data_adapters/ → tools/data_adapters/

• Step 1: git mv

git mv ai_mouse/eval tools/eval
git mv ai_mouse/data_adapters tools/data_adapters

• Step 2: Fix imports in moved files

In tools/eval/__main__.py:

• from ai_mouse.eval.report import build_report → from tools.eval.report import build_report

In tools/eval/report.py:

• from ai_mouse.eval.metrics import ... → from tools.eval.metrics import ...

In tools/data_adapters/__main__.py:

• from ai_mouse.data_adapters.balabit import main → from tools.data_adapters.balabit import main

In tools/data_adapters/balabit.py:

• from ai_mouse.config import BalabitAdapterConfig → from tools.config import BalabitAdapterConfig

• Step 3: Update test imports

In tests/test_eval_metrics.py:

• from ai_mouse.eval.metrics import ... → from tools.eval.metrics import ...

In tests/test_balabit_adapter.py:

• from ai_mouse.data_adapters.balabit import ... → from tools.data_adapters.balabit import ...

• Step 4: Run tests

uv run pytest tests/test_eval_metrics.py tests/test_balabit_adapter.py -v

Expected: pass.

• Step 5: Commit

git add -A
git commit -m "refactor: move eval/ and data_adapters/ to tools/"

---

Task 1.7: Move CLI dispatcher

Files:

• Move: ai_mouse/__main__.py → tools/__main__.py

• Modify: tools/__main__.py (update internal subcommand wiring)

• Step 1: git mv

git mv ai_mouse/__main__.py tools/__main__.py

• Step 2: Update internal imports

In tools/__main__.py:

• from ai_mouse.trainer import train → from tools.trainer import train

• from ai_mouse.eval.__main__ import main as eval_main → from tools.eval.__main__ import main as eval_main

• from ai_mouse.data_adapters.balabit import main as bal_main → from tools.data_adapters.balabit import main as bal_main

• Step 3: Verify CLI dispatch

uv run python -m tools --help

Expected: prints help showing train, eval, balabit-adapter subcommands.

uv run python -m tools train --help

Expected: prints train-specific args.

• Step 4: Commit

git add -A
git commit -m "refactor: move CLI dispatcher to tools/__main__.py"

---

Task 1.8: Convert root main.py to tools/serve.py

Files:

• Move: main.py → tools/serve.py

• Step 1: git mv

git mv main.py tools/serve.py

• Step 2: Fix imports in tools/serve.py

from tools.server import create_app

(was from ai_mouse.server import create_app)

• Step 3: Verify it starts

In one terminal: uv run python tools/serve.py. In another: curl http://127.0.0.1:8765/api/status (or similar status endpoint). Kill the server. Then:

uv run python -c "from tools.serve import app; print('app:', app)"

Expected: prints app: <FastAPI ...> without error.

• Step 4: Commit

git add -A
git commit -m "refactor: move web entry main.py to tools/serve.py"

---

Task 1.9: Split tests into tests/unit/ and tests/tools/

Files:

• Move test files based on dependency:
  • tests/unit/: test_coord.py, test_generator.py (still uses torch via current generate() — KEEP in unit; will be rewritten in Phase 4)
  • tests/tools/: test_trainer.py, test_models.py, test_server.py, test_scroll_*.py, test_eval_metrics.py, test_balabit_adapter.py

Special case: test_generator.py and test_coord.py test the library API — they belong in tests/unit/. They depend on torch transitively today (via the current generator.py) but in Phase 4 they will not. Move them now to tests/unit/; they will keep working through both phases.

• Step 1: Create test dirs and split

mkdir -p tests/unit tests/tools
git mv tests/test_coord.py tests/unit/test_coord.py
git mv tests/test_generator.py tests/unit/test_generator.py
git mv tests/test_trainer.py tests/tools/test_trainer.py
git mv tests/test_models.py tests/tools/test_models.py
git mv tests/test_server.py tests/tools/test_server.py
git mv tests/test_scroll_collector.py tests/tools/test_scroll_collector.py
git mv tests/test_scroll_generator.py tests/unit/test_scroll_generator.py
git mv tests/test_scroll_models.py tests/tools/test_scroll_models.py
git mv tests/test_scroll_trainer.py tests/tools/test_scroll_trainer.py
git mv tests/test_eval_metrics.py tests/tools/test_eval_metrics.py
git mv tests/test_balabit_adapter.py tests/tools/test_balabit_adapter.py

• Step 2: Split conftest.py

Current tests/conftest.py provides model_dir and scroll_model_dir fixtures that use torch. These are used by tests that will end up in tests/tools/ (the torch-using ones). Move them there:

git mv tests/conftest.py tests/tools/conftest.py

Create empty tests/unit/conftest.py:

"""Fixtures for library-only tests (no torch)."""

• Step 3: Add init.py if pytest needs them

touch tests/unit/__init__.py tests/tools/__init__.py

(tests/__init__.py already exists.)

• Step 4: Run both directories separately

uv run pytest tests/unit -v
uv run pytest tests/tools -v

Expected: both pass. Some tests in tests/unit may still touch torch indirectly via the current generator.py — that's OK, will be cleared in Phase 4.

• Step 5: Commit

git add -A
git commit -m "refactor(tests): split into tests/unit and tests/tools"

---

Task 1.10: Verify whole Phase 1 outcome

• Step 1: Inspect package surface

ls ai_mouse/

Expected (Phase 1 end state): __init__.py, coord.py, generator.py, scroll/ (with __init__.py, generator.py only).

ai_mouse/scroll/:

ls ai_mouse/scroll/

Expected: __init__.py, generator.py.

• Step 2: Verify imports from each side still work

uv run python -c "
from ai_mouse import generate, generate_scroll
print('Library import OK')
from tools.trainer import train
from tools.scroll.trainer import train as st
from tools.server import create_app
from tools.eval.metrics import compute_speed
print('Tools imports OK')
"

Expected: prints both OK lines.

• Step 3: Run full test suite

uv run pytest tests/ -v

Expected: all green.

---

Phase 2: Switch to src/ layout + tighten pyproject

Task 2.1: git mv ai_mouse → src/ai_mouse

Files:

• Move: ai_mouse/ → src/ai_mouse/

• Step 1: Move the package

mkdir -p src
git mv ai_mouse src/ai_mouse

• Step 2: Verify nothing inside needs path updates

The package code uses absolute imports like from ai_mouse.coord import .... After the move, ai_mouse is still importable (because src/ becomes a path entry for setuptools/hatchling). Sanity-check there are no hard-coded paths in the source:

grep -rn "ai_mouse/" src/ai_mouse/ --include="*.py"

Expected: only string matches inside docstrings/comments, no live Path("ai_mouse/...") constructions.

• Step 3: Commit

git add -A
git commit -m "refactor: switch to src/ layout"

---

Task 2.2: Rewrite pyproject.toml (hatchling + tightened deps)

Files:

• Modify: pyproject.toml

• Step 1: Backup current pyproject

cp pyproject.toml pyproject.toml.bak

• Step 2: Write the new pyproject.toml

Replace the entire file with:

[project]
name = "ai-mouse"
version = "0.2.0"
description = "Human-like mouse trajectory and scroll wheel event generator (ONNX Runtime SDK)."
requires-python = ">=3.12,<3.14"
dependencies = [
    "numpy>=1.26.0",
    "onnxruntime>=1.17.0",
]

[project.urls]
Repository = "https://github.com/<owner>/ai_mouse"

[dependency-groups]
dev = [
    "torch>=2.2.0",
    "fastapi>=0.111.0",
    "uvicorn>=0.29.0",
    "scipy>=1.10.0",
    "matplotlib>=3.8.0",
    "pytest>=8.0.0",
    "pytest-asyncio>=0.23.0",
    "httpx>=0.27.0",
    "onnx>=1.15.0",
]

[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[tool.hatch.build.targets.wheel]
packages = ["src/ai_mouse"]

[tool.hatch.build.targets.wheel.force-include]
"src/ai_mouse/assets" = "ai_mouse/assets"

[tool.pytest.ini_options]
asyncio_mode = "auto"
testpaths = ["tests"]

Notes:

• onnx is in [dev] only because the export script in tools/export_onnx.py uses it; runtime doesn't.

• The force-include line is harmless before assets/ exists; it becomes load-bearing in Phase 3.

• Step 3: Re-sync dev environment

uv sync --group dev

Expected: completes without error. uv.lock updates.

• Step 4: Run all tests

uv run pytest tests/ -v

Expected: all pass.

• Step 5: Delete the backup

rm pyproject.toml.bak

• Step 6: Commit

git add pyproject.toml uv.lock
git commit -m "build: switch to hatchling + src layout; tighten runtime deps"

---

Task 2.3: Smoke-test wheel build

Files:

• (no files modified, just verify build)

• Step 1: Build the wheel

uv build

Expected: produces dist/ai_mouse-0.2.0-py3-none-any.whl and dist/ai_mouse-0.2.0.tar.gz.

• Step 2: Inspect wheel contents

uv run python -c "
import zipfile
with zipfile.ZipFile('dist/ai_mouse-0.2.0-py3-none-any.whl') as z:
    for n in z.namelist():
        print(n)
"

Expected: shows ai_mouse/__init__.py, ai_mouse/generator.py, ai_mouse/coord.py, etc. No tools/ content; no tests/.

• Step 3: Try installing into a clean venv

uv venv .venv-clean
.venv-clean/Scripts/python -m pip install dist/ai_mouse-0.2.0-py3-none-any.whl
.venv-clean/Scripts/python -c "import ai_mouse; print(ai_mouse.__file__)"

Expected: import works. Note torch is NOT installed in this venv, so from ai_mouse import generate will FAIL right now (current generator.py still imports torch). That's expected pre-Phase-4 — just confirm import ai_mouse itself succeeds (it doesn't trigger generator.py).

Actually ai_mouse/__init__.py does from ai_mouse.generator import generate, which transitively imports torch. So this import WILL fail. Expected outcome:

ModuleNotFoundError: No module named 'torch'

Confirms the wheel content is correct but the runtime promise isn't met yet — exactly the state we expect at Phase 2 end. Document this in commit message.

• Step 4: Clean up

rm -rf .venv-clean dist/

• Step 5: No commit needed (verification only)

---

Phase 3: ONNX export

Task 3.1: Write the mouse-model export portion of tools/export_onnx.py

Files:

• Create: tools/export_onnx.py

• Step 1: Create the file with imports and helpers

Create tools/export_onnx.py:

"""Export trained PyTorch checkpoints to ONNX for the inference SDK.

Usage:
    uv run python tools/export_onnx.py \
        --flow-ckpt data/models_v2 \
        --scroll-ckpt data/scroll_models \
        --output src/ai_mouse/assets/

Produces:
    <output>/flow_model.onnx
    <output>/scroll_decoder.onnx
    <output>/click_dist.json
    <output>/duration_dist.json
    <output>/train_config.json
    <output>/scroll_config.json

A PyTorch vs ONNX Runtime parity check runs at the end. If parity fails
the .onnx files are deleted to prevent shipping broken weights.
"""
from __future__ import annotations

import argparse
import json
import logging
import shutil
import sys
from pathlib import Path

import numpy as np
import torch

logger = logging.getLogger(__name__)

_ATOL = 1e-4

• Step 2: Add export_flow_model function

Append to tools/export_onnx.py:

def export_flow_model(ckpt_dir: Path, out_dir: Path) -> Path:
    """Export TrajectoryFlowModel to ONNX.

    Args:
        ckpt_dir: directory with flow_model.pt and train_config.json.
        out_dir: destination directory (created if missing).

    Returns:
        Path to the written flow_model.onnx.
    """
    from tools.models import TrajectoryFlowModel

    config_path = ckpt_dir / "train_config.json"
    cfg = json.loads(config_path.read_text())
    seq_len = int(cfg["seq_len"])
    d_model = int(cfg["d_model"])
    nhead = int(cfg["nhead"])
    num_layers = int(cfg["num_layers"])
    dim_feedforward = int(cfg["dim_feedforward"])
    cond_dim = int(cfg.get("cond_dim", 3))

    model = TrajectoryFlowModel(
        seq_len=seq_len,
        d_model=d_model,
        nhead=nhead,
        num_layers=num_layers,
        dim_feedforward=dim_feedforward,
        cond_dim=cond_dim,
        dropout=0.0,  # disable dropout for export
    )
    state = torch.load(ckpt_dir / "flow_model.pt", map_location="cpu", weights_only=True)
    model.load_state_dict(state)
    model.eval()

    out_dir.mkdir(parents=True, exist_ok=True)
    out_path = out_dir / "flow_model.onnx"

    dummy_x = torch.zeros(1, seq_len, 3, dtype=torch.float32)
    dummy_t = torch.zeros(1, dtype=torch.float32)
    dummy_cond = torch.zeros(1, cond_dim, dtype=torch.float32)

    torch.onnx.export(
        model,
        (dummy_x, dummy_t, dummy_cond),
        str(out_path),
        input_names=["x_t", "t", "cond"],
        output_names=["v"],
        dynamic_axes={
            "x_t": {0: "batch"},
            "t": {0: "batch"},
            "cond": {0: "batch"},
            "v": {0: "batch"},
        },
        opset_version=17,
        do_constant_folding=True,
    )
    logger.info("Wrote %s (%.1f MB)", out_path, out_path.stat().st_size / 1e6)
    return out_path

• Step 3: Add a quick sanity test (manual run)

In a python shell:

uv run python -c "
from pathlib import Path
from tools.export_onnx import export_flow_model
out = export_flow_model(Path('data/models_v2'), Path('/tmp/test_export'))
print('Wrote:', out)
"

Expected: prints the output path and a size like 2-3 MB. The file exists.

• Step 4: Commit

git add tools/export_onnx.py
git commit -m "feat(tools): add export_flow_model for ONNX export"

---

Task 3.2: Add scroll-decoder export

Files:

• Modify: tools/export_onnx.py

• Step 1: Define ScrollDecoder wrapper module

Append to tools/export_onnx.py:

class _ScrollDecoder(torch.nn.Module):
    """Wraps ScrollCVAE.decode for ONNX export.

    The full ScrollCVAE is encoder+decoder; inference only needs decoder.
    """

    def __init__(self, dec_h0, dec_gru, dec_out, seq_len: int, hidden: int):
        super().__init__()
        self.dec_h0 = dec_h0
        self.dec_gru = dec_gru
        self.dec_out = dec_out
        self.seq_len = seq_len
        self.hidden = hidden

    def forward(self, z: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
        b = z.shape[0]
        zc = torch.cat([z, cond], dim=-1)
        h0_flat = self.dec_h0(zc)
        h0 = h0_flat.view(b, 2, self.hidden).permute(1, 0, 2).contiguous()
        inp = zc.unsqueeze(1).expand(b, self.seq_len, -1)
        out, _ = self.dec_gru(inp, h0)
        return self.dec_out(out)

• Step 2: Add export_scroll_decoder function

Append:

def export_scroll_decoder(ckpt_dir: Path, out_dir: Path) -> Path:
    """Export ScrollCVAE decoder to ONNX."""
    from tools.scroll.models import ScrollCVAE

    config_path = ckpt_dir / "scroll_config.json"
    cfg = json.loads(config_path.read_text())
    seq_len = int(cfg["seq_len"])
    latent_dim = int(cfg["latent_dim"])
    hidden = int(cfg["hidden"])
    cond_dim = int(cfg["cond_dim"])

    full = ScrollCVAE(
        seq_len=seq_len, latent_dim=latent_dim, hidden=hidden, cond_dim=cond_dim
    )
    state = torch.load(ckpt_dir / "scroll_model.pt", map_location="cpu", weights_only=True)
    full.load_state_dict(state)
    full.eval()

    decoder = _ScrollDecoder(
        dec_h0=full.dec_h0,
        dec_gru=full.dec_gru,
        dec_out=full.dec_out,
        seq_len=seq_len,
        hidden=hidden,
    )
    decoder.eval()

    out_dir.mkdir(parents=True, exist_ok=True)
    out_path = out_dir / "scroll_decoder.onnx"

    dummy_z = torch.zeros(1, latent_dim, dtype=torch.float32)
    dummy_cond = torch.zeros(1, cond_dim, dtype=torch.float32)

    torch.onnx.export(
        decoder,
        (dummy_z, dummy_cond),
        str(out_path),
        input_names=["z", "cond"],
        output_names=["seq"],
        dynamic_axes={
            "z": {0: "batch"},
            "cond": {0: "batch"},
            "seq": {0: "batch"},
        },
        opset_version=17,
        do_constant_folding=True,
    )
    logger.info("Wrote %s (%.1f KB)", out_path, out_path.stat().st_size / 1e3)
    return out_path

• Step 3: Manual sanity test

uv run python -c "
from pathlib import Path
from tools.export_onnx import export_scroll_decoder
out = export_scroll_decoder(Path('data/scroll_models'), Path('/tmp/test_export'))
print('Wrote:', out)
"

Expected: prints path; file <300 KB.

• Step 4: Commit

git add tools/export_onnx.py
git commit -m "feat(tools): add export_scroll_decoder for ONNX export"

---

Task 3.3: Add PyTorch vs ORT parity check

Files:

• Modify: tools/export_onnx.py

• Step 1: Add parity helpers

Append to tools/export_onnx.py:

def _check_flow_parity(ckpt_dir: Path, onnx_path: Path) -> None:
    """Verify ONNX flow model matches PyTorch output on random input."""
    import onnxruntime as ort
    from tools.models import TrajectoryFlowModel

    cfg = json.loads((ckpt_dir / "train_config.json").read_text())
    seq_len = int(cfg["seq_len"])
    cond_dim = int(cfg.get("cond_dim", 3))

    model = TrajectoryFlowModel(
        seq_len=seq_len,
        d_model=int(cfg["d_model"]),
        nhead=int(cfg["nhead"]),
        num_layers=int(cfg["num_layers"]),
        dim_feedforward=int(cfg["dim_feedforward"]),
        cond_dim=cond_dim,
        dropout=0.0,
    )
    model.load_state_dict(
        torch.load(ckpt_dir / "flow_model.pt", map_location="cpu", weights_only=True)
    )
    model.eval()

    torch.manual_seed(42)
    np.random.seed(42)
    x = torch.randn(2, seq_len, 3, dtype=torch.float32)
    t = torch.tensor([0.0, 0.5], dtype=torch.float32)
    cond = torch.randn(2, cond_dim, dtype=torch.float32)

    with torch.no_grad():
        torch_out = model(x, t, cond).numpy()

    sess = ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"])
    ort_out = sess.run(
        ["v"],
        {
            "x_t": x.numpy(),
            "t": t.numpy(),
            "cond": cond.numpy(),
        },
    )[0]

    if not np.allclose(torch_out, ort_out, atol=_ATOL, rtol=1e-3):
        max_diff = float(np.abs(torch_out - ort_out).max())
        raise RuntimeError(
            f"Flow model ORT/PyTorch parity FAILED: max abs diff = {max_diff:.2e} > {_ATOL:.2e}"
        )
    logger.info("Flow model parity OK (atol=%.0e)", _ATOL)


def _check_scroll_parity(ckpt_dir: Path, onnx_path: Path) -> None:
    """Verify ONNX scroll decoder matches PyTorch decoder output."""
    import onnxruntime as ort
    from tools.scroll.models import ScrollCVAE

    cfg = json.loads((ckpt_dir / "scroll_config.json").read_text())
    seq_len = int(cfg["seq_len"])
    latent_dim = int(cfg["latent_dim"])
    cond_dim = int(cfg["cond_dim"])

    full = ScrollCVAE(
        seq_len=seq_len,
        latent_dim=latent_dim,
        hidden=int(cfg["hidden"]),
        cond_dim=cond_dim,
    )
    full.load_state_dict(
        torch.load(ckpt_dir / "scroll_model.pt", map_location="cpu", weights_only=True)
    )
    full.eval()

    torch.manual_seed(7)
    z = torch.randn(2, latent_dim, dtype=torch.float32)
    cond = torch.randn(2, cond_dim, dtype=torch.float32)

    with torch.no_grad():
        torch_out = full.decode(z, cond).numpy()

    sess = ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"])
    ort_out = sess.run(["seq"], {"z": z.numpy(), "cond": cond.numpy()})[0]

    if not np.allclose(torch_out, ort_out, atol=_ATOL, rtol=1e-3):
        max_diff = float(np.abs(torch_out - ort_out).max())
        raise RuntimeError(
            f"Scroll decoder parity FAILED: max abs diff = {max_diff:.2e} > {_ATOL:.2e}"
        )
    logger.info("Scroll decoder parity OK (atol=%.0e)", _ATOL)

• Step 2: Manual test the checks

uv run python -c "
from pathlib import Path
from tools.export_onnx import (
    export_flow_model, export_scroll_decoder,
    _check_flow_parity, _check_scroll_parity,
)
import logging; logging.basicConfig(level=logging.INFO)
out = Path('/tmp/test_export')
export_flow_model(Path('data/models_v2'), out)
_check_flow_parity(Path('data/models_v2'), out / 'flow_model.onnx')
export_scroll_decoder(Path('data/scroll_models'), out)
_check_scroll_parity(Path('data/scroll_models'), out / 'scroll_decoder.onnx')
"

Expected: prints two "parity OK" lines, no exceptions.

• Step 3: Commit

git add tools/export_onnx.py
git commit -m "feat(tools): add ORT vs PyTorch parity check for exports"

---

Task 3.4: Add CLI main() to tools/export_onnx.py

Files:

• Modify: tools/export_onnx.py

• Step 1: Add main() and main guard

Append to tools/export_onnx.py:

def _copy_metadata(flow_dir: Path, scroll_dir: Path, out_dir: Path) -> None:
    """Copy JSON metadata files alongside the ONNX models."""
    for name in ("click_dist.json", "duration_dist.json", "train_config.json"):
        src = flow_dir / name
        if not src.exists():
            raise FileNotFoundError(f"Required metadata missing: {src}")
        shutil.copy2(src, out_dir / name)
    src = scroll_dir / "scroll_config.json"
    if not src.exists():
        raise FileNotFoundError(f"Required metadata missing: {src}")
    shutil.copy2(src, out_dir / "scroll_config.json")


def main(argv: list[str] | None = None) -> int:
    p = argparse.ArgumentParser(prog="export_onnx", description=__doc__.splitlines()[0])
    p.add_argument("--flow-ckpt", type=Path, required=True)
    p.add_argument("--scroll-ckpt", type=Path, required=True)
    p.add_argument("--output", type=Path, required=True)
    args = p.parse_args(argv)

    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")

    args.output.mkdir(parents=True, exist_ok=True)

    flow_onnx = export_flow_model(args.flow_ckpt, args.output)
    scroll_onnx = export_scroll_decoder(args.scroll_ckpt, args.output)

    try:
        _check_flow_parity(args.flow_ckpt, flow_onnx)
        _check_scroll_parity(args.scroll_ckpt, scroll_onnx)
    except RuntimeError as exc:
        logger.error("Parity check failed: %s", exc)
        flow_onnx.unlink(missing_ok=True)
        scroll_onnx.unlink(missing_ok=True)
        return 1

    _copy_metadata(args.flow_ckpt, args.scroll_ckpt, args.output)
    logger.info("Export complete: %s", args.output)
    return 0


if __name__ == "__main__":
    sys.exit(main())

• Step 2: Run the full export to produce assets

mkdir -p src/ai_mouse/assets
uv run python tools/export_onnx.py \
    --flow-ckpt data/models_v2 \
    --scroll-ckpt data/scroll_models \
    --output src/ai_mouse/assets/

Expected output (timestamps elided):

... INFO Wrote src/ai_mouse/assets/flow_model.onnx (2.x MB)
... INFO Wrote src/ai_mouse/assets/scroll_decoder.onnx (0.x KB)
... INFO Flow model parity OK (atol=1e-04)
... INFO Scroll decoder parity OK (atol=1e-04)
... INFO Export complete: src/ai_mouse/assets

• Step 3: Verify assets directory

ls src/ai_mouse/assets/

Expected: flow_model.onnx, scroll_decoder.onnx, click_dist.json, duration_dist.json, train_config.json, scroll_config.json.

• Step 4: Commit assets + script main()

git add tools/export_onnx.py src/ai_mouse/assets/
git commit -m "feat: export ONNX weights and metadata into src/ai_mouse/assets/"

---

Task 3.5: Test ONNX export with toy model

Files:

• Create: tests/tools/test_export_onnx.py

• Step 1: Write the failing test

Create tests/tools/test_export_onnx.py:

"""Validate tools.export_onnx with a tiny synthetic model."""
from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import pytest
import torch

from tools.export_onnx import (
    _check_flow_parity,
    _check_scroll_parity,
    export_flow_model,
    export_scroll_decoder,
)


@pytest.fixture
def tiny_flow_ckpt(tmp_path: Path) -> Path:
    """A flow model with seq_len=8, d_model=16, 1 layer — small but valid."""
    from tools.models import TrajectoryFlowModel

    cfg = {
        "seq_len": 8,
        "d_model": 16,
        "nhead": 2,
        "num_layers": 1,
        "dim_feedforward": 32,
        "cond_dim": 3,
    }
    model = TrajectoryFlowModel(**cfg, dropout=0.0)
    model.eval()
    out = tmp_path / "flow_ckpt"
    out.mkdir()
    torch.save(model.state_dict(), out / "flow_model.pt")
    (out / "train_config.json").write_text(json.dumps(cfg))
    return out


@pytest.fixture
def tiny_scroll_ckpt(tmp_path: Path) -> Path:
    """A scroll model with seq_len=4, latent=4, hidden=8."""
    from tools.scroll.models import ScrollCVAE

    cfg = {"seq_len": 4, "latent_dim": 4, "hidden": 8, "cond_dim": 7}
    model = ScrollCVAE(**cfg)
    model.eval()
    out = tmp_path / "scroll_ckpt"
    out.mkdir()
    torch.save(model.state_dict(), out / "scroll_model.pt")
    (out / "scroll_config.json").write_text(json.dumps(cfg))
    return out


def test_export_flow_model_parity(tiny_flow_ckpt: Path, tmp_path: Path) -> None:
    out_dir = tmp_path / "out"
    onnx_path = export_flow_model(tiny_flow_ckpt, out_dir)
    assert onnx_path.exists()
    _check_flow_parity(tiny_flow_ckpt, onnx_path)  # raises on failure


def test_export_scroll_decoder_parity(tiny_scroll_ckpt: Path, tmp_path: Path) -> None:
    out_dir = tmp_path / "out"
    onnx_path = export_scroll_decoder(tiny_scroll_ckpt, out_dir)
    assert onnx_path.exists()
    _check_scroll_parity(tiny_scroll_ckpt, onnx_path)

• Step 2: Run the tests

uv run pytest tests/tools/test_export_onnx.py -v

Expected: both pass.

• Step 3: Commit

git add tests/tools/test_export_onnx.py
git commit -m "test(tools): cover export_onnx with tiny synthetic models"

---

Phase 4: Rewrite library in NumPy + ORT

Task 4.1: Create _coord.py (private numpy coordinate transforms)

Files:

• Create: src/ai_mouse/_coord.py

• Keep (for now): src/ai_mouse/coord.py — tools/ still imports it; deleted at end of Phase 4

• Step 1: Copy coord.py to _coord.py

cp src/ai_mouse/coord.py src/ai_mouse/_coord.py

• Step 2: No content edits needed

The file is already pure numpy. Verify:

grep -E "^import|^from" src/ai_mouse/_coord.py

Expected: only import math and import numpy as np.

• Step 3: Write the test

Create tests/unit/test__coord.py:

"""Test the private numpy coordinate transforms."""
from __future__ import annotations

import numpy as np

from ai_mouse._coord import decode_trajectory, encode_trajectory


def test_encode_decode_roundtrip() -> None:
    points = np.array([[100.0, 200.0], [300.0, 250.0], [500.0, 300.0]])
    start = (100, 200)
    end = (500, 300)
    encoded = encode_trajectory(points, start, end)
    decoded = decode_trajectory(encoded, start, end)
    assert np.allclose(decoded, points, atol=1e-6)


def test_encode_endpoints() -> None:
    """Start should encode to (0,0); end should encode to (1,0)."""
    points = np.array([[100.0, 200.0], [500.0, 300.0]])
    encoded = encode_trajectory(points, (100, 200), (500, 300))
    assert np.allclose(encoded[0], [0.0, 0.0], atol=1e-6)
    assert np.allclose(encoded[1], [1.0, 0.0], atol=1e-6)


def test_zero_distance_returns_zeros() -> None:
    points = np.array([[100.0, 200.0]])
    encoded = encode_trajectory(points, (100, 200), (100, 200))
    assert encoded.shape == (1, 2)
    assert np.all(encoded == 0)

• Step 4: Run test

uv run pytest tests/unit/test__coord.py -v

Expected: 3 pass.

• Step 5: Commit

git add src/ai_mouse/_coord.py tests/unit/test__coord.py
git commit -m "feat(lib): add private _coord.py with numpy transforms"

---

Task 4.2: Create errors.py

Files:

• Create: src/ai_mouse/errors.py

• Step 1: Write the failing test

Create tests/unit/test_errors.py:

"""Test the error hierarchy."""
from __future__ import annotations

import pytest

from ai_mouse import errors


def test_model_load_error_is_aimouse_error() -> None:
    assert issubclass(errors.ModelLoadError, errors.AiMouseError)


def test_generation_error_is_aimouse_error() -> None:
    assert issubclass(errors.GenerationError, errors.AiMouseError)


def test_can_catch_specific_with_general() -> None:
    with pytest.raises(errors.AiMouseError):
        raise errors.ModelLoadError("test")

• Step 2: Run test, observe failure

uv run pytest tests/unit/test_errors.py -v

Expected: ImportError on from ai_mouse import errors.

• Step 3: Create the module

Create src/ai_mouse/errors.py:

"""Exception hierarchy for the ai_mouse library.

Downstream consumers can catch the umbrella :class:`AiMouseError`
or the specific subclasses for finer control.
"""
from __future__ import annotations


class AiMouseError(Exception):
    """Base class for all ai_mouse errors."""


class ModelLoadError(AiMouseError):
    """Raised when ONNX weights / metadata cannot be loaded."""


class GenerationError(AiMouseError):
    """Raised when inference produces an invalid result (e.g. NaN)."""

• Step 4: Run test, observe pass

uv run pytest tests/unit/test_errors.py -v

Expected: 3 pass.

• Step 5: Commit

git add src/ai_mouse/errors.py tests/unit/test_errors.py
git commit -m "feat(lib): add errors module"

---

Task 4.3: Create _assets.py (importlib.resources loader)

Files:

• Create: src/ai_mouse/_assets.py

• Step 1: Write the failing test

Create tests/unit/test_assets.py:

"""Test the asset path resolver."""
from __future__ import annotations

import json
from pathlib import Path

import pytest

from ai_mouse import _assets
from ai_mouse.errors import ModelLoadError


def test_bundled_flow_model_exists() -> None:
    p = _assets.bundled_path("flow_model.onnx")
    assert p.exists()
    assert p.suffix == ".onnx"


def test_bundled_train_config_loadable() -> None:
    p = _assets.bundled_path("train_config.json")
    cfg = json.loads(p.read_text())
    assert "seq_len" in cfg
    assert "d_model" in cfg


def test_resolve_with_custom_dir(tmp_path: Path) -> None:
    (tmp_path / "flow_model.onnx").write_bytes(b"x")
    p = _assets.resolve(tmp_path, "flow_model.onnx")
    assert p == tmp_path / "flow_model.onnx"


def test_missing_asset_raises_model_load_error(tmp_path: Path) -> None:
    with pytest.raises(ModelLoadError, match="missing"):
        _assets.resolve(tmp_path, "nonexistent.onnx")

• Step 2: Run test, observe failure

uv run pytest tests/unit/test_assets.py -v

Expected: ImportError on from ai_mouse import _assets.

• Step 3: Create the module

Create src/ai_mouse/_assets.py:

"""Asset path resolution for bundled ONNX weights and JSON metadata.

Uses :mod:`importlib.resources` to locate files inside the installed
package, falling back to a user-supplied directory if provided.
"""
from __future__ import annotations

from importlib.resources import as_file, files
from pathlib import Path

from ai_mouse.errors import ModelLoadError

_PACKAGE_ASSETS = "ai_mouse.assets"


def bundled_path(name: str) -> Path:
    """Return a filesystem path to a bundled asset.

    Args:
        name: filename inside the assets/ directory.

    Returns:
        A concrete :class:`pathlib.Path`. Note: for zipapp installations
        this materialises a temp file; for normal site-packages installs
        it points into the package directly.
    """
    ref = files(_PACKAGE_ASSETS) / name
    # as_file is the canonical way; for non-zip installs this is a no-op
    # context that yields the actual path.
    with as_file(ref) as p:
        # We're inside the with-block; the contextmanager keeps the
        # temp file alive only while open. For zip installs we'd need
        # to extract to a stable location. For now, all our installs
        # are wheel-based (non-zip), so the path is stable after exit.
        return Path(p)


def resolve(model_path: Path | None, filename: str) -> Path:
    """Locate an asset given an optional user-supplied directory.

    Args:
        model_path: user-supplied directory, or None to use bundled assets.
        filename: file to locate inside the directory.

    Returns:
        Absolute path to the asset.

    Raises:
        ModelLoadError: if the file does not exist.
    """
    if model_path is None:
        p = bundled_path(filename)
    else:
        p = Path(model_path) / filename
    if not p.exists():
        raise ModelLoadError(f"Required asset missing: {p}")
    return p

• Step 4: Run test, observe pass

uv run pytest tests/unit/test_assets.py -v

Expected: 4 pass.

• Step 5: Commit

git add src/ai_mouse/_assets.py tests/unit/test_assets.py
git commit -m "feat(lib): add _assets module for bundled-weight resolution"

---

Task 4.4: Create _postprocess.py skeleton + gaussian_smooth

Files:

• Create: src/ai_mouse/_postprocess.py

• Create: tests/unit/test_postprocess.py

• Step 1: Write failing test

Create tests/unit/test_postprocess.py:

"""Tests for trajectory post-processing primitives."""
from __future__ import annotations

import numpy as np

from ai_mouse._postprocess import gaussian_smooth


def test_gaussian_smooth_preserves_endpoints() -> None:
    x = np.array([1.0, 5.0, 3.0, 8.0, 2.0, 6.0, 4.0])
    result = gaussian_smooth(x, sigma=1.0)
    assert result[0] == 1.0
    assert result[-1] == 4.0


def test_gaussian_smooth_short_input_unchanged() -> None:
    x = np.array([1.0, 2.0, 3.0])
    result = gaussian_smooth(x, sigma=1.0)
    assert np.array_equal(result, x)


def test_gaussian_smooth_constant_unchanged() -> None:
    x = np.full(20, 7.5)
    result = gaussian_smooth(x, sigma=1.0)
    assert np.allclose(result, x, atol=1e-6)

• Step 2: Run, observe failure

uv run pytest tests/unit/test_postprocess.py -v

Expected: ImportError.

• Step 3: Create module + function

Create src/ai_mouse/_postprocess.py:

"""Pure-numpy post-processing primitives for trajectory generation.

All functions are pure (no I/O, no global state) and accept an explicit
:class:`numpy.random.Generator` when randomness is involved.
"""
from __future__ import annotations

import numpy as np


def gaussian_smooth(x: np.ndarray, sigma: float = 1.0) -> np.ndarray:
    """5-tap gaussian smoothing along a 1-D array; endpoints preserved.

    Args:
        x: 1-D input array.
        sigma: gaussian std. Default 1.0 gives weights ≈
               [0.054, 0.244, 0.403, 0.244, 0.054].

    Returns:
        Smoothed array of the same shape. ``x[0]`` and ``x[-1]`` unchanged.
        If ``len(x) < 5`` returns a copy of ``x`` (kernel won't fit).
    """
    if len(x) < 5:
        return x.copy()
    kernel = np.exp(-0.5 * (np.arange(-2, 3) / sigma) ** 2)
    kernel /= kernel.sum()
    padded = np.pad(x, pad_width=2, mode="edge")
    smoothed = np.convolve(padded, kernel, mode="valid")
    smoothed[0] = x[0]
    smoothed[-1] = x[-1]
    return smoothed

• Step 4: Run, observe pass

uv run pytest tests/unit/test_postprocess.py -v

Expected: 3 pass.

• Step 5: Commit

git add src/ai_mouse/_postprocess.py tests/unit/test_postprocess.py
git commit -m "feat(lib): add gaussian_smooth to _postprocess"

---

Task 4.5: Add snap_endpoints

Files:

• Modify: src/ai_mouse/_postprocess.py

• Modify: tests/unit/test_postprocess.py

• Step 1: Write failing test (append)

Append to tests/unit/test_postprocess.py:

from ai_mouse._postprocess import snap_endpoints


def test_snap_endpoints_pins_first_and_last() -> None:
    forward = np.linspace(0.1, 0.9, 16)
    lateral = np.full(16, 0.5)
    f, l = snap_endpoints(forward.copy(), lateral.copy(), seq_len=16)
    assert f[0] == 0.0
    assert l[0] == 0.0
    assert f[-1] == 1.0
    assert l[-1] == 0.0


def test_snap_endpoints_preserves_middle() -> None:
    forward = np.linspace(0.0, 1.0, 16)
    lateral = np.zeros(16)
    f, _ = snap_endpoints(forward.copy(), lateral.copy(), seq_len=16, n_snap=4)
    # Points before the last n_snap should be unchanged
    assert np.allclose(f[1 : 16 - 4], forward[1 : 16 - 4], atol=1e-6)

• Step 2: Run, observe failure

uv run pytest tests/unit/test_postprocess.py::test_snap_endpoints_pins_first_and_last -v

Expected: ImportError.

• Step 3: Implement

Append to src/ai_mouse/_postprocess.py:

def snap_endpoints(
    forward: np.ndarray,
    lateral: np.ndarray,
    seq_len: int,
    n_snap: int = 6,
) -> tuple[np.ndarray, np.ndarray]:
    """Force first point to (0,0) and last point to (1,0) with quadratic ease.

    The last ``n_snap`` points are linearly interpolated towards (1, 0)
    with quadratic easing, then the first/last points are pinned exactly.

    Args:
        forward: (T,) forward coordinates (modified in place).
        lateral: (T,) lateral coordinates (modified in place).
        seq_len: length of forward/lateral.
        n_snap: number of trailing points to ease (capped at seq_len//4).

    Returns:
        ``(forward, lateral)`` after modification.
    """
    n_snap = min(n_snap, seq_len // 4)
    for i in range(n_snap):
        alpha = ((i + 1) / n_snap) ** 2
        k = seq_len - n_snap + i
        forward[k] = forward[k] * (1.0 - alpha) + 1.0 * alpha
        lateral[k] = lateral[k] * (1.0 - alpha) + 0.0 * alpha
    forward[0], lateral[0] = 0.0, 0.0
    forward[-1], lateral[-1] = 1.0, 0.0
    return forward, lateral

• Step 4: Run all postprocess tests

uv run pytest tests/unit/test_postprocess.py -v

Expected: all pass.

• Step 5: Commit

git add src/ai_mouse/_postprocess.py tests/unit/test_postprocess.py
git commit -m "feat(lib): add snap_endpoints to _postprocess"

---

Task 4.6: Add smooth_start, enforce_forward_monotonic

Files:

• Modify: src/ai_mouse/_postprocess.py, tests/unit/test_postprocess.py

• Step 1: Write tests (append)

from ai_mouse._postprocess import enforce_forward_monotonic, smooth_start


def test_smooth_start_dampens_lateral() -> None:
    forward = np.linspace(0, 1, 16)
    lateral = np.full(16, 1.0)
    forward[0] = lateral[0] = 0.0  # invariant: snap already done
    _, l = smooth_start(forward.copy(), lateral.copy(), n=4)
    # Lateral at points 1-4 should be < original (dampened)
    assert l[1] < 1.0
    assert l[4] < 1.0
    # Lateral at point 5+ unchanged
    assert l[5] == 1.0


def test_enforce_forward_monotonic_repairs_inversions() -> None:
    f = np.array([0.0, 0.4, 0.3, 0.6, 0.5, 1.0])
    out = enforce_forward_monotonic(f.copy())
    assert np.all(np.diff(out) > 0), out


def test_enforce_forward_monotonic_clips_to_unit_interval() -> None:
    f = np.array([-0.1, 0.5, 1.2])
    out = enforce_forward_monotonic(f.copy())
    assert out[0] == 0.0
    assert out[-1] == 1.0

• Step 2: Run, observe failure

uv run pytest tests/unit/test_postprocess.py -v

• Step 3: Implement (append to _postprocess.py)

def smooth_start(
    forward: np.ndarray,
    lateral: np.ndarray,
    n: int = 4,
) -> tuple[np.ndarray, np.ndarray]:
    """Dampen lateral oscillation in the first ``n`` points.

    Assumes :func:`snap_endpoints` has already pinned (0,0). Forward is
    forced non-decreasing locally; lateral is linearly damped towards 0.
    """
    n_start_fix = min(n, len(forward) // 4)
    for i in range(1, n_start_fix + 1):
        blend = i / (n_start_fix + 1)
        forward[i] = max(forward[i], forward[i - 1])
        lateral[i] = lateral[i] * blend
    return forward, lateral


def enforce_forward_monotonic(forward: np.ndarray) -> np.ndarray:
    """Force ``forward`` non-decreasing, clip to [0,1], pin endpoints."""
    seq_len = len(forward)
    for i in range(1, seq_len - 1):
        if forward[i] < forward[i - 1]:
            forward[i] = forward[i - 1] + 0.001
    forward = np.clip(forward, 0.0, 1.0)
    forward[0] = 0.0
    forward[-1] = 1.0
    return forward

• Step 4: Test

uv run pytest tests/unit/test_postprocess.py -v

Expected: all pass.

• Step 5: Commit

git add src/ai_mouse/_postprocess.py tests/unit/test_postprocess.py
git commit -m "feat(lib): add smooth_start, enforce_forward_monotonic"

---

Task 4.7: Add resample_arc, build_timestamps

Files:

• Modify: src/ai_mouse/_postprocess.py, tests/unit/test_postprocess.py

• Step 1: Tests

Append:

from ai_mouse._postprocess import build_timestamps, resample_arc


def test_resample_arc_identity_when_same_length() -> None:
    pts = np.array([[0.0, 0.0], [1.0, 1.0], [2.0, 0.0], [3.0, 1.0]])
    out = resample_arc(pts, 4)
    assert np.allclose(out, pts, atol=1e-6)


def test_resample_arc_changes_length() -> None:
    pts = np.array([[float(i), 0.0] for i in range(10)])
    out = resample_arc(pts, 5)
    assert out.shape == (5, 2)
    # Endpoints preserved
    assert np.allclose(out[0], pts[0])
    assert np.allclose(out[-1], pts[-1])


def test_build_timestamps_strictly_increasing() -> None:
    log_dt = np.array([0.0, 2.0, 2.5, 3.0, 2.0])
    ts = build_timestamps(log_dt, total_duration_ms=200.0)
    assert ts[0] == 0
    assert np.all(np.diff(ts) >= 1)  # at least 1 ms apart


def test_build_timestamps_total_close_to_target() -> None:
    log_dt = np.array([1.0] * 10)
    ts = build_timestamps(log_dt, total_duration_ms=300.0)
    # Last timestamp should be roughly total - one slot
    assert abs(ts[-1] - 270) < 60  # tolerant of clipping

• Step 2: Run, observe failure

uv run pytest tests/unit/test_postprocess.py::test_resample_arc_identity_when_same_length -v

• Step 3: Implement

Append to _postprocess.py:

def resample_arc(xy: np.ndarray, n_points: int) -> np.ndarray:
    """Resample a 2-D polyline to ``n_points`` along cumulative arc length."""
    arc = np.concatenate(
        [[0], np.cumsum(np.linalg.norm(np.diff(xy, axis=0), axis=1))]
    )
    s_new = np.linspace(0, arc[-1], n_points)
    return np.stack(
        [np.interp(s_new, arc, xy[:, 0]), np.interp(s_new, arc, xy[:, 1])],
        axis=1,
    )


def build_timestamps(
    log_dt: np.ndarray,
    total_duration_ms: float,
    dt_clip: tuple[float, float] = (2.0, 150.0),
) -> np.ndarray:
    """Convert per-step log_dt + total duration to cumulative ms timestamps.

    Args:
        log_dt: (N,) array of natural-log step intervals.
        total_duration_ms: target total span. The output is scaled so the
            sum approximately matches this (modulo dt_clip).
        dt_clip: (min, max) per-step clamp in milliseconds.

    Returns:
        (N,) integer-rounded cumulative timestamps starting at 0,
        strictly increasing.
    """
    n = len(log_dt)
    dt_raw = np.clip(np.exp(log_dt), 0.0, None)
    dt_sum = dt_raw.sum()
    if dt_sum > 1e-6:
        scale = total_duration_ms / dt_sum
    else:
        scale = total_duration_ms / max(n, 1)
    dt_ms = np.clip(dt_raw * scale, dt_clip[0], dt_clip[1])

    t_abs = np.cumsum(dt_ms)
    t_abs = np.concatenate([[0.0], t_abs[:-1]])

    for i in range(1, n):
        if t_abs[i] <= t_abs[i - 1]:
            t_abs[i] = t_abs[i - 1] + 1.0
    return t_abs

• Step 4: Run

uv run pytest tests/unit/test_postprocess.py -v

• Step 5: Commit

git add src/ai_mouse/_postprocess.py tests/unit/test_postprocess.py
git commit -m "feat(lib): add resample_arc, build_timestamps"

---

Task 4.8: Add sample_duration + truncnorm_sample

Files:

• Modify: src/ai_mouse/_postprocess.py, tests/unit/test_postprocess.py

• Step 1: Tests

from ai_mouse._postprocess import sample_duration, truncnorm_sample


def test_truncnorm_sample_within_bounds() -> None:
    rng = np.random.default_rng(0)
    samples = [
        truncnorm_sample(80.0, 30.0, 20.0, 300.0, rng) for _ in range(500)
    ]
    arr = np.array(samples)
    assert arr.min() >= 20.0
    assert arr.max() <= 300.0
    # Mean roughly close to mu
    assert abs(arr.mean() - 80.0) < 5.0


def test_truncnorm_sample_far_outside_falls_back_to_clip() -> None:
    rng = np.random.default_rng(0)
    # mu far outside [low, high] — rejection will fail
    v = truncnorm_sample(mu=1000.0, sigma=1.0, low=20.0, high=30.0, rng=rng)
    assert 20.0 <= v <= 30.0


def test_sample_duration_uses_correct_bin() -> None:
    dist_dict = {
        "bins": [0, 50, 100, 200, 400, 600, 800, 1200, float("inf")],
        "params": [
            {"mu_log": 4.0, "sigma_log": 0.01},  # bin 0: dist < 50
            {"mu_log": 5.0, "sigma_log": 0.01},  # bin 1: 50 <= dist < 100
            {"mu_log": 6.0, "sigma_log": 0.01},  # bin 2: 100 <= dist < 200
        ] + [{"mu_log": 7.0, "sigma_log": 0.01}] * 5,
    }
    rng = np.random.default_rng(0)
    v = sample_duration(dist_dict, 150.0, rng)
    # exp(6) ~ 403, with tiny sigma we should land near there
    assert 350 < v < 460

• Step 2: Run, observe failure

uv run pytest tests/unit/test_postprocess.py -v

• Step 3: Implement

Append to _postprocess.py:

def truncnorm_sample(
    mu: float,
    sigma: float,
    low: float,
    high: float,
    rng: np.random.Generator,
    max_tries: int = 32,
) -> float:
    """Sample from N(mu, sigma) truncated to [low, high] via rejection.

    Falls back to clipping if rejection fails ``max_tries`` times.
    """
    for _ in range(max_tries):
        v = rng.normal(mu, sigma)
        if low <= v <= high:
            return float(v)
    return float(np.clip(rng.normal(mu, sigma), low, high))


def sample_duration(
    duration_dist: dict,
    dist: float,
    rng: np.random.Generator,
) -> float:
    """Sample total trajectory duration (ms) for the given pixel distance.

    Uses per-bin log-normal parameters in ``duration_dist``.
    """
    bins = duration_dist["bins"]
    params = duration_dist["params"]
    bin_idx = len(bins) - 1
    for i in range(len(bins) - 1):
        if dist < bins[i + 1]:
            bin_idx = i
            break
    bin_idx = min(bin_idx, len(params) - 1)
    mu_log = params[bin_idx]["mu_log"]
    sigma_log = params[bin_idx]["sigma_log"]
    return float(np.exp(rng.normal(mu_log, sigma_log)))

• Step 4: Test

uv run pytest tests/unit/test_postprocess.py -v

• Step 5: Commit

git add src/ai_mouse/_postprocess.py tests/unit/test_postprocess.py
git commit -m "feat(lib): add sample_duration, truncnorm_sample (no scipy)"

---

Task 4.9: Write mouse.py (MouseModel + _get_default_mouse_model)

Files:

• Create: src/ai_mouse/mouse.py

• Step 1: Write test scaffolding

Create tests/unit/test_mouse.py:

"""Tests for MouseModel and ai_mouse.generate()."""
from __future__ import annotations

import numpy as np
import pytest


def test_mouse_model_init_default() -> None:
    from ai_mouse.mouse import MouseModel
    m = MouseModel()
    assert m._seq_len > 0
    assert m._session is not None
    m.close()


def test_mouse_model_generate_returns_correct_shape() -> None:
    from ai_mouse.mouse import MouseModel
    m = MouseModel()
    pts = m.generate((100, 200), (900, 400))
    assert len(pts) == 66  # 64 moves + 2 clicks
    for x, y, t in pts:
        assert isinstance(x, int)
        assert isinstance(y, int)
        assert isinstance(t, int)


def test_mouse_model_click_false_omits_clicks() -> None:
    from ai_mouse.mouse import MouseModel
    m = MouseModel()
    pts = m.generate((100, 200), (900, 400), click=False)
    assert len(pts) == 64


def test_mouse_model_seed_reproducibility() -> None:
    from ai_mouse.mouse import MouseModel
    m = MouseModel()
    a = m.generate((100, 200), (900, 400), seed=42)
    b = m.generate((100, 200), (900, 400), seed=42)
    assert a == b


def test_mouse_model_invalid_path_raises_model_load_error() -> None:
    from ai_mouse.mouse import MouseModel
    from ai_mouse.errors import ModelLoadError
    with pytest.raises(ModelLoadError):
        MouseModel(model_path="/nonexistent/path/here")

• Step 2: Run, observe failure

uv run pytest tests/unit/test_mouse.py -v

Expected: ImportError.

• Step 3: Implement mouse.py

Create src/ai_mouse/mouse.py:

"""MouseModel — ONNX Runtime-backed mouse trajectory generation."""
from __future__ import annotations

import json
import math
from collections.abc import Sequence
from pathlib import Path
from typing import Optional

import numpy as np
import onnxruntime as ort

from ai_mouse._assets import resolve
from ai_mouse._coord import decode_trajectory
from ai_mouse._postprocess import (
    build_timestamps,
    enforce_forward_monotonic,
    gaussian_smooth,
    resample_arc,
    sample_duration,
    smooth_start,
    snap_endpoints,
    truncnorm_sample,
)
from ai_mouse.errors import GenerationError, ModelLoadError

_N_EULER_STEPS = 10


class MouseModel:
    """Persistent ONNX Runtime session for mouse trajectory generation.

    Construct once and reuse across calls — the underlying
    ``InferenceSession`` is created lazily in ``__init__`` and kept alive
    until :meth:`close` is called.
    """

    def __init__(
        self,
        model_path: str | Path | None = None,
        providers: Sequence[str] | None = None,
        seed: int | None = None,
    ) -> None:
        path_obj: Optional[Path] = Path(model_path) if model_path is not None else None

        onnx_path = resolve(path_obj, "flow_model.onnx")
        cfg_path = resolve(path_obj, "train_config.json")
        click_path = resolve(path_obj, "click_dist.json")
        dur_path = resolve(path_obj, "duration_dist.json")

        cfg = json.loads(cfg_path.read_text())
        self._seq_len = int(cfg["seq_len"])
        self._cond_dim = int(cfg.get("cond_dim", 3))
        self._click_params = json.loads(click_path.read_text())
        self._duration_dist = json.loads(dur_path.read_text())

        try:
            self._session = ort.InferenceSession(
                str(onnx_path),
                providers=list(providers) if providers else ["CPUExecutionProvider"],
            )
        except Exception as exc:
            raise ModelLoadError(f"Failed to load ONNX session: {exc}") from exc

        self._default_seed = seed
        self._rng = np.random.default_rng(seed)

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def generate(
        self,
        start: tuple[int, int],
        end: tuple[int, int],
        n_points: int = 64,
        speed: float | None = None,
        click: bool = True,
        seed: int | None = None,
    ) -> list[tuple[int, int, int]]:
        """Generate a human-like mouse trajectory from ``start`` to ``end``.

        Args:
            start: (x, y) starting pixel coordinate.
            end: (x, y) target pixel coordinate.
            n_points: number of move points (default 64).
            speed: optional multiplier; ``speed=2`` halves the duration.
            click: append mouse-down and mouse-up events at the end.
            seed: per-call seed overriding the instance seed.

        Returns:
            List of (x, y, t_ms) tuples. ``t_ms`` is cumulative from 0.
        """
        rng = np.random.default_rng(seed) if seed is not None else self._rng

        sx, sy = float(start[0]), float(start[1])
        ex, ey = float(end[0]), float(end[1])
        dist = max(math.hypot(ex - sx, ey - sy), 1.0)

        total_duration = sample_duration(self._duration_dist, dist, rng)
        if speed is not None and speed > 0:
            total_duration /= speed
        total_duration = max(total_duration, 10.0)

        cond = np.array(
            [
                dist / 2000.0,
                math.log(dist / 100.0),
                math.log(total_duration / 500.0),
            ],
            dtype=np.float32,
        )[None]

        # Euler ODE
        x = rng.standard_normal((1, self._seq_len, 3)).astype(np.float32)
        dt = 1.0 / _N_EULER_STEPS
        for step in range(_N_EULER_STEPS):
            t = np.full((1,), step * dt, dtype=np.float32)
            v = self._session.run(
                ["v"], {"x_t": x, "t": t, "cond": cond}
            )[0]
            x = x + v * dt

        if not np.all(np.isfinite(x)):
            raise GenerationError("Trajectory contains NaN/Inf after Euler integration")

        forward = x[0, :, 0].copy()
        lateral = x[0, :, 1].copy()
        log_dt = x[0, :, 2].copy()

        # Spatial post-processing
        forward, lateral = snap_endpoints(forward, lateral, self._seq_len)
        forward, lateral = smooth_start(forward, lateral)
        forward = enforce_forward_monotonic(forward)
        lateral = gaussian_smooth(lateral, sigma=1.0)

        # Temporal post-processing
        log_dt = np.clip(log_dt, 0.0, 5.0)
        log_dt[0] = 0.0

        # Decode to pixels
        normalised = np.stack([forward, lateral], axis=1)
        pixels = decode_trajectory(normalised, start, end)

        if n_points != self._seq_len:
            pixels = resample_arc(pixels, n_points)
            log_dt = np.interp(
                np.linspace(0, 1, n_points),
                np.linspace(0, 1, self._seq_len),
                log_dt,
            )

        ts = build_timestamps(log_dt, total_duration)

        moves: list[tuple[int, int, int]] = [
            (int(round(pixels[i, 0])), int(round(pixels[i, 1])), int(round(ts[i])))
            for i in range(n_points)
        ]
        if not click:
            return moves

        click_dur = int(
            truncnorm_sample(
                float(self._click_params["mu"]),
                float(self._click_params["sigma"]),
                float(self._click_params["low"]),
                float(self._click_params["high"]),
                rng,
            )
        )
        click_dur = max(click_dur, int(float(self._click_params["low"])))
        last_t = moves[-1][2]
        cx, cy = moves[-1][0], moves[-1][1]
        return moves + [(cx, cy, last_t), (cx, cy, last_t + click_dur)]

    def sample_click_duration_ms(self, seed: int | None = None) -> int:
        """Sample one click hold duration from the bundled distribution."""
        rng = np.random.default_rng(seed) if seed is not None else self._rng
        v = truncnorm_sample(
            float(self._click_params["mu"]),
            float(self._click_params["sigma"]),
            float(self._click_params["low"]),
            float(self._click_params["high"]),
            rng,
        )
        return max(int(v), int(float(self._click_params["low"])))

    def close(self) -> None:
        """Release the ONNX session."""
        self._session = None  # type: ignore[assignment]

    def __enter__(self) -> "MouseModel":
        return self

    def __exit__(self, *exc) -> None:
        self.close()

• Step 4: Run tests

uv run pytest tests/unit/test_mouse.py -v

Expected: 5 pass. If test_mouse_model_seed_reproducibility fails because the bundled ONNX model produces different results across two runs with the same seed, that's a bug in MouseModel. Verify the rng is properly seeded.

• Step 5: Commit

git add src/ai_mouse/mouse.py tests/unit/test_mouse.py
git commit -m "feat(lib): add MouseModel (numpy + ONNX Runtime)"

---

Task 4.10: Write scroll.py (ScrollModel)

Files:

• Create: src/ai_mouse/scroll.py

• Create: tests/unit/test_scroll.py

• Step 1: Write tests

Create tests/unit/test_scroll.py:

"""Tests for ScrollModel and ai_mouse.generate_scroll()."""
from __future__ import annotations

import pytest


def test_scroll_model_init_default() -> None:
    from ai_mouse.scroll import ScrollModel
    m = ScrollModel()
    assert m._seq_len > 0
    m.close()


def test_scroll_model_generate_target_mode() -> None:
    from ai_mouse.scroll import ScrollModel
    m = ScrollModel()
    events = m.generate(0, 1500, mode="target")
    assert len(events) >= 5
    total = sum(e["deltaY"] for e in events)
    # Should approach but not necessarily equal 1500 exactly
    assert 1000 <= total <= 2000  # broad — quantisation can drift
    assert events[0]["t"] == 0
    assert all(e["deltaMode"] == 0 for e in events)


def test_scroll_model_direction() -> None:
    from ai_mouse.scroll import ScrollModel
    m = ScrollModel()
    events = m.generate(2000, 0, mode="target")  # upward
    total = sum(e["deltaY"] for e in events)
    assert total < 0


def test_scroll_invalid_path() -> None:
    from ai_mouse.errors import ModelLoadError
    from ai_mouse.scroll import ScrollModel
    with pytest.raises(ModelLoadError):
        ScrollModel(model_path="/no/such/path")

• Step 2: Run, observe failure

uv run pytest tests/unit/test_scroll.py -v

• Step 3: Implement scroll.py

Create src/ai_mouse/scroll.py:

"""ScrollModel — ONNX Runtime-backed scroll event generation."""
from __future__ import annotations

import json
import math
from collections.abc import Sequence
from pathlib import Path
from typing import Literal, Optional

import numpy as np
import onnxruntime as ort

from ai_mouse._assets import resolve
from ai_mouse.errors import ModelLoadError

_DURATION_TABLE = {
    "fast":    lambda d: d * 0.2 + 100.0,
    "precise": lambda d: d * 1.5 + 300.0,
    "target":  lambda d: d * 0.4 + 200.0,
}

_QUANTUM = {"precise": 40, "fast": 120, "target": 120}


class ScrollModel:
    """Persistent ONNX Runtime session for scroll event generation."""

    def __init__(
        self,
        model_path: str | Path | None = None,
        providers: Sequence[str] | None = None,
        seed: int | None = None,
    ) -> None:
        path_obj: Optional[Path] = Path(model_path) if model_path is not None else None

        onnx_path = resolve(path_obj, "scroll_decoder.onnx")
        cfg_path = resolve(path_obj, "scroll_config.json")
        cfg = json.loads(cfg_path.read_text())

        self._seq_len = int(cfg["seq_len"])
        self._latent_dim = int(cfg["latent_dim"])
        self._cond_dim = int(cfg["cond_dim"])

        try:
            self._session = ort.InferenceSession(
                str(onnx_path),
                providers=list(providers) if providers else ["CPUExecutionProvider"],
            )
        except Exception as exc:
            raise ModelLoadError(f"Failed to load scroll ONNX session: {exc}") from exc

        self._rng = np.random.default_rng(seed)

    def generate(
        self,
        start_scroll_y: int,
        target_scroll_y: int,
        mode: Literal["target", "fast", "precise"] = "target",
        viewport_height: int = 800,
        seed: int | None = None,
    ) -> list[dict]:
        """Generate a sequence of mouse-wheel events.

        Returns a list of ``{"deltaY": int, "deltaMode": 0, "t": int}``
        dicts. Positive ``deltaY`` = scroll down.
        """
        rng = np.random.default_rng(seed) if seed is not None else self._rng

        distance = abs(target_scroll_y - start_scroll_y)
        direction = 1 if target_scroll_y > start_scroll_y else -1
        distance = max(distance, 10)

        cond = self._build_condition(float(distance), direction, mode, viewport_height)
        z = rng.standard_normal((1, self._latent_dim)).astype(np.float32)
        decoded = self._session.run(["seq"], {"z": z, "cond": cond[None]})[0][0]

        delta_norm = decoded[:, 0]
        log_dt = decoded[:, 1]

        # Softmax-like normalisation; scale to target distance
        delta_weights = np.exp(delta_norm)
        delta_weights /= delta_weights.sum()
        delta_px = delta_weights * distance * direction

        quantum = _QUANTUM[mode]
        delta_q = np.round(delta_px / quantum) * quantum
        for i in range(len(delta_q)):
            if delta_q[i] == 0:
                delta_q[i] = quantum * direction
        # Adjust last event so total matches target distance
        delta_q[-1] += (distance * direction) - delta_q.sum()

        # Timestamp building
        if len(log_dt) > 3:
            median_log = float(np.median(log_dt))
            log_dt[:2] = np.clip(log_dt[:2], None, median_log + 0.5)
            log_dt[-2:] = np.clip(log_dt[-2:], None, median_log + 0.5)
        dt_ms = np.clip(np.exp(log_dt), 5, 80)
        expected = _DURATION_TABLE[mode](distance)
        dt_ms = np.clip(dt_ms * (expected / max(dt_ms.sum(), 1.0)), 5, 80)

        t_abs = np.cumsum(dt_ms).astype(int)
        t_abs = np.concatenate([[0], t_abs[:-1]])
        for i in range(1, len(t_abs)):
            if t_abs[i] <= t_abs[i - 1]:
                t_abs[i] = t_abs[i - 1] + 5

        events: list[dict] = []
        for i in range(self._seq_len):
            dy = int(delta_q[i])
            if dy != 0 or len(events) < 5:
                events.append({"deltaY": dy, "deltaMode": 0, "t": int(t_abs[i])})
        return events

    def _build_condition(
        self,
        distance: float,
        direction: int,
        mode: str,
        viewport_height: int,
    ) -> np.ndarray:
        mode_onehot = [0.0, 0.0, 0.0]
        if mode == "target":
            mode_onehot[0] = 1.0
        elif mode == "fast":
            mode_onehot[1] = 1.0
        elif mode == "precise":
            mode_onehot[2] = 1.0
        return np.array(
            [
                distance / 5000.0,
                math.log(max(distance, 1.0) / 500.0),
                float(direction),
                viewport_height / 1000.0,
                *mode_onehot,
            ],
            dtype=np.float32,
        )

    def close(self) -> None:
        self._session = None  # type: ignore[assignment]

    def __enter__(self) -> "ScrollModel":
        return self

    def __exit__(self, *exc) -> None:
        self.close()

• Step 4: Run tests

uv run pytest tests/unit/test_scroll.py -v

Expected: 4 pass.

• Step 5: Commit

git add src/ai_mouse/scroll.py tests/unit/test_scroll.py
git commit -m "feat(lib): add ScrollModel (numpy + ONNX Runtime)"

---

Task 4.11: Rewrite __init__.py with cached singleton functions

Files:

• Modify: src/ai_mouse/__init__.py

• Step 1: Write tests for the public surface

Create tests/unit/test_public_api.py:

"""Tests for the public package-level API."""
from __future__ import annotations


def test_public_symbols_importable() -> None:
    from ai_mouse import (
        MouseModel,
        ScrollModel,
        generate,
        generate_scroll,
        errors,
    )
    assert MouseModel is not None
    assert ScrollModel is not None
    assert callable(generate)
    assert callable(generate_scroll)
    assert hasattr(errors, "ModelLoadError")


def test_generate_function_returns_list_of_tuples() -> None:
    from ai_mouse import generate
    pts = generate((100, 100), (300, 200))
    assert isinstance(pts, list)
    assert len(pts) > 0
    assert isinstance(pts[0], tuple)
    assert len(pts[0]) == 3


def test_generate_singleton_reused() -> None:
    from ai_mouse import generate
    from ai_mouse import _model_cache
    _model_cache._get_mouse_model.cache_clear()
    generate((0, 0), (100, 100))
    info_after_first = _model_cache._get_mouse_model.cache_info()
    generate((0, 0), (200, 200))
    info_after_second = _model_cache._get_mouse_model.cache_info()
    assert info_after_second.hits > info_after_first.hits


def test_version_present() -> None:
    import ai_mouse
    assert hasattr(ai_mouse, "__version__")
    assert isinstance(ai_mouse.__version__, str)

• Step 2: Run, observe failure

uv run pytest tests/unit/test_public_api.py -v

• Step 3: Create _model_cache.py

Create src/ai_mouse/_model_cache.py:

"""Process-level lru_cache for default MouseModel / ScrollModel instances."""
from __future__ import annotations

from collections.abc import Sequence
from functools import lru_cache
from pathlib import Path

from ai_mouse.mouse import MouseModel
from ai_mouse.scroll import ScrollModel


@lru_cache(maxsize=4)
def _get_mouse_model(
    model_key: str,
    providers_key: tuple[str, ...],
) -> MouseModel:
    path = None if model_key == "__bundled__" else Path(model_key)
    providers = list(providers_key) if providers_key else None
    return MouseModel(model_path=path, providers=providers)


@lru_cache(maxsize=4)
def _get_scroll_model(
    model_key: str,
    providers_key: tuple[str, ...],
) -> ScrollModel:
    path = None if model_key == "__bundled__" else Path(model_key)
    providers = list(providers_key) if providers_key else None
    return ScrollModel(model_path=path, providers=providers)


def get_mouse_model(
    model_path: str | Path | None,
    providers: Sequence[str] | None,
) -> MouseModel:
    key = "__bundled__" if model_path is None else str(model_path)
    return _get_mouse_model(key, tuple(providers or ()))


def get_scroll_model(
    model_path: str | Path | None,
    providers: Sequence[str] | None,
) -> ScrollModel:
    key = "__bundled__" if model_path is None else str(model_path)
    return _get_scroll_model(key, tuple(providers or ()))

• Step 4: Rewrite __init__.py

Replace src/ai_mouse/__init__.py entirely:

"""ai_mouse — ONNX Runtime SDK for human-like mouse trajectories and scroll events.

Public API:

    from ai_mouse import generate, generate_scroll, MouseModel, ScrollModel

See https://github.com/<owner>/ai_mouse for usage examples.
"""
from __future__ import annotations

from collections.abc import Sequence
from pathlib import Path
from typing import Literal

from ai_mouse import errors
from ai_mouse._model_cache import get_mouse_model, get_scroll_model
from ai_mouse.mouse import MouseModel
from ai_mouse.scroll import ScrollModel

__version__ = "0.2.0"

__all__ = [
    "MouseModel",
    "ScrollModel",
    "errors",
    "generate",
    "generate_scroll",
    "__version__",
]


def generate(
    start: tuple[int, int],
    end: tuple[int, int],
    *,
    n_points: int = 64,
    speed: float | None = None,
    click: bool = True,
    seed: int | None = None,
    model_path: str | Path | None = None,
    providers: Sequence[str] | None = None,
) -> list[tuple[int, int, int]]:
    """Generate a human-like mouse trajectory.

    See :class:`MouseModel.generate` for argument semantics.
    The underlying :class:`MouseModel` is cached process-wide; repeat
    calls with the same ``(model_path, providers)`` reuse the session.
    """
    model = get_mouse_model(model_path, providers)
    return model.generate(
        start=start,
        end=end,
        n_points=n_points,
        speed=speed,
        click=click,
        seed=seed,
    )


def generate_scroll(
    start_scroll_y: int,
    target_scroll_y: int,
    *,
    mode: Literal["target", "fast", "precise"] = "target",
    viewport_height: int = 800,
    seed: int | None = None,
    model_path: str | Path | None = None,
    providers: Sequence[str] | None = None,
) -> list[dict]:
    """Generate a sequence of mouse-wheel events. See :class:`ScrollModel.generate`."""
    model = get_scroll_model(model_path, providers)
    return model.generate(
        start_scroll_y=start_scroll_y,
        target_scroll_y=target_scroll_y,
        mode=mode,
        viewport_height=viewport_height,
        seed=seed,
    )

• Step 5: Run all unit tests

uv run pytest tests/unit -v

Expected: all pass.

• Step 6: Commit

git add src/ai_mouse/__init__.py src/ai_mouse/_model_cache.py tests/unit/test_public_api.py
git commit -m "feat(lib): rewrite __init__.py with cached singleton entrypoints"

---

Task 4.12: Add py.typed marker

Files:

• Create: src/ai_mouse/py.typed

• Step 1: Create marker file

touch src/ai_mouse/py.typed

• Step 2: Commit

git add src/ai_mouse/py.typed
git commit -m "feat(lib): add py.typed marker (PEP 561)"

---

Task 4.13: Add golden regression tests

Files:

• Create: tests/unit/test_golden.py

• Step 1: Write the test

Create tests/unit/test_golden.py:

"""Golden regression tests — lock library output against pre-migration captures.

Tolerance: pixels and ms allowed ±2 due to ORT/PyTorch fp accumulation
and rounding differences. Update goldens only via an explicit recapture.
"""
from __future__ import annotations

from pathlib import Path

import numpy as np
import pytest

from ai_mouse import generate, generate_scroll

_GOLDEN_DIR = Path(__file__).parent / "data"

_MOUSE_CASES: list[tuple[tuple[int, int], tuple[int, int]]] = [
    ((100, 200), (900, 400)),
    ((500, 500), (500, 100)),
    ((200, 600), (800, 200)),
    ((100, 100), (130, 110)),
    ((50, 50),   (1500, 900)),
    ((400, 300), (500, 300)),
    ((300, 300), (700, 700)),
    ((600, 400), (200, 100)),
]

_SCROLL_CASES: list[tuple[int, int, str]] = [
    (0,    1500, "target"),
    (0,    500,  "precise"),
    (0,    5000, "fast"),
    (2000, 0,    "target"),
    (0,    800,  "precise"),
    (0,    3500, "fast"),
    (1000, 1200, "precise"),
    (0,    10000, "fast"),
]


@pytest.mark.parametrize("case_idx", range(8))
@pytest.mark.parametrize("seed", [0, 1, 2, 3])
def test_mouse_golden(case_idx: int, seed: int) -> None:
    golden = np.load(_GOLDEN_DIR / "golden_mouse.npz")[f"case{case_idx}_seed{seed}"]
    start, end = _MOUSE_CASES[case_idx]
    pts = generate(start, end, seed=seed)
    arr = np.array(pts, dtype=np.int64)
    assert arr.shape == golden.shape, f"shape mismatch: {arr.shape} vs {golden.shape}"
    diff = np.abs(arr - golden)
    assert diff.max() <= 2, (
        f"case{case_idx} seed{seed} max diff {diff.max()} > 2; "
        f"first diff row: arr[?]=..., golden[?]=..."
    )


@pytest.mark.parametrize("case_idx", range(8))
@pytest.mark.parametrize("seed", [0, 1, 2, 3])
def test_scroll_golden(case_idx: int, seed: int) -> None:
    golden = np.load(_GOLDEN_DIR / "golden_scroll.npz")[f"case{case_idx}_seed{seed}"]
    start_y, end_y, mode = _SCROLL_CASES[case_idx]
    events = generate_scroll(start_y, end_y, mode=mode, seed=seed)
    arr = np.array(
        [[e["deltaY"], e["deltaMode"], e["t"]] for e in events],
        dtype=np.int64,
    )
    # Scroll uses VAE prior sampling — looser tolerance.
    # Allow ±1 wheel quantum (40 or 120 px) for deltaY; ±10 ms for t.
    quantum = 120 if mode != "precise" else 40
    if arr.shape != golden.shape:
        pytest.skip(
            f"event count diverged: {arr.shape[0]} vs {golden.shape[0]} "
            f"(quantisation boundary sensitivity)"
        )
    delta_diff = np.abs(arr[:, 0] - golden[:, 0])
    t_diff = np.abs(arr[:, 2] - golden[:, 2])
    assert delta_diff.max() <= quantum, f"deltaY diverged > 1 quantum"
    assert t_diff.max() <= 20, f"t diverged > 20ms"

• Step 2: Run goldens

uv run pytest tests/unit/test_golden.py -v

Expected outcomes:

• 32 mouse golden cases run; some failures are expected because the post-migration randomness differs from torch (different RNG instance, different floating-point path). Inspect failures.
• If max diff is large (>10), there's a real bug — investigate.
• If max diff is in the 3-8 range, bump the tolerance in the test (from 2 to a value that lets all pass) with a comment explaining why, then re-commit.

This is the moment of truth for the migration: a passing golden suite says the rewrite preserved semantics.

• Step 3: Decide on tolerance

If you needed to widen the tolerance, edit test_golden.py and document it. For example, if max diff observed is 5, change assert diff.max() <= 2 to assert diff.max() <= 6, ... with a comment:

# Tolerance 6: ORT/PyTorch numeric path differs slightly; observed max diff 5.
assert diff.max() <= 6, (
    ...
)

• Step 4: Commit

git add tests/unit/test_golden.py
git commit -m "test(lib): add golden regression suite for mouse + scroll"

---

Task 4.14: Delete obsolete legacy modules

Files:

• Delete: src/ai_mouse/generator.py

• Delete: src/ai_mouse/scroll/generator.py and src/ai_mouse/scroll/__init__.py

• Delete: src/ai_mouse/scroll/ directory entirely (replaced by src/ai_mouse/scroll.py)

• Delete: src/ai_mouse/coord.py (replaced by _coord.py)

• Delete: any remaining files in src/ai_mouse/ not in the spec's final layout

• Move (clean up): scripts/build_golden_*.py → can be deleted now that goldens are captured

• Step 1: Check current state of src/ai_mouse/

ls src/ai_mouse/
ls src/ai_mouse/scroll/ 2>/dev/null

Expected at this point: a mix of new files (__init__.py, mouse.py, scroll.py, _coord.py, _postprocess.py, _assets.py, _model_cache.py, errors.py, py.typed, assets/) and leftover legacy (generator.py, coord.py, scroll/).

• Step 2: Delete legacy files

git rm src/ai_mouse/generator.py
git rm src/ai_mouse/coord.py
git rm -r src/ai_mouse/scroll/

• Step 3: Delete temporary scripts

git rm scripts/build_golden_mouse.py scripts/build_golden_scroll.py
rmdir scripts/ 2>/dev/null  # only succeeds if empty

• Step 4: Verify package layout

ls src/ai_mouse/

Expected:

__init__.py
_assets.py
_coord.py
_model_cache.py
_postprocess.py
errors.py
mouse.py
py.typed
scroll.py
assets/

(scroll/ directory removed; replaced by scroll.py module.)

• Step 5: Run full library test suite

uv run pytest tests/unit -v

Expected: all pass.

• Step 6: Verify tools/ still works (it now imports from src/ai_mouse private modules)

The tools-side trainer's import of from ai_mouse.coord import encode_trajectory will break (we deleted that file). Fix:

grep -rn "from ai_mouse.coord" tools/ --include="*.py"

For each hit, replace with from ai_mouse._coord import .... The spec explicitly allows tools/ to depend on ai_mouse._* private modules.

• Step 7: Run tools tests

uv run pytest tests/tools -v

Expected: all pass.

• Step 8: Commit

git add -A
git commit -m "refactor(lib): remove legacy generator.py / coord.py / scroll/ package"

---

Task 4.15: Verify clean install has no torch

Files:

• (verification only)

• Step 1: Build fresh wheel

uv build

• Step 2: Install into a clean venv with NO torch

uv venv .venv-test
.venv-test/Scripts/python -m pip install dist/ai_mouse-0.2.0-py3-none-any.whl

• Step 3: Smoke test

.venv-test/Scripts/python -c "
from ai_mouse import generate
pts = generate((100, 200), (900, 400), seed=0)
print(f'Got {len(pts)} events')
print('First 3:', pts[:3])
print('Last 2 (clicks):', pts[-2:])
print('No torch needed!')
"

Expected: prints output without ImportError. Verifies the "pure inference SDK" promise.

• Step 4: Confirm torch absent

.venv-test/Scripts/python -c "import torch" 2>&1 | head -1

Expected: ModuleNotFoundError: No module named 'torch'

• Step 5: Clean up

rm -rf .venv-test dist/

• Step 6: No commit needed (verification only — but if anything failed, fix forward)

---

Phase 5: Docs + cleanup

Task 5.1: Rewrite README.md

Files:

• Modify: README.md (overwrite — if it exists; create if not)

• Step 1: Check current README

ls README.md 2>/dev/null && head -20 README.md

• Step 2: Write the new README

Create/overwrite README.md:

# ai_mouse

Human-like mouse trajectory and scroll wheel event generator. Inference runs on
ONNX Runtime; the only runtime dependencies are `numpy` and `onnxruntime`.

## Install

```bash
pip install git+https://github.com/<owner>/ai_mouse.git
```

For GPU inference (optional), replace `onnxruntime` with the GPU variant:

```bash
pip install onnxruntime-gpu      # CUDA / TensorRT
# or
pip install onnxruntime-directml # Windows DirectML
```

## Quick start

### Mouse trajectory

```python
from ai_mouse import generate

points = generate(start=(100, 200), end=(900, 400))
# [(x, y, t_ms), ..., (cx, cy, t_down), (cx, cy, t_up)]
```

### Scroll wheel

```python
from ai_mouse import generate_scroll

events = generate_scroll(start_scroll_y=0, target_scroll_y=2000)
# [{"deltaY": 120, "deltaMode": 0, "t": 32}, ...]
```

### Class API (recommended for repeated calls)

```python
from ai_mouse import MouseModel

m = MouseModel()              # session created once
for target in target_list:
    pts = m.generate((cx, cy), target)
```

### Custom providers / GPU

```python
from ai_mouse import MouseModel

m = MouseModel(providers=["CUDAExecutionProvider", "CPUExecutionProvider"])
# or
m = MouseModel(providers=["DmlExecutionProvider"])
```

### Reproducibility

```python
m.generate(start, end, seed=42)
```

## API summary

| Name | Purpose |
|---|---|
| `generate(start, end, *, n_points=64, speed=None, click=True, seed=None)` | One-shot call; internal lru_cache singleton |
| `MouseModel(model_path=None, providers=None, seed=None)` | Persistent session |
| `generate_scroll(...)` / `ScrollModel(...)` | Same shape for scroll |
| `ai_mouse.errors.{ModelLoadError, GenerationError}` | Exception hierarchy |

## Thread safety

`MouseModel.generate` and `ScrollModel.generate` are safe to call concurrently
from multiple threads — ORT `InferenceSession` is itself thread-safe.

## Development

The repo contains optional dev-only tooling under `tools/` for training your
own models, running the FastAPI web UI, and evaluating output quality. Install
with the `dev` group:

```bash
uv sync --group dev
```

Common commands:

```bash
# Web UI (collect + train + verify in browser)
uv run python tools/serve.py

# Training (after collecting your own data)
uv run python -m tools train --data data/traces.jsonl --output data/models_v2

# Convert Balabit corpus to trace format
uv run python -m tools balabit-adapter --input data/balabit_raw \
    --output data/pretrain_traces.jsonl

# Eval report
uv run python -m tools eval --model-dir data/models_v2 \
    --reference data/pretrain_traces.jsonl --output data/eval_reports/report.md

# Re-export ONNX after retraining
uv run python tools/export_onnx.py --flow-ckpt data/models_v2 \
    --scroll-ckpt data/scroll_models --output src/ai_mouse/assets/
```

Tests:

```bash
uv run pytest tests/unit    # library-only (no torch)
uv run pytest tests/tools   # full dev suite
```

• Step 3: Commit

git add README.md
git commit -m "docs: rewrite README from SDK-consumer perspective"

---

Task 5.2: Create CHANGELOG.md

Files:

• Create: CHANGELOG.md

• Step 1: Write CHANGELOG

Create CHANGELOG.md:

# Changelog

All notable changes to this project will be documented here. Format follows
[Keep a Changelog](https://keepachangelog.com/en/1.1.0/); versioning follows
[Semantic Versioning](https://semver.org/).

## [0.2.0] - 2026-05-11

### Changed (breaking)

- Inference no longer requires PyTorch. Runtime dependencies are now
  `numpy + onnxruntime` only.
- Public API additions: `MouseModel` and `ScrollModel` classes wrapping a
  persistent ORT `InferenceSession`.
- Function signatures `generate()` and `generate_scroll()` are now keyword-only
  past the positional `start`/`end` (or `start_scroll_y`/`target_scroll_y`).
- New parameters: `click=True` (mouse), `seed=` (both), `viewport_height=` (scroll).
- Removed `config=` parameter; use kwargs directly.
- `model_dir=` renamed to `model_path=`; accepts `str` or `pathlib.Path`.
- Training, web UI, collector, eval, and data adapter code moved to repo-level
  `tools/`; no longer packaged in the wheel.

### Added

- ONNX-format pre-trained weights bundled inside the wheel via
  `importlib.resources` (~3 MB).
- `tools/export_onnx.py` script with PyTorch/ORT parity check.
- Errors namespace `ai_mouse.errors` with `AiMouseError`, `ModelLoadError`,
  `GenerationError`.
- Custom ORT providers parameter for GPU / DirectML.
- Per-process `lru_cache` so `generate()` / `generate_scroll()` reuse the
  default model across calls.

### Removed

- Legacy `JointCVAE` class.
- `ai_mouse.config.GenerateConfig` top-level export (parameters moved to kwargs).
- Source dependency on `scipy.stats.truncnorm` (replaced by numpy rejection sampling).

• Step 2: Commit

git add CHANGELOG.md
git commit -m "docs: add CHANGELOG with 0.2.0 entry"

---

Task 5.3: Create examples/quickstart.py

Files:

• Create: examples/quickstart.py

• Step 1: Create the example

mkdir -p examples

Create examples/quickstart.py:

"""Minimal example: generate one trajectory + click event.

Run:    uv run python examples/quickstart.py
"""
from __future__ import annotations

from ai_mouse import generate

start = (100, 200)
end = (900, 400)
points = generate(start, end, seed=0)

print(f"Generated {len(points)} events:")
print(f"  first move:  {points[0]}")
print(f"  middle move: {points[len(points) // 2]}")
print(f"  last move:   {points[-3]}")
print(f"  click-down:  {points[-2]}")
print(f"  click-up:    {points[-1]}")

# Typical replay loop pattern. t_ms is cumulative from the start of the trace,
# so block your sender thread until time-since-start reaches each event's t_ms.
#
# import time
# t0 = time.monotonic()
# for x, y, t_ms in points:
#     target_wallclock = t0 + t_ms / 1000.0
#     while time.monotonic() < target_wallclock:
#         pass
#     # replace this with pyautogui / pynput / win32 mouse_event:
#     # send_mouse_move(x, y)

• Step 2: Run it

uv run python examples/quickstart.py

Expected: prints 5 event lines.

• Step 3: Commit

git add examples/quickstart.py
git commit -m "docs: add examples/quickstart.py"

---

Task 5.4: Update CLAUDE.md

Files:

• Modify: CLAUDE.md

• Step 1: Read current CLAUDE.md

cat CLAUDE.md

• Step 2: Rewrite to match new layout

Replace CLAUDE.md with content that reflects the new structure. Key changes:

• All python -m ai_mouse <cmd> references → python -m tools <cmd>
• "Bundled weights live in data/models_v2/" → "Bundled weights live in src/ai_mouse/assets/"
• Add a "Library vs tools boundary" section: library code in src/ai_mouse/ MUST NOT import torch; training code in tools/ may import library private modules
• Test commands split: pytest tests/unit vs pytest tests/tools
• main.py reference → tools/serve.py

Suggested new content:

# CLAUDE.md

This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.

## Project

`ai_mouse` is an ONNX-Runtime SDK that generates human-like mouse trajectories
and scroll wheel events. Runtime dependencies are `numpy + onnxruntime` only;
training and the FastAPI web UI live under `tools/` and are not packaged.

Package manager: **uv**, Python 3.12-3.13.

## Library vs tools — hard boundary

- **`src/ai_mouse/`** — wheel content. NEVER add a `import torch` /
  `import fastapi` / `import scipy` / `import matplotlib` here. CI's
  `library` job installs only runtime deps and would break.
- **`tools/`** — repo-only dev code (training, server, collector, eval,
  data adapters, ONNX export). May `import` library private modules
  (`ai_mouse._coord`, `ai_mouse._postprocess`) freely — they co-evolve.
- **Bundled assets**: `src/ai_mouse/assets/{flow_model,scroll_decoder}.onnx`
  plus four JSON metadata files. Re-generated by
  `tools/export_onnx.py` after retraining.

## Commands

```bash
# Web UI (collect + train + verify in browser)
uv run python tools/serve.py

# Tools CLI dispatch
uv run python -m tools train --data data/traces.jsonl --output data/models_v2
uv run python -m tools eval --model-dir data/models_v2 \
    --reference data/pretrain_traces.jsonl --output data/eval_reports/r.md
uv run python -m tools balabit-adapter --input data/balabit_raw \
    --output data/pretrain_traces.jsonl

# Re-export ONNX (after retraining)
uv run python tools/export_onnx.py --flow-ckpt data/models_v2 \
    --scroll-ckpt data/scroll_models --output src/ai_mouse/assets/

# Tests
uv run pytest tests/unit            # library-only (no torch)
uv run pytest tests/tools           # full dev suite (needs [dev] group)
uv run pytest tests/unit/test_mouse.py::test_mouse_model_seed_reproducibility

# Dependency sync
uv sync                              # runtime only
uv sync --group dev                  # dev-everything

Architecture

Two parallel ML subsystems share a collect → train → export → serve flow.

Mouse trajectories (src/ai_mouse/mouse.py library; tools/trainer.py training)

• Model: TrajectoryFlowModel (Conditional Flow Matching with 4-layer pre-norm Transformer, d_model=128, defined in tools/models.py)
• Inference: 10-step Euler ODE in Python; each step runs session.run(...) on src/ai_mouse/assets/flow_model.onnx. Followed by numpy post-processing in _postprocess.py (endpoint snapping, forward monotonicity, gaussian smoothing, log_dt → cumulative timestamps, truncated-normal click duration).
• Rotated coordinate frame (_coord.py): trajectories normalised so start → (0, 0), end → (1, 0). Makes the model angle/distance invariant.

Scroll wheel (src/ai_mouse/scroll.py; tools/scroll/trainer.py)

• Model: ScrollCVAE (bidirectional-GRU encoder + GRU decoder VAE, tools/scroll/models.py). Only the decoder is exported to ONNX (scroll_decoder.onnx); encoder is training-only.
• Inference: sample z ~ N(0, 1) in numpy → one session.run(...) → softmax-normalise deltas → quantise (40 px precise / 120 px otherwise) → scale to target distance → cumulative timestamps.

Server (tools/server/) and frontend (static/)

Unchanged from before. App factory create_app() mounts four routers under /api. Frontend is vanilla Vue 3 + axios + ECharts via CDN.

Config

tools/config.py holds the training-side dataclasses (TrainConfig, ScrollTrainConfig, etc.). The library does NOT use these — its only "configuration" is what's embedded in src/ai_mouse/assets/train_config.json (architecture params needed to know seq_len etc. at inference time).

Tests

• tests/unit/conftest.py — fixtures for library-only tests, no torch.
• tests/tools/conftest.py — model_dir and scroll_model_dir fixtures that produce untrained torch weights in a temp dir. Used by training- /server-side tests.
• tests/unit/test_golden.py — regression suite that pins library output against tests/unit/data/golden_{mouse,scroll}.npz captured before the ONNX migration. Tolerance: ±2 pixels/ms for mouse, ±1 quantum for scroll.

Server tests use httpx.ASGITransport(app=create_app()) with pytest-asyncio — no live socket.

- [ ] **Step 3: Commit**

```bash
git add CLAUDE.md
git commit -m "docs: update CLAUDE.md for new src/tools layout"

---

Task 5.5: Add GitHub Actions CI

Files:

• Create: .github/workflows/ci.yml

• Step 1: Create directory

mkdir -p .github/workflows

• Step 2: Write the workflow

Create .github/workflows/ci.yml:

name: CI

on:
  push:
    branches: [main]
  pull_request:
    branches: [main]

jobs:
  library:
    name: Library tests (no torch)
    runs-on: ${{ matrix.os }}
    strategy:
      fail-fast: false
      matrix:
        os: [ubuntu-latest, windows-latest]
        python: ["3.12", "3.13"]
    steps:
      - uses: actions/checkout@v4
      - uses: astral-sh/setup-uv@v3
      - run: uv venv --python ${{ matrix.python }}
      - run: uv pip install -e . pytest
      - run: uv run pytest tests/unit -v

  dev:
    name: Full dev suite (with torch)
    runs-on: ${{ matrix.os }}
    strategy:
      fail-fast: false
      matrix:
        os: [ubuntu-latest, windows-latest]
        python: ["3.12", "3.13"]
    steps:
      - uses: actions/checkout@v4
      - uses: astral-sh/setup-uv@v3
      - run: uv sync --group dev --python ${{ matrix.python }}
      - run: uv run pytest tests/ -v

• Step 3: Commit

git add .github/workflows/ci.yml
git commit -m "ci: add GitHub Actions workflow (library + dev jobs)"

---

Task 5.6: Delete remaining legacy artefacts

Files:

• Delete: legacy JointCVAE class (it's in tools/models.py now; spec says delete it)

• Delete: any leftover bundled-models path string referencing data/models_v2/ in the library

• Step 1: Check if JointCVAE is referenced anywhere

grep -rn "JointCVAE" --include="*.py"

If any tools/ file still references it (e.g., tools/models.py exports it), remove the class definition and the export.

• Step 2: Edit tools/models.py

Open tools/models.py, find the class JointCVAE block, delete it. Also delete the legacy from torch.distributions import Normal import if it's only used there.

• Step 3: Verify tools still pass

uv run pytest tests/tools -v

Expected: all pass.

• Step 4: Commit

git add tools/models.py
git commit -m "chore: remove legacy JointCVAE"

---

Task 5.7: Final verification — full sweep

• Step 1: Clean rebuild + install

uv venv .venv-final
.venv-final/Scripts/python -m pip install -e .
.venv-final/Scripts/python -m pip install pytest
.venv-final/Scripts/python -m pytest tests/unit -v

Expected: all unit tests pass; no torch installed in this venv.

• Step 2: Run dev suite separately

uv sync --group dev
uv run pytest tests/ -v

Expected: all tests pass.

• Step 3: Build wheel and inspect contents

uv build
unzip -l dist/ai_mouse-0.2.0-py3-none-any.whl | grep -v "^\$"

Expected file list (rough):

• ai_mouse/__init__.py
• ai_mouse/_assets.py, _coord.py, _model_cache.py, _postprocess.py
• ai_mouse/errors.py, mouse.py, scroll.py, py.typed
• ai_mouse/assets/flow_model.onnx
• ai_mouse/assets/scroll_decoder.onnx
• ai_mouse/assets/{click_dist,duration_dist,train_config,scroll_config}.json
• ai_mouse-0.2.0.dist-info/{METADATA,RECORD,WHEEL}

No tools/, tests/, data/, static/, or docs/.

• Step 4: Smoke test the wheel

uv venv .venv-wheel
.venv-wheel/Scripts/python -m pip install dist/ai_mouse-0.2.0-py3-none-any.whl
.venv-wheel/Scripts/python examples/quickstart.py

Expected: prints 5 event lines without error.

• Step 5: Clean up

rm -rf .venv-final .venv-wheel dist/

• Step 6: Document the final state

The migration is complete. Update PR description / branch message with:

ai_mouse 0.2.0 refactor complete:
- src/ai_mouse/ ships only numpy + ONNX Runtime runtime deps
- tools/ holds all training/server/eval code; not packaged
- 3 MB wheel includes ONNX weights for both mouse and scroll
- Golden regression suite locks behavior across the migration
- README, CHANGELOG, CLAUDE.md updated

No commit needed unless you adjust docs further.

---

Self-Review Notes

(Performed during plan authoring per writing-plans skill instructions.)

Spec coverage check:

• §1 Public API ⇒ Tasks 4.9, 4.10, 4.11
• §2 ONNX export ⇒ Tasks 3.1–3.5
• §3 NumPy rewrite ⇒ Tasks 4.4–4.8 (each postprocess fn) + 4.9, 4.10 (using them)
• §4 Migration phasing ⇒ Phases 1, 2 match the spec's 5 stages
• §5 Test strategy ⇒ Golden capture in Phase 0; per-fn unit tests in Phase 4; golden regression in 4.13; ONNX parity in 3.5; CI in 5.5
• §6 Documentation ⇒ Tasks 5.1–5.4 + examples/quickstart.py in 5.3

Placeholder scan: No "TBD", no unspecified test code, no "similar to Task N" shortcuts. Every code block is self-contained.

Type consistency: Function signatures in _postprocess.py referenced across Tasks 4.5–4.10 use consistent names (snap_endpoints, smooth_start, enforce_forward_monotonic, gaussian_smooth, resample_arc, build_timestamps, sample_duration, truncnorm_sample). MouseModel and ScrollModel constructor signatures match the spec verbatim. _get_mouse_model in _model_cache.py is used by Task 4.11 with matching signature.

Scope: This is one cohesive refactor — five phases with clear hand-offs but one logical goal. Not splittable into independent plans.