@10et/cli 1.0.0

package/scripts/train/v2/generate_balanced.py

@@ -0,0 +1,439 @@
+"""
+Generate balanced synthetic training data without LLM API calls.
+
+Uses handcrafted templates to fill class imbalance gaps.
+Each tool gets realistic (state, goal) pairs based on actual JFL agent patterns.
+
+Targets: tools with <50 examples in training buffer need augmentation.
+"""
+
+import json
+import os
+import sys
+import random
+import argparse
+from pathlib import Path
+
+
+# ============================================================================
+# Templates per tool — realistic state/goal combinations
+# ============================================================================
+
+TEMPLATES = {
+    "fix_bug": {
+        "agents": ["error-fixer", "hub-sentinel", "test-coverage", "system-health"],
+        "goals": [
+            "Fix TypeError in session-manager.ts — null reference on line 142",
+            "Hub process crashes with SIGTERM when 3+ agents connect simultaneously",
+            "Race condition in worktree allocation causes duplicate locks",
+            "Training buffer append fails silently when disk is full",
+            "Eval score regression after config change — investigate root cause",
+            "WebSocket connection drops after 60s idle — implement keepalive",
+            "PolicyHead inference returns NaN when state embedding has zero vector",
+            "Git merge conflict in auto-generated files blocks CI pipeline",
+            "Agent stranding when hub restarts — sessions not recovered",
+            "Memory leak in event polling loop — RSS grows 50MB per hour",
+            "Fix broken symlink detection in project scanner",
+            "Subprocess timeout not honored — zombie processes accumulate",
+            "JSONL parser crashes on malformed entries instead of skipping",
+            "Auth token refresh race condition causes 401 cascade",
+            "Incorrect path resolution when CWD differs from project root",
+            "Fix off-by-one in pagination causing missing eval entries",
+            "Build fails on CI but passes locally — env var not forwarded",
+            "Hub crashes under load — 100% agent stranding across all projects",
+            "Corrupted checkpoint file causes training loop to hang",
+            "Session cleanup fails to release file locks on abnormal exit",
+            "Fix flaky test that depends on filesystem ordering",
+            "API endpoint returns 500 when request body exceeds 1MB",
+            "Eval composite calculation wrong — missing dimension weight normalization",
+            "Context hub port conflict when multiple projects run simultaneously",
+            "Webhook delivery fails silently — no retry on HTTP 503",
+            "Training transform drops entries with unicode in description field",
+            "Config file watcher triggers infinite reload loop",
+            "Agent scheduling deadlock when two agents request same worktree",
+            "Stack overflow in recursive directory scanner on deep symlink chains",
+            "Metric aggregation double-counts entries spanning midnight boundary",
+        ],
+        "state_ranges": {"composite": (0.3, 0.75), "tests_ratio": (0.5, 0.9), "trajectory": (1, 8)},
+        "dimensions": ["test_pass_rate", "build_health", "error_rate", "hub_health"],
+    },
+    "refactor_code": {
+        "agents": ["code-quality", "feature-quality", "system-health"],
+        "goals": [
+            "peter.ts is 2400 lines — extract experiment selection into separate module",
+            "Reduce cyclomatic complexity in meta-orchestrator scheduling logic",
+            "Extract shared embedding logic from policy-head.ts and stratus-client.ts",
+            "Consolidate 3 different config loading patterns into single utility",
+            "Split monolithic train.ts into per-subcommand files",
+            "Replace nested callbacks with async/await in hub-client.ts",
+            "De-duplicate state capture code between sync and async versions",
+            "Extract common JSONL read/write into shared storage utility",
+            "Move inline type definitions to dedicated types/ files",
+            "Simplify 4-level deep if/else chain in agent scheduling",
+            "Replace magic numbers with named constants in eval scoring",
+            "Convert class-based API to functional composition for testability",
+            "Extract worktree management from peter.ts into worktree-manager.ts",
+            "Reduce coupling between eval-store and training-buffer modules",
+            "Consolidate error handling patterns — 5 different approaches in codebase",
+            "Decompose 300-line function in counterfactual-engine into smaller helpers",
+            "Replace string-based action types with enum for type safety",
+            "Extract telemetry event emission into decorator pattern",
+            "Simplify planning loop scoring — too many weight parameters",
+            "Break circular dependency between hub-client and state-capture",
+            "Flatten deeply nested Promise chains in dynamics-model",
+            "Extract CLI argument validation into reusable validators",
+            "Consolidate 3 different JSON serialization approaches",
+            "Module has 15 exported functions — group into logical classes",
+            "Remove dead code paths from v1 migration that are no longer used",
+            "Refactor test helpers — too much duplication across test files",
+            "Replace manual retries with shared retry utility",
+            "Extract benchmark scenario generation into data-driven config",
+            "Simplify imports — 40% of files import from 8+ modules",
+            "Restructure commands/ directory — flat list of 20 files hard to navigate",
+        ],
+        "state_ranges": {"composite": (0.65, 0.85), "tests_ratio": (0.9, 1.0), "trajectory": (3, 10)},
+        "dimensions": ["code_quality", "test_pass_rate", "build_health", "maintainability"],
+    },
+    "add_tests": {
+        "agents": ["test-coverage", "feature-quality", "code-quality"],
+        "goals": [
+            "Test coverage at 45% — add unit tests for uncovered modules",
+            "New planning-loop.ts has no tests — add comprehensive test suite",
+            "Add integration tests for the full training pipeline flow",
+            "Write regression tests for the hub crash recovery path",
+            "Coverage dropped from 80% to 65% after refactor — restore tests",
+            "Add property-based tests for state similarity computation",
+            "Missing edge case tests for empty training buffer scenarios",
+            "Add tests for counterfactual bridge ratio targeting logic",
+            "Write snapshot tests for CLI output formatting",
+            "Test online learning rollback behavior under degradation",
+            "Add benchmark scenario tests to prevent scenario drift",
+            "Missing tests for World Model store JSONL rotation",
+            "Write tests for invariant monitor with multiple simultaneous violations",
+            "Add tests for PolicyHead v1/v2 backward compatibility",
+            "Test experience replay buffer reservoir sampling distribution",
+            "Add end-to-end test for eval → training tuple → retrain flow",
+            "Missing tests for config file migration between versions",
+            "Write tests for agent scheduling fairness guarantees",
+            "Test WebSocket reconnection behavior under network flapping",
+            "Add chaos tests — what happens when hub dies mid-operation",
+            "Write tests for the v2 data transform with malformed input",
+            "Test that benchmark scenarios cover all tools in domain.json",
+            "Add tests for counterfactual quality scoring and ranking",
+            "Missing tests for training threshold calculation edge cases",
+            "Write performance tests — ensure eval completes under 30s",
+            "Test concurrent agent access to shared training buffer file",
+            "Add tests for git worktree cleanup on session abort",
+            "Write tests for embedding cache hit/miss behavior",
+            "Test that online training properly mixes historical and new data",
+            "Add tests for planning loop timeout behavior",
+        ],
+        "state_ranges": {"composite": (0.4, 0.75), "tests_ratio": (0.3, 0.7), "trajectory": (1, 6)},
+        "dimensions": ["test_coverage", "test_pass_rate", "code_quality"],
+    },
+    "update_docs": {
+        "agents": ["docs-updater", "feature-quality", "code-quality"],
+        "goals": [
+            "3 new CLI commands (train online, train counterfactual, train benchmark) have no docs",
+            "README doesn't mention the planning loop architecture",
+            "API changed but usage examples in docs still show old interface",
+            "SPEC.md doesn't reflect Drew's architecture decisions",
+            "Missing migration guide for v1 to v2 policy head upgrade",
+            "Add architecture diagram for PH → DM → InvariantMonitor flow",
+            "Document the 70/30 real/synthetic training data ratio decision",
+            "Getting Started guide doesn't cover agent setup with eval",
+            "Add troubleshooting section for common hub crash scenarios",
+            "Document environment variables needed for Stratus integration",
+            "CHANGELOG hasn't been updated in 3 weeks — summarize changes",
+            "Add JSDoc comments to all exported functions in planning-loop.ts",
+            "Document benchmark scenario format for contributors",
+            "Write onboarding doc for new team members joining the project",
+            "API reference for World Model types is missing",
+            "Document the nightly training pipeline flow and cron schedule",
+            "Add examples to CLI help text for complex commands",
+            "Write decision record for transformer vs MLP architecture choice",
+            "Document cold start behavior for new agents without training data",
+            "Update ROADMAP.md to reflect Q1 achievements and Q2 goals",
+            "Missing deployment guide for Fly.io production setup",
+            "Add performance benchmarks to documentation",
+            "Document the experience replay buffer sampling algorithm",
+            "Write contributor guide for adding new benchmark scenarios",
+            "Update architecture overview — 4 new modules added this month",
+            "Document inference latency expectations for CPU vs GPU",
+            "Add FAQ section for common PolicyHead questions",
+            "Document the invariant monitoring TLA+ spec mapping",
+            "Write runbook for training pipeline failures",
+            "Update team configuration docs with new role definitions",
+        ],
+        "state_ranges": {"composite": (0.7, 0.9), "tests_ratio": (0.9, 1.0), "trajectory": (1, 5)},
+        "dimensions": ["doc_coverage", "test_pass_rate", "code_quality"],
+    },
+    "optimize_performance": {
+        "agents": ["jfl-cli-speed", "build-optimizer", "system-performance-rl"],
+        "goals": [
+            "CLI startup takes 450ms — target sub-100ms for interactive use",
+            "Build time increased from 15s to 45s after adding new modules",
+            "Eval run takes 3 minutes — optimize test execution order",
+            "PolicyHead inference on CPU is 500ms — need under 100ms for planning loop",
+            "Bundle size grew to 2.4MB — tree-shake unused dependencies",
+            "Hub memory usage grows 200MB/day — find and fix the leak",
+            "Training pipeline takes 20 minutes — parallelize data loading",
+            "State capture function scans entire filesystem — add caching",
+            "JSONL file reads are O(n) — add index for faster lookups",
+            "Embedding precompute makes 1850 API calls — batch more aggressively",
+            "CI pipeline runs 4 minutes — parallelize test suites",
+            "Agent scheduling loop has O(n²) complexity — optimize to O(n log n)",
+            "Database queries for eval history scan full table — add time index",
+            "Webpack hot reload takes 8s — switch to esbuild or turbopack",
+            "Docker image is 1.2GB — multi-stage build to reduce to 200MB",
+            "API response time p95 is 2s — add response caching",
+            "Training data loading allocates 500MB — stream instead of load-all",
+            "Git operations block main thread — move to worker",
+            "Reduce TypeScript compile time by splitting into project references",
+            "Embedding cache lookup is O(n) string comparison — use hash map",
+            "Stratus API calls average 1.2s — add connection pooling",
+            "File watcher triggers on generated files — add ignore patterns",
+            "Reduce cold start time for new agent sessions",
+            "Test suite runs sequentially — parallelize with vitest threads",
+            "JSON serialization of large state objects takes 200ms — use msgpack",
+            "Worktree creation takes 5s per agent — pre-create pool",
+            "Reduce memory footprint of PolicyHead inference from 400MB to 100MB",
+            "Optimize counterfactual generation — currently 10s per transition",
+            "Planning loop timeout at 5s but most time spent in state capture",
+            "Startup module loading — lazy-import heavy dependencies",
+        ],
+        "state_ranges": {"composite": (0.5, 0.8), "tests_ratio": (0.85, 1.0), "trajectory": (2, 8)},
+        "dimensions": ["build_health", "p90_ms", "test_pass_rate"],
+    },
+    "dependency_update": {
+        "agents": ["dependency-updater", "system-health", "security"],
+        "goals": [
+            "3 critical CVEs in production dependencies — update immediately",
+            "Node.js 18 reaches EOL — migrate to Node.js 22",
+            "TypeScript 5.2 to 5.8 upgrade needed for new language features",
+            "Lodash dependency is 600KB — replace with native alternatives",
+            "Conflicting peer dependency warnings on npm install",
+            "Security audit shows 12 moderate vulnerabilities in dep tree",
+            "Vitest major version update available — check breaking changes",
+            "React 19 release — plan migration from React 18",
+            "Python requirements.txt has unpinned versions — lock them",
+            "Chalk v5 is ESM-only — update import statements",
+            "Commander.js deprecated API warnings on every CLI run",
+            "torch 2.10 has MPS performance improvements — update",
+            "npm audit fix --force broke 3 packages — resolve manually",
+            "Lockfile merge conflicts on every PR — regenerate clean lockfile",
+            "Drop support for Node.js 16 — remove compatibility shims",
+            "Update GitHub Actions runners from ubuntu-20.04 to ubuntu-24.04",
+            "Dependabot PRs piling up — batch review and merge",
+            "ESLint config uses deprecated rules — migrate to flat config",
+            "prettier and eslint conflict on formatting — resolve config",
+            "Replace moment.js (deprecated) with date-fns",
+            "Upgrade Docker base image from node:18-alpine to node:22-alpine",
+            "Pin transitive dependency that introduced breaking change",
+            "Remove unused dependencies flagged by depcheck",
+            "Update Fly.io CLI to latest version for new deployment features",
+            "Migrate from ts-node to tsx for faster script execution",
+            "Update pre-commit hooks to latest versions",
+            "Resolve yarn.lock and package-lock.json coexistence issue",
+            "Update SSH library for CVE-2026-1234 fix",
+            "Migrate from legacy @types packages to built-in TypeScript types",
+            "Update embedding model SDK for new Stratus API version",
+        ],
+        "state_ranges": {"composite": (0.7, 0.85), "tests_ratio": (0.9, 1.0), "trajectory": (1, 4)},
+        "dimensions": ["security_score", "test_pass_rate", "build_health"],
+    },
+    "data_pipeline": {
+        "agents": ["data-engineer", "telemetry-rl", "feature-quality"],
+        "goals": [
+            "Training buffer v1→v2 transform dropping 20% of entries — fix conversion",
+            "Embedding precompute OOM on datasets over 5000 texts — chunk processing",
+            "Counterfactual generation produces invalid probability masses > 1.0",
+            "Training data has duplicate entries — add deduplication step",
+            "JSONL rotation not triggering — daily files growing unbounded",
+            "Eval results not flowing to training buffer — broken event pipeline",
+            "Missing data validation between transform and train steps",
+            "Telemetry events lost during hub restart — add persistent queue",
+            "Training buffer merge from multiple agents has ID collisions",
+            "Embedding cache invalidation needed when model version changes",
+            "State snapshot hourly capture skipping overnight hours",
+            "Transition store index.json growing unbounded — add compaction",
+            "Training data export to CSV has encoding issues with unicode",
+            "Counterfactual bridge not respecting 70/30 ratio target",
+            "Benchmark JSONL not included in precompute scan — missing embeddings",
+            "Mining tuples from journals produces low-quality entries — add filters",
+            "Training buffer readAll() doesn't handle missing service directories",
+            "Eval scored events not triggering auto-retrain flow",
+            "Dynamics cache TTL not enforced — stale predictions served",
+            "Training data split not stratified — rare classes absent from val/test",
+            "ETL pipeline silently succeeds with 0 outputs — add empty check",
+            "Counterfactual store getAllCounterfactuals has no pagination",
+            "Training buffer JSONL corrupted by concurrent writes — add file locking",
+            "State capture fails silently and returns empty state — add validation",
+            "Precompute script doesn't resume after API rate limit — add backoff",
+            "Training history not preserved across model versions",
+            "Missing data lineage tracking — can't trace training tuple to source",
+            "Batch embedding API returns different dimension than cached — mismatch",
+            "Transform script doesn't handle new action types added to domain.json",
+            "Journal mining extracts wrong reward signal from session-end entries",
+        ],
+        "state_ranges": {"composite": (0.4, 0.7), "tests_ratio": (0.7, 0.9), "trajectory": (3, 9)},
+        "dimensions": ["pipeline_health", "data_quality", "test_pass_rate"],
+    },
+    "add_monitoring": {
+        "agents": ["observability", "system-health", "telemetry-rl"],
+        "goals": [
+            "No visibility into agent performance — add per-agent metrics dashboard",
+            "Hub crashes go undetected for hours — add health check with alerting",
+            "Training pipeline has no metrics — add loss/accuracy/time tracking",
+            "Missing observability for planning loop decision quality",
+            "Add prometheus metrics endpoint for hub connection count",
+            "No alerting when eval scores drop below threshold",
+            "Agent session duration not tracked — add timing telemetry",
+            "Missing structured logging in production — add log levels and context",
+            "Training buffer growth rate not monitored — could fill disk",
+            "Add tracing for Stratus API calls — latency and error rates",
+            "No metrics for counterfactual generation quality",
+            "Hub memory usage not tracked — only noticed leak when OOM killed",
+            "Add CI pipeline duration metrics to catch performance regressions",
+            "Missing audit log for config changes across sessions",
+            "No dashboard showing agent scheduling decisions and outcomes",
+            "Add embedding cache hit rate monitoring",
+            "Training data quality metrics not tracked over time",
+            "Missing error rate tracking per CLI command",
+            "Add uptime monitoring for all registered services",
+            "No metrics for planning loop latency by phase (capture/propose/rollout/select)",
+            "Add disk usage monitoring for world-model-data directory",
+            "Missing notification when nightly training pipeline fails",
+            "Add agent health heartbeat mechanism",
+            "No tracking of counterfactual prediction accuracy vs actual outcomes",
+            "Missing SLA metrics for hub availability",
+            "Add logging for experience replay buffer composition per training run",
+            "No visibility into online learning rollback frequency",
+            "Missing git operation timing — merges, rebases, worktree creation",
+            "Add cost tracking for Stratus API usage per agent",
+            "No alerting on training data distribution drift",
+        ],
+        "state_ranges": {"composite": (0.6, 0.8), "tests_ratio": (0.85, 1.0), "trajectory": (2, 6)},
+        "dimensions": ["observability", "test_pass_rate", "build_health"],
+    },
+    "security_hardening": {
+        "agents": ["security", "system-health", "code-quality"],
+        "goals": [
+            "API keys stored in plaintext config files — use secure keychain",
+            "No input validation on CLI arguments — potential injection",
+            "Hub WebSocket accepts connections without authentication",
+            "Training data contains embedded API keys — sanitize before storage",
+            "Missing rate limiting on eval submission endpoint",
+            "Subprocess commands built from user input without escaping",
+            "JWT tokens don't expire — add expiration and refresh flow",
+            "Agent-to-agent communication unencrypted — add TLS",
+            "Secrets leaked in CI logs — mask sensitive environment variables",
+            "File permissions too permissive on checkpoint files",
+            "No CORS configuration on hub HTTP endpoints",
+            "Missing Content-Security-Policy headers on web dashboard",
+            "Git credentials cached in plaintext in home directory",
+            "Webhook endpoint accepts requests without signature verification",
+            "Training buffer contains PII from user sessions — add redaction",
+            "No audit trail for who accessed admin endpoints",
+            "Missing RBAC — all authenticated users have full access",
+            "Dependency has known RCE vulnerability — update urgently",
+            "SSH key rotation not automated — some keys 2 years old",
+            "SQL injection possible in raw query builder utility",
+            "Missing CSRF protection on form submissions",
+            "API error messages expose internal stack traces",
+            "No network segmentation between agent sandbox and main system",
+            "Stratus API key rotation not implemented — single key for all",
+            "Missing integrity checks on downloaded model checkpoints",
+            "Plaintext passwords in legacy migration scripts",
+            "No timeout on external HTTP requests — DoS vector",
+            "Missing certificate pinning for critical API connections",
+            "Agent can read files outside its worktree scope",
+            "Eval results endpoint returns data for all agents without authz",
+        ],
+        "state_ranges": {"composite": (0.6, 0.8), "tests_ratio": (0.8, 1.0), "trajectory": (1, 5)},
+        "dimensions": ["security_score", "test_pass_rate", "code_quality"],
+    },
+}
+
+
+def generate_state(agent: str, ranges: dict, dims: list[str], rng: random.Random) -> str:
+    composite = rng.uniform(*ranges["composite"])
+    tests_ratio = rng.uniform(*ranges["tests_ratio"])
+    trajectory = rng.randint(*ranges["trajectory"])
+    tests_total = rng.choice([15, 20, 25, 30, 35, 40, 45, 50])
+    tests_passing = int(tests_total * tests_ratio)
+
+    dim_scores = {}
+    for d in dims:
+        dim_scores[d] = rng.uniform(0.2, 0.95)
+
+    dims_str = ", ".join(f"{k}={v:.4f}" for k, v in dim_scores.items())
+
+    n_deltas = rng.randint(1, 4)
+    deltas = [rng.uniform(-0.05, 0.05) for _ in range(n_deltas)]
+    deltas_str = ", ".join(f"{'+' if d >= 0 else ''}{d:.4f}" for d in deltas)
+
+    return "\n".join([
+        f"Agent: {agent}",
+        f"Composite: {composite:.4f}",
+        f"Tests: {tests_passing}/{tests_total}",
+        f"Trajectory: {trajectory}",
+        f"Dimensions: {dims_str}",
+        f"Recent deltas: {deltas_str}",
+    ])
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Generate balanced synthetic training data")
+    parser.add_argument("--output", default=".jfl/v2-data", help="Output directory")
+    parser.add_argument("--per-tool", type=int, default=60, help="Examples per tool (default: 60)")
+    parser.add_argument("--seed", type=int, default=42, help="Random seed")
+    args = parser.parse_args()
+
+    rng = random.Random(args.seed)
+    examples = []
+
+    for tool_name, tmpl in TEMPLATES.items():
+        for i in range(args.per_tool):
+            agent = rng.choice(tmpl["agents"])
+            goal = tmpl["goals"][i % len(tmpl["goals"])]
+            state = generate_state(agent, tmpl["state_ranges"], tmpl["dimensions"], rng)
+
+            examples.append({
+                "current_state": state,
+                "goal": goal,
+                "correct_tool": tool_name,
+                "source": "synthetic_balanced",
+            })
+
+    rng.shuffle(examples)
+
+    n = len(examples)
+    train_end = int(n * 0.7)
+    val_end = int(n * 0.85)
+
+    splits = {
+        "train": examples[:train_end],
+        "val": examples[train_end:val_end],
+        "test": examples[val_end:],
+    }
+
+    os.makedirs(args.output, exist_ok=True)
+
+    for split_name, split_data in splits.items():
+        path = os.path.join(args.output, f"{split_name}.jsonl")
+        # Append to existing split files
+        with open(path, "a") as f:
+            for ex in split_data:
+                f.write(json.dumps(ex) + "\n")
+        print(f"  {split_name}: +{len(split_data)} synthetic → {path}")
+
+    from collections import Counter
+    tool_counts = Counter(ex["correct_tool"] for ex in examples)
+    print(f"\nGenerated {n} synthetic examples:")
+    for tool, count in tool_counts.most_common():
+        print(f"  {tool:25s} {count:4d}")
+
+
+if __name__ == "__main__":
+    main()

