rune/scripts/research_event_relay.py

#!/usr/bin/env python3
"""research_event_relay.py — live event relay from mimo to canonical state (RSC8).

Consumes mimo's structured JSON events WHILE the process runs and publishes
factual ResearchEvents into the ResearchTaskTree state. No private CoT.

Event types emitted:
  StageStarted, PromptCompiled, QueryIssued, ToolStarted, ToolCompleted,
  SourceDiscovered, SourceRead, EvidenceAccepted, ToolFailed, StageCompleted, StageFailed
"""

from __future__ import annotations

import json
import os
import subprocess
import time
from pathlib import Path
from typing import Callable


class ResearchEventRelay:
    """Reads mimo JSON events in real-time and updates canonical state."""

    def __init__(self, tree, events_path: str, mimo_cmd: list[str], err_path: str, attempt=None):
        self.tree = tree
        self.events_path = events_path
        self.mimo_cmd = mimo_cmd
        self.err_path = err_path
        # OBS-1/D3 (CORRECTION-2): observational telemetry handle. The relay
        # marks tool windows; it never changes classification or control flow.
        self.attempt = attempt
        self.event_counts = {
            "queries": 0,
            "sources_discovered": 0,
            "sources_read": 0,
            "evidence_accepted": 0,
            "tool_failures": 0,
            "tool_calls": 0,
        }
        # Real tool receipts (RP3 evidence chain: QUERY → SOURCE → READ → CLAIM)
        self.receipts = []
        self.last_useful_event_time = time.time()
        self.stage_start_time = time.time()

    def _classify_event(self, ev: dict) -> str | None:
        """Classify a mimo JSON event into a ResearchEvent type."""
        ev_type = ev.get("type", "")
        part = ev.get("part", {})
        part_type = part.get("type", "")

        if ev_type == "tool" or part_type == "tool":
            tool = part.get("tool", "")
            state = part.get("state", {})
            status = state.get("status", "")
            tool_input = state.get("input", {})

            if self.attempt is not None and status == "running":
                self.attempt.mark_tool_start()
            if self.attempt is not None and status in ("completed", "error", "failed"):
                self.attempt.mark_tool_end()

            if status == "completed":
                self.event_counts["tool_calls"] += 1
                if "url" in tool_input:
                    url = tool_input["url"]
                    if "duckduckgo" in url or "search" in url:
                        self.event_counts["queries"] += 1
                        self.receipts.append(
                            {
                                "kind": "QUERY",
                                "tool": tool,
                                "url": url,
                                "ts": time.time(),
                            }
                        )
                        return "QueryIssued"
                    else:
                        self.event_counts["sources_read"] += 1
                        self.receipts.append(
                            {
                                "kind": "SOURCE_READ",
                                "tool": tool,
                                "url": url,
                                "ts": time.time(),
                            }
                        )
                        return "SourceRead"
                self.receipts.append(
                    {
                        "kind": "TOOL",
                        "tool": tool,
                        "input_keys": sorted(tool_input.keys()) if isinstance(tool_input, dict) else [],
                        "ts": time.time(),
                    }
                )
                return "ToolCompleted"
            elif status in ("error", "failed"):
                self.event_counts["tool_failures"] += 1
                self.receipts.append(
                    {
                        "kind": "TOOL_FAILED",
                        "tool": tool,
                        "input_keys": sorted(tool_input.keys()) if isinstance(tool_input, dict) else [],
                        "ts": time.time(),
                    }
                )
                return "ToolFailed"
        return None

    def run(self) -> dict:
        """Run mimo, consume events in real-time, return the result."""
        os.makedirs(os.path.dirname(self.events_path) or ".", exist_ok=True)

        proc = subprocess.Popen(
            self.mimo_cmd,
            stdout=open(self.events_path, "wb"),
            stderr=open(self.err_path, "wb"),
        )

        texts = []
        offset = 0

        while proc.poll() is None:
            # Read new lines from the events file
            try:
                with open(self.events_path, "r") as f:
                    f.seek(offset)
                    new_lines = f.readlines()
                    offset = f.tell()
            except FileNotFoundError:
                new_lines = []

            for line in new_lines:
                try:
                    ev = json.loads(line.strip())
                except json.JSONDecodeError:
                    continue

                event_type = self._classify_event(ev)
                if event_type:
                    self.last_useful_event_time = time.time()
                    print(f"  [{event_type}] {self._event_summary(ev)}", flush=True)

                # Collect text parts
                if ev.get("type") == "text":
                    t = ev.get("part", {}).get("text", "")
                    if t.strip():
                        texts.append(t)
                        self.last_useful_event_time = time.time()

            time.sleep(1)  # poll interval for file reads (not Gitea polling)

        proc.wait()
        return {
            "returncode": proc.returncode,
            "texts": texts,
            "event_counts": self.event_counts,
            "receipts": self.receipts,
            "elapsed": time.time() - self.stage_start_time,
        }

    def _event_summary(self, ev: dict) -> str:
        part = ev.get("part", {})
        tool = part.get("tool", "?")
        state = part.get("state", {})
        inp = state.get("input", {})
        if "url" in inp:
            return f"{tool} {inp['url'][:60]}"
        return f"{tool}"


def run_stage_with_relay(tree, out_dir: str) -> dict:
    """Run one research stage with live event relay."""
    from research_prompt_compiler import compile_research_prompt

    stage = tree.current_stage
    prompt = compile_research_prompt(tree, stage)

    events_path = os.path.join(out_dir, f"stage_{stage}.events.jsonl")
    err_path = os.path.join(out_dir, f"stage_{stage}.mimo.err.log")

    mimo_bin = os.environ.get("MIMO_BIN", "/var/lib/gitea-desktop-runner/bin/mimo")
    from rune.provider_context import require_model
    mimo_model = require_model()  # T16 §16F: provider layer owns selection
    mimo_agent = os.environ.get("MIMO_AGENT", "build")

    cmd = [
        mimo_bin,
        "run",
        "--agent",
        mimo_agent,
        "-m",
        mimo_model,
        "--dangerously-skip-permissions",
        "--format",
        "json",
        prompt,
    ]

    relay = ResearchEventRelay(tree, events_path, cmd, err_path)
    result = relay.run()

    if result["returncode"] != 0:
        tree.fail(f"mimo exited {result['returncode']}")
        return {"stage": stage, "status": "FAILED", "error": f"mimo exited {result['returncode']}"}

    if not result["texts"]:
        tree.fail("no text output from mimo")
        return {"stage": stage, "status": "FAILED", "error": "no text output"}

    artifact_text = "\n\n".join(result["texts"])
    artifact = {
        "text": artifact_text,
        "tool_calls": result["event_counts"]["tool_calls"],
        "event_counts": result["event_counts"],
        "elapsed": result["elapsed"],
    }

    next_stage = tree.advance(artifact)
    print(f"  stage {stage} completed → next: {next_stage}")
    print(f"  events: {json.dumps(result['event_counts'])}")
    print(f"  elapsed: {result['elapsed']:.0f}s")

    return {
        "stage": stage,
        "status": "COMPLETED",
        "next_stage": next_stage,
        "artifact": artifact,
        "event_counts": result["event_counts"],
    }