API Reference for rune-agent¶

def should_retry(state: RuneState) -> Literal["generate", "save_trajectory"]:
    """Route from reflect: retry if attempts remain and tests failed, else save.

    This is a REAL implementation, not a stub:
    - If tests passed -> save_trajectory (success)
    - If attempt_count >= max_attempts -> save_trajectory (exhausted)
    - Otherwise -> generate (retry)

    Args:
        state: Current agent state after reflection.

    Returns:
        Next node name: 'generate' to retry or 'save_trajectory' to finish.

    Example:
        >>> state = {"tests_passed": False, "attempt_count": 0, "max_attempts": 3}
        >>> should_retry(state)
        'generate'
        >>> state2 = {"tests_passed": True, "attempt_count": 1, "max_attempts": 3}
        >>> should_retry(state2)
        'save_trajectory'
    """
    if state["tests_passed"]:
        return "save_trajectory"
    if state["attempt_count"] >= state["max_attempts"]:
        return "save_trajectory"
    return "generate"

def create_graph() -> Any:
    """Create and compile the Rune agent workflow graph.

    The graph follows this flow:
    START -> generate -> execute -> reflect -> [should_retry]
        -> generate (retry) OR save_trajectory -> END

    Returns:
        Compiled StateGraph ready for execution.

    Example:
        >>> graph = create_graph()
        >>> graph is not None
        True
    """
    workflow = StateGraph(RuneState)

    workflow.add_node("generate", generate_node)
    workflow.add_node("execute", execute_node)
    workflow.add_node("reflect", reflect_node)
    workflow.add_node("save_trajectory", save_trajectory_node)

    workflow.add_edge(START, "generate")
    workflow.add_edge("generate", "execute")
    workflow.add_edge("execute", "reflect")
    workflow.add_conditional_edges("reflect", should_retry)
    workflow.add_edge("save_trajectory", END)

    return workflow.compile()

def create_single_iteration_graph() -> Any:
    """Create a graph for one iteration: generate → execute → reflect → END.

    Used by rune_runner's outer iteration loop. Each iteration runs one
    generate/execute/reflect cycle, then returns control to the outer loop
    which runs the hypernetwork to produce a new adapter for the next iteration.

    Returns:
        Compiled StateGraph for single-iteration execution.
    """
    workflow = StateGraph(RuneState)

    workflow.add_node("generate", generate_node)
    workflow.add_node("execute", execute_node)
    workflow.add_node("reflect", reflect_node)

    workflow.add_edge(START, "generate")
    workflow.add_edge("generate", "execute")
    workflow.add_edge("execute", "reflect")
    workflow.add_edge("reflect", END)

    return workflow.compile()

def get_graph() -> Any:
    """Get or create the compiled graph instance (thread-safe singleton).

    Uses double-checked locking so the expensive create_graph() call is only
    made once, even if multiple threads call get_graph() concurrently.

    Returns:
        Compiled StateGraph instance.
    """
    global _graph
    if _graph is None:
        with _graph_lock:
            if _graph is None:
                _graph = create_graph()
    return _graph

async def generate_node(state: RuneState) -> dict[str, Any]:
    """Generate code for the given task using the inference layer.

    Calls the LLM (with optional LoRA adapters) to produce a code solution
    for the task description.

    Args:
        state: Current agent state with task description and context.

    Returns:
        State update dict with generated_code key.

    Example:
        >>> state = {"task_description": "Write fibonacci", "task_type": "function",
        ...          "test_suite": "assert fib(5) == 5", "adapter_ids": []}
        >>> result = await generate_node(state)
        >>> 'generated_code' in result
        True
    """
    # Read env vars inside function body so monkeypatch.setenv() works in tests
    model: str = os.environ.get("RUNE_MODEL", DEFAULT_MODEL)
    adapter_id: str | None = state["adapter_ids"][0] if state["adapter_ids"] else None

    provider = get_provider()
    user_prompt = _build_prompt(state)
    phase = state.get("phase")
    system_prompt = _PHASE_SYSTEM_PROMPTS.get(phase or "", DEFAULT_SYSTEM_PROMPT)

    max_tokens = int(os.environ.get("RUNE_MAX_TOKENS", "1024"))

    result: GenerationResult = await provider.generate(
        prompt=user_prompt,
        model=model,
        adapter_id=adapter_id,
        max_tokens=max_tokens,
        system_prompt=system_prompt,
    )

    extracted = _extract_code(result.text)
    logger.info(
        "generate_node: attempt=%d, model=%s, adapter_id=%s, tokens=%d, finish=%s",
        state["attempt_count"],
        result.model,
        result.adapter_id,
        result.token_count,
        result.finish_reason,
    )

    return {"generated_code": extracted, "finish_reason": result.finish_reason}

