Add run derive subcommand

Author: aliroberts

tests/test_derive_handler.py

@@ -50,3 +50,145 @@ def test_plain_optimization_ui_best_respects_minimize_goal():
 
     assert any("best so far: 1.5" in line for line in lines)
     assert any("Best metric value: 1.5" in line for line in lines)
+
+
+def _make_plain_ui_with_capture() -> tuple[PlainOptimizationUI, list[str]]:
+    """Construct a PlainOptimizationUI that records output to a list instead of stdout."""
+    ui = PlainOptimizationUI(
+        run_id="run-1",
+        run_name="demo",
+        total_steps=5,
+        dashboard_url="https://example.com",
+        model="gpt-4",
+        metric_name="accuracy",
+    )
+    lines: list[str] = []
+    ui._print = lines.append
+    return ui, lines
+
+
+def test_plain_ui_on_init_prints_header():
+    """on_init prints the run banner exactly once. (Header used to be printed
+    by __enter__; the move to on_init keeps the same observable behavior for
+    non-derived runs.)"""
+    ui, lines = _make_plain_ui_with_capture()
+
+    ui.on_init()
+
+    output = "\n".join(lines)
+    assert "WECO OPTIMIZATION RUN" in output
+    assert "Run ID: run-1" in output
+    assert "Run Name: demo" in output
+    assert "Dashboard: https://example.com" in output
+    assert "Model: gpt-4" in output
+    assert "Metric: accuracy" in output
+    assert "Total Steps: 5" in output
+    # Non-derived run: no "Derived from" line
+    assert not any("Derived from" in line for line in lines)
+
+
+def test_plain_ui_on_init_includes_derived_from_line():
+    ui, lines = _make_plain_ui_with_capture()
+
+    ui.on_init(derived_from={
+        "run_id": "parent-uuid",
+        "node_id": "node-uuid",
+        "step": 7,
+        "metric_value": 0.842,
+    })
+
+    derived_lines = [line for line in lines if "Derived from" in line]
+    assert len(derived_lines) == 1
+    assert "parent-uuid" in derived_lines[0]
+    assert "step 7" in derived_lines[0]
+    assert "0.842" in derived_lines[0]
+
+
+def test_plain_ui_on_init_handles_derived_from_without_metric():
+    """A node with no metric_value (e.g., still pending eval) shouldn't crash
+    the header rendering."""
+    ui, lines = _make_plain_ui_with_capture()
+
+    ui.on_init(derived_from={
+        "run_id": "parent-uuid",
+        "node_id": "node-uuid",
+        "step": 0,
+        "metric_value": None,
+    })
+
+    derived_lines = [line for line in lines if "Derived from" in line]
+    assert len(derived_lines) == 1
+    assert "parent-uuid" in derived_lines[0]
+    assert "step 0" in derived_lines[0]
+    # No "(metric: ...)" suffix when metric_value is None. Specific to the
+    # suffix's literal form so the assertion is robust to metric_name values
+    # that happen to contain the substring "metric".
+    assert "(metric:" not in derived_lines[0]
+
+
+def test_plain_ui_enter_no_longer_prints_header():
+    """Header printing must happen via on_init now, not __enter__. This guards
+    against accidentally re-introducing the auto-print and double-printing the
+    header."""
+    ui, lines = _make_plain_ui_with_capture()
+
+    with ui:
+        pass
+
+    assert lines == []
+
+
+def test_live_ui_on_init_renders_derived_from_row():
+    """The Live panel grid gains a "From" row when derived_from is set."""
+    ui = LiveOptimizationUI(
+        console=Console(force_terminal=False, color_system=None),
+        run_id="run-1",
+        run_name="demo",
+        total_steps=3,
+        dashboard_url="https://example.com",
+        metric_name="accuracy",
+    )
+
+    ui.on_init(derived_from={
+        "run_id": "parent-uuid",
+        "node_id": "node-uuid",
+        "step": 4,
+        "metric_value": 0.91,
+    })
+
+    text = _render_to_text(ui._render())
+    assert "parent-uuid" in text
+    assert "step 4" in text
+    assert "0.91" in text
+
+
+def test_live_ui_on_init_without_derived_from_does_not_render_from_row():
+    """Negative cousin to test_live_ui_on_init_renders_derived_from_row.
+
+    Renders the same UI both with and without ``derived_from`` and asserts
+    the parent reference (the marker added by the "From" row) appears only
+    in the derived render. Comparing the two renders directly avoids
+    fragile substring checks against unrelated panel chrome.
+    """
+
+    def render(derived_from):
+        ui = LiveOptimizationUI(
+            console=Console(force_terminal=False, color_system=None),
+            run_id="run-1",
+            run_name="demo",
+            total_steps=3,
+            dashboard_url="https://example.com",
+        )
+        ui.on_init(derived_from=derived_from)
+        return _render_to_text(ui._render())
+
+    derived_text = render({"run_id": "parent-uuid", "node_id": "n", "step": 4, "metric_value": 0.91})
+    plain_text = render(None)
+
+    # The parent reference is added only by the "From" row, so its presence
+    # in one render and absence in the other proves the row is conditional.
+    assert "parent-uuid" in derived_text
+    assert "parent-uuid" not in plain_text
+    # Belt and braces: the row label itself only appears in the derived
+    # render (no other panel label contains "From" as a substring).
+    assert "From" not in plain_text