package/scripts/train/v2/generate_data.py

@@ -0,0 +1,219 @@
+"""
+Synthetic training data generator for v2 policy head.
+
+Generates (current_state, goal, correct_tool) tuples using LLM API.
+Adapted from Drew's generate_data.py for JFL's action taxonomy.
+"""
+
+import json
+import os
+import sys
+import random
+import argparse
+import time
+from pathlib import Path
+
+def load_domain(path: str) -> dict:
+    with open(path) as f:
+        return json.load(f)
+
+def get_llm_client():
+    api_key = os.environ.get("OPENAI_API_KEY") or os.environ.get("ANTHROPIC_API_KEY")
+    if not api_key:
+        print("Set OPENAI_API_KEY or ANTHROPIC_API_KEY for synthetic data generation")
+        sys.exit(1)
+
+    if os.environ.get("OPENAI_API_KEY"):
+        from openai import OpenAI
+        return OpenAI(), "openai"
+    else:
+        from anthropic import Anthropic
+        return Anthropic(), "anthropic"
+
+def generate_goals_for_tool(client, client_type: str, tool: dict, all_tools: list, n: int = 50) -> list:
+    tool_list_str = "\n".join(
+        f"- {t['name']}: {t['description']}" for t in all_tools
+    )
+
+    prompt = f"""You are generating training data for an AI action-routing system in software development.
+
+Given these {len(all_tools)} available action types:
+{tool_list_str}
+
+Generate exactly {n} realistic goals/situations that should route to:
+  Action: {tool["name"]}
+  Description: {tool["description"]}
+
+Requirements:
+1. Each goal should describe a real software development situation or user request
+2. Include variety: specific bug reports, vague feature requests, urgent fixes, routine maintenance
+3. Some should be indirect (describe the symptom, not the action needed)
+4. Some should use technical jargon (CI/CD, linting, test coverage, etc.)
+5. Some should include context about current project state
+6. Include goals that could be confused with OTHER action types
+7. Do NOT include the action type name in the goal text
+8. Goals should reflect what an AI coding agent would encounter
+
+Output as a JSON object with a "goals" key containing an array of strings. No explanation."""
+
+    try:
+        if client_type == "openai":
+            response = client.chat.completions.create(
+                model="gpt-4o-mini",
+                messages=[{"role": "user", "content": prompt}],
+                temperature=0.9,
+                response_format={"type": "json_object"},
+            )
+            result = json.loads(response.choices[0].message.content)
+        else:
+            response = client.messages.create(
+                model="claude-sonnet-4-5-20250514",
+                max_tokens=4096,
+                messages=[{"role": "user", "content": prompt}],
+            )
+            text = response.content[0].text
+            json_match = text[text.index("{"):text.rindex("}") + 1]
+            result = json.loads(json_match)
+
+        goals = result.get("goals", result.get("queries", []))
+        return goals[:n] if isinstance(goals, list) else []
+    except Exception as e:
+        print(f"  Error generating goals for {tool['name']}: {e}")
+        return []
+
+def generate_states_for_goal(client, client_type: str, goal: str, domain_desc: str, n: int = 2) -> list:
+    prompt = f"""You are generating training data for an AI planning system in:
+"{domain_desc}"
+
+Given this development goal: "{goal}"
+
+Generate {n} different realistic "current state" descriptions that an AI coding agent might be in
+when this goal arises. Each state should describe:
+- Current codebase state (test results, scores, recent changes)
+- What the agent currently knows about the project
+- Any relevant trajectory information (recent actions taken)
+
+States should vary in how much context is available and what prior work has been done.
+
+Output as a JSON object with a "states" key containing an array of strings (1-3 sentences each)."""
+
+    try:
+        if client_type == "openai":
+            response = client.chat.completions.create(
+                model="gpt-4o-mini",
+                messages=[{"role": "user", "content": prompt}],
+                temperature=0.8,
+                response_format={"type": "json_object"},
+            )
+            result = json.loads(response.choices[0].message.content)
+        else:
+            response = client.messages.create(
+                model="claude-sonnet-4-5-20250514",
+                max_tokens=2048,
+                messages=[{"role": "user", "content": prompt}],
+            )
+            text = response.content[0].text
+            json_match = text[text.index("{"):text.rindex("}") + 1]
+            result = json.loads(json_match)
+
+        states = result.get("states", [])
+        return states[:n]
+    except Exception as e:
+        print(f"  Error generating states: {e}")
+        return [f"Agent working on codebase. Goal context: {goal[:100]}"]
+
+def generate_dataset(
+    domain_path: str,
+    goals_per_tool: int = 50,
+    states_per_goal: int = 2,
+    output_dir: str = "data",
+    seed: int = 42,
+):
+    domain = load_domain(domain_path)
+    tools = domain["tools"]
+    domain_desc = domain["description"]
+
+    client, client_type = get_llm_client()
+    print(f"Using {client_type} API for generation")
+
+    all_examples = []
+
+    for tool_idx, tool in enumerate(tools):
+        print(f"[{tool_idx + 1}/{len(tools)}] Generating goals for: {tool['name']}")
+
+        goals = generate_goals_for_tool(client, client_type, tool, tools, n=goals_per_tool)
+        print(f"  Generated {len(goals)} goals")
+
+        for goal_idx, goal in enumerate(goals):
+            if goal_idx % 10 == 0:
+                print(f"  Generating states: {goal_idx}/{len(goals)}")
+
+            states = generate_states_for_goal(client, client_type, goal, domain_desc, n=states_per_goal)
+
+            for state in states:
+                example = {
+                    "current_state": state,
+                    "goal": goal,
+                    "correct_tool": tool["name"],
+                    "tool_category": tool["category"],
+                    "source": "synthetic",
+                }
+                all_examples.append(example)
+
+            time.sleep(0.1)
+
+    random.seed(seed)
+    random.shuffle(all_examples)
+
+    n = len(all_examples)
+    train_end = int(n * 0.7)
+    val_end = int(n * 0.85)
+
+    splits = {
+        "train": all_examples[:train_end],
+        "val": all_examples[train_end:val_end],
+        "test": all_examples[val_end:],
+    }
+
+    os.makedirs(output_dir, exist_ok=True)
+
+    for split_name, split_data in splits.items():
+        path = os.path.join(output_dir, f"{split_name}.jsonl")
+        with open(path, "w") as f:
+            for ex in split_data:
+                f.write(json.dumps(ex) + "\n")
+        print(f"  {split_name}: {len(split_data)} examples -> {path}")
+
+    print(f"\nTotal: {n} examples across {len(tools)} tools")
+    print(f"  Per tool: ~{n // len(tools)} examples")
+    return splits
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Generate synthetic training data for v2 policy head")
+    parser.add_argument("--domain", default=None, help="Path to domain.json")
+    parser.add_argument("--output", default=".jfl/v2-data", help="Output directory")
+    parser.add_argument("--goals-per-tool", type=int, default=50, help="Goals to generate per tool")
+    parser.add_argument("--states-per-goal", type=int, default=2, help="States to generate per goal")
+    parser.add_argument("--seed", type=int, default=42, help="Random seed")
+    args = parser.parse_args()
+
+    domain_path = args.domain
+    if domain_path is None:
+        domain_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "domain.json")
+
+    if not os.path.exists(domain_path):
+        print(f"Domain file not found: {domain_path}")
+        sys.exit(1)
+
+    generate_dataset(
+        domain_path=domain_path,
+        goals_per_tool=args.goals_per_tool,
+        states_per_goal=args.states_per_goal,
+        output_dir=args.output,
+        seed=args.seed,
+    )
+
+
+if __name__ == "__main__":
+    main()