async def execute_node(state: RuneState) -> dict[str, Any]:
    """Execute the generated code in a sandboxed environment.

    Runs the generated code against the test suite in an isolated subprocess
    and captures stdout, stderr, and exit code.

    Args:
        state: Current agent state with generated code and test suite.

    Returns:
        State update dict with stdout, stderr, exit_code, and tests_passed keys.

    Example:
        >>> state = {"generated_code": "def fib(n): return n", "test_suite": ""}
        >>> result = await execute_node(state)
        >>> 'tests_passed' in result
        True
    """
    # Read env var inside function body so monkeypatch.setenv() works in tests
    timeout: int = int(os.environ.get("RUNE_EXEC_TIMEOUT", DEFAULT_TIMEOUT))

    script = state["generated_code"] + "\n\n" + state["test_suite"]

    # Auto-inject unittest.main() if TestCase classes exist but no runner
    if has_unittest_classes(script) and not re.search(r"unittest\.main\s*\(", script):
        script += (
            "\n\nimport unittest\nif __name__ == '__main__':\n    unittest.main()\n"
        )

    backend = get_sandbox_backend()
    result = backend.run(script, timeout=timeout)

    stdout = result.stdout
    stderr = result.stderr
    exit_code = result.exit_code

    # Determine test validation based on unittest output
    combined = state["generated_code"] + "\n" + state["test_suite"]
    has_tests = has_unittest_classes(combined)
    passed_count, total_count = count_test_results(stdout, stderr)
    tests_ran = total_count > 0

    if has_tests and tests_ran and exit_code == 0 and not result.is_timed_out:
        tests_passed = True
    elif has_tests and not tests_ran:
        tests_passed = False
    elif not has_tests:
        tests_passed = False
    else:
        tests_passed = False

    logger.info(
        "execute_node: exit_code=%d, tests_passed=%s, test_count=%d, tests_ran=%s",
        exit_code,
        tests_passed,
        total_count,
        tests_ran,
    )

    return {
        "stdout": stdout,
        "stderr": stderr,
        "exit_code": exit_code,
        "tests_passed": tests_passed,
        "test_count": total_count,
        "tests_ran": tests_ran,
    }

async def reflect_node(state: RuneState) -> dict[str, Any]:
    """Reflect on execution results and record the attempt in trajectory.

    Increments the attempt counter and appends the current attempt's data
    to the trajectory list. Does not make any LLM calls.

    Args:
        state: Current agent state with execution results.

    Returns:
        State update dict with incremented attempt_count and extended trajectory.

    Example:
        >>> state = {"attempt_count": 0, "generated_code": "def fib(n): pass",
        ...          "exit_code": 1, "tests_passed": False, "trajectory": [],
        ...          "stdout": "", "stderr": ""}
        >>> result = await reflect_node(state)
        >>> result['attempt_count']
        1
    """
    step: dict[str, Any] = {
        "generated_code": state["generated_code"],
        "stdout": state["stdout"],
        "stderr": state["stderr"],
        "exit_code": state["exit_code"],
        "tests_passed": state["tests_passed"],
    }

    # Use list concatenation (not .append()) — LangGraph requires immutable state
    new_trajectory: list[dict[str, Any]] = state["trajectory"] + [step]
    new_attempt_count = state["attempt_count"] + 1

    logger.info(
        "reflect_node: attempt=%d -> %d, trajectory_length=%d",
        state["attempt_count"],
        new_attempt_count,
        len(new_trajectory),
    )

    return {
        "attempt_count": new_attempt_count,
        "trajectory": new_trajectory,
    }

async def save_trajectory_node(state: RuneState) -> dict[str, Any]:
    """Save the completed trajectory for parametric memory training.

    Persists the trajectory to disk via record_trajectory() and determines
    the final outcome based on whether tests passed.

    Args:
        state: Current agent state with complete trajectory and outcome.

    Returns:
        State update dict with outcome key ('success' or 'exhausted').

    Example:
        >>> state = {"tests_passed": True, "session_id": "abc", "trajectory": [],
        ...          "task_description": "", "task_type": "", "adapter_ids": []}
        >>> result = await save_trajectory_node(state)
        >>> result['outcome']
        'success'
    """
    outcome = "success" if state["tests_passed"] else "exhausted"

    record_trajectory(
        session_id=state["session_id"],
        steps=state["trajectory"],
        outcome=outcome,
        task_description=state["task_description"],
        task_type=state["task_type"],
        adapter_ids=state["adapter_ids"],
    )

    logger.info(
        "save_trajectory_node: session_id=%s, outcome=%s",
        state["session_id"],
        outcome,
    )

    return {"outcome": outcome}