tests/test_ui.py

@@ -226,6 +226,30 @@ def _configure_run_subcommands(run_parser: argparse.ArgumentParser) -> None:
     revise_source.add_argument("-s", "--source", type=str, help="Path to a single source file")
     revise_source.add_argument("--sources", nargs="+", type=str, help="Paths to multiple source files")
 
+    # weco run derive <run-id>
+    p = subs.add_parser("derive", help="Create a new run derived from an existing run's step")
+    p.add_argument("run_id", type=str, help="Parent run UUID")
+    p.add_argument(
+        "--from-step", type=str, default="best", help="'best' (lineage-best, default), 'run-best', or a step number"
+    )
+    p.add_argument("-n", "--steps", type=int, default=None, help="Override step count for the derived run")
+    p.add_argument(
+        "-i",
+        "--additional-instructions",
+        type=str,
+        default=None,
+        help="Steering instructions for the new run (inline text or path to a file). "
+        "If omitted, the parent run's instructions are inherited.",
+    )
+    p.add_argument("--api-key", nargs="+", type=str, default=None, help="API keys in provider=key format")
+    p.add_argument(
+        "--output",
+        type=str,
+        choices=["rich", "plain"],
+        default="rich",
+        help="Output mode: 'rich' for interactive UI, 'plain' for machine-readable output",
+    )
+
     # weco run submit <run-id> --node <id>
     p = subs.add_parser("submit", help="Submit a pending node for evaluation (review mode)")
     p.add_argument("run_id", type=str, help="Run UUID")
@@ -350,7 +374,7 @@ def configure_resume_parser(resume_parser: argparse.ArgumentParser) -> None:
 
 def _dispatch_run_subcommand(sub: str, args: argparse.Namespace) -> None:
     """Dispatch ``weco run <subcommand>`` to the appropriate handler."""
-    from .commands.run import status, results, show, diff, stop, instruct, review, revise, submit
+    from .commands.run import status, results, show, diff, stop, instruct, review, revise, submit, derive
 
     def _collect_source_paths() -> list[str] | None:
         if getattr(args, "sources", None):
@@ -371,6 +395,15 @@ def _collect_source_paths() -> list[str] | None:
         ),
         "show": lambda: show.handle(run_id=args.run_id, step=args.step, console=console),
         "diff": lambda: diff.handle(run_id=args.run_id, step=args.step, against=args.against, console=console),
+        "derive": lambda: derive.handle(
+            run_id=args.run_id,
+            from_step=args.from_step,
+            steps=args.steps,
+            additional_instructions=args.additional_instructions,
+            api_keys=parse_api_keys(args.api_key),
+            output_mode=args.output,
+            console=console,
+        ),
         "stop": lambda: stop.handle(run_id=args.run_id, console=console),
         "instruct": lambda: instruct.handle(run_id=args.run_id, instructions=args.instructions, console=console),
         "review": lambda: review.handle(run_id=args.run_id, console=console),
@@ -390,8 +423,10 @@ def _collect_source_paths() -> list[str] | None:
     if handler is None:
         console.print(f"[bold red]Unknown run subcommand: {sub}[/]")
         sys.exit(1)
-    handler()
-    sys.exit(0)
+    # Handlers that drive an optimization loop (e.g. derive) return a bool to
+    # signal success/failure. Read-only handlers return None and exit cleanly.
+    result = handler()
+    sys.exit(0 if result is not False else 1)
 
 
 def execute_run_command(args: argparse.Namespace) -> None:
@@ -526,7 +561,7 @@ def main() -> None:
 def _main() -> None:
     """Internal main function containing the CLI logic."""
     parser = argparse.ArgumentParser(
-        description="[bold cyan]Weco CLI[/]\nEnhance your code with AI-driven optimization.",
+        description="Weco CLI\nEnhance your code with AI-driven optimization.",
         formatter_class=argparse.RawDescriptionHelpFormatter,
     )