kaiju_harness/router.py

#!/usr/bin/env python3
"""Unified Kaiju harness router.

This is the product-facing layer. Customers should not need to know whether a
task is a website, business document, one-file app, project starter, or repo
patch. The router picks the harness, writes artifacts, validates them, and
returns a manifest.
"""

from __future__ import annotations

import datetime as dt
import json
import re
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from kaiju_harness import app, business, business_suite, code_project, coding, repo_patch, website
from kaiju_harness.model_spec import request_json_spec
from kaiju_harness.verification import failed_checks, verify_output


TASK_TYPES = {"website", "business_document", "business_suite", "app", "code_project", "repo_patch", "coding"}
FORBIDDEN_TOKENS = ["sk_live_", "sk_test_", "rk_live_", "pplx-", "AIza", "anthropic_api_key"]
ROOT = Path(__file__).resolve().parents[1]
PROMPT_FILES = {
    "website": ROOT / "prompts/kaiju-website-spec-system.md",
    "business_document": ROOT / "prompts/kaiju-business-spec-system.md",
    "app": ROOT / "prompts/kaiju-app-spec-system.md",
    "code_project": ROOT / "prompts/kaiju-code-project-spec-system.md",
    "repo_patch": ROOT / "prompts/kaiju-repo-patch-spec-system.md",
}


@dataclass
class RouterResult:
    task_type: str
    artifact_type: str
    artifact_path: Path | None
    project_dir: Path | None
    changed_files: list[str]
    spec: dict[str, Any]
    manifest_path: Path
    manifest: dict[str, Any]
    response_text: str
    errors: list[str]


def slugify(value: str) -> str:
    return re.sub(r"[^a-z0-9]+", "-", value.lower()).strip("-")[:70] or "kaiju-task"


def now_stamp() -> str:
    return dt.datetime.now(dt.UTC).strftime("%Y%m%dT%H%M%SZ")


def has_any(lower: str, terms: list[str]) -> bool:
    return any(term in lower for term in terms)


def route_prompt(prompt: str, repo: Path | None = None, kind: str = "auto") -> str:
    if kind != "auto":
        if kind not in TASK_TYPES:
            raise ValueError(f"Unsupported task type: {kind}")
        return kind

    lower = prompt.lower()
    if repo is not None:
        return "repo_patch"

    business_suite_terms = [
        "kiyomi 7.7.7",
        "kiyomi777",
        "full kiyomi",
        "ai company",
        "business owner operating system",
        "owner-ready ai company",
        "launch kit",
        "connector pack",
        "the workshop",
        "teach-once",
        "teach once",
        "/kiyomi-do",
        "/kiyomi",
    ]
    business_suite_module_terms = [
        "connectors",
        "intake",
        "crm",
        "reporting",
        "lead",
        "sales",
        "roi",
        "workshop",
        "operator training",
        "content",
    ]
    if has_any(lower, business_suite_terms) and (
        has_any(lower, business_suite_module_terms)
        or has_any(lower, ["full", "setup", "business owner", "business stack", "growth engine"])
    ):
        return "business_suite"

    direct_coding_terms = [
        "write a production-ready typescript",
        "write a robust typescript",
        "write a safe node.js",
        "write a safe nodejs",
        "write a node.js",
        "write a typescript",
        "token bucket",
        "sse parser",
        "streaming response",
        "artifact writer",
        "path traversal",
        "design and implement a typescript",
        "implementation, types, usage example",
        "small vitest test suite",
    ]
    if has_any(lower, direct_coding_terms):
        return "coding"

    technical_plan_terms = [
        "diagnose",
        "likely root causes",
        "patch plan",
        "verification checklist",
        "release checklist",
        "test plan",
        "pseudo-code",
        "pseudocode",
        "state model",
        "workflow to upload",
        "computer-use workflow",
        "backend for",
        "safe web-search proxy",
        "telegram coding-agent bridge",
        "local model fleet router",
        "fleet router",
        "license gate",
        "update-check service",
        "safe update service",
        "artifact-writing layer",
        "artifact writing layer",
        "provide file structure",
    ]
    if has_any(lower, technical_plan_terms):
        return "business_document"

    product_app_terms = [
        "quote builder",
        "estimate builder",
        "quote app",
        "estimate app",
        "invoice app",
        "invoice builder",
        "invoice tracker",
        "invoice generator",
        "content calendar",
        "content planner",
        "expense tracker",
        "budget tracker",
        "lead tracker",
        "task board",
    ]

    if has_any(lower, ["next.js", "nextjs", "starter project", "multi-file", "full project", "repo", "repository"]) and has_any(lower, product_app_terms):
        return "code_project"

    if has_any(lower, product_app_terms):
        return "app"

    if has_any(lower, ["invoice", "proposal", "quote", "launch plan", "marketing plan", "email sequence", "welcome email", "follow-up email", "business brief"]):
        return "business_document"

    if has_any(lower, ["existing repo", "patch this repo", "fix this repo", "modify this repo", "add to this repo"]):
        return "repo_patch"

    if has_any(lower, ["next.js", "nextjs", "starter project", "multi-file", "full project", "repo", "repository", "stripe checkout", "webhook", "api route", "cloudflare worker", "worker api", "wrangler", "d1", "r2", "durable object"]):
        return "code_project"

    if has_any(lower, ["booking app", "tracker app", "crm", "lead tracker", "invoice tracker", "inventory", "task board", "local app", "simple app", "tool to track", "estimate builder", "quote builder", "quote app", "content calendar", "content planner", "expense tracker", "budget tracker"]):
        return "app"

    if has_any(lower, ["website", "landing page", "homepage", "one-page site", "web page", "site for"]):
        return "website"

    return "business_document"


def safe_read(path: Path | None, limit: int = 80_000) -> str:
    if path is None or not path.exists() or not path.is_file():
        return ""
    content = path.read_text(encoding="utf-8")
    return content[:limit]


def contains_forbidden_token(text: str) -> bool:
    lower = text.lower()
    return any(token.lower() in lower for token in FORBIDDEN_TOKENS)


def open_command_for(task_type: str, artifact_path: Path | None, project_dir: Path | None) -> str | None:
    if artifact_path and artifact_path.suffix.lower() == ".html":
        return f"open {artifact_path}"
    if task_type == "business_suite" and project_dir:
        readme = project_dir / "README.md"
        return f"open {readme if readme.exists() else project_dir}"
    if task_type == "business_document" and artifact_path:
        return f"open {artifact_path}"
    if task_type == "coding" and artifact_path:
        return f"open {artifact_path}"
    if project_dir:
        return f"cd {project_dir} && npm install && npm run test"
    return None


def write_manifest(path: Path, manifest: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(manifest, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")


def request_planner_spec(
    *,
    task_type: str,
    prompt: str,
    openai_base_url: str | None,
    model: str | None,
    api_key_env: str,
    timeout: int,
    max_tokens: int = 224,
) -> tuple[dict[str, Any] | None, list[str], float]:
    if not openai_base_url or not model:
        return None, [], 0.0
    started = time.time()
    prompt_file = PROMPT_FILES.get(task_type)
    try:
        raw_spec = request_json_spec(
            base_url=openai_base_url,
            model=model,
            prompt=prompt,
            api_key_env=api_key_env,
            system_prompt_file=prompt_file if prompt_file and prompt_file.exists() else None,
            timeout=timeout,
            max_tokens=max_tokens,
            temperature=0.0,
        )
        return raw_spec, [], round(time.time() - started, 3)
    except Exception as exc:
        return None, [f"model spec planner failed, used deterministic fallback: {exc}"], round(time.time() - started, 3)


def build_manifest(
    *,
    task_type: str,
    prompt: str,
    spec: dict[str, Any],
    artifact_type: str,
    artifact_path: Path | None,
    project_dir: Path | None,
    changed_files: list[str],
    verification_results: list[dict[str, Any]],
    plan_source: str,
    planner_elapsed_s: float,
    planner_warnings: list[str],
    errors: list[str],
    elapsed_s: float,
) -> dict[str, Any]:
    open_command = open_command_for(task_type, artifact_path, project_dir)
    checks_run = [item["name"] for item in verification_results]
    return {
        "router": "kaiju_unified_router_v0",
        "status": "passed" if not errors else "failed",
        "task_type": task_type,
        "artifact_type": artifact_type,
        "prompt": prompt,
        "spec": spec,
        "plan_source": plan_source,
        "planner_elapsed_s": planner_elapsed_s,
        "planner_warnings": planner_warnings,
        "artifact_path": str(artifact_path) if artifact_path else None,
        "project_dir": str(project_dir) if project_dir else None,
        "changed_files": changed_files,
        "checks_run": checks_run + ["manifest_write"],
        "verification_results": verification_results,
        "errors": errors,
        "elapsed_s": round(elapsed_s, 3),
        "open_command": open_command,
    }


def run_task(
    prompt: str,
    out_dir: Path,
    repo: Path | None = None,
    kind: str = "auto",
    openai_base_url: str | None = None,
    model: str | None = None,
    api_key_env: str = "KAIJU_EVAL_API_KEY",
    planner_timeout: int = 90,
    planner_max_tokens: int = 224,
) -> RouterResult:
    started = time.time()
    task_type = route_prompt(prompt, repo=repo, kind=kind)
    raw_spec, planner_warnings, planner_elapsed_s = request_planner_spec(
        task_type=task_type,
        prompt=prompt,
        openai_base_url=openai_base_url,
        model=model,
        api_key_env=api_key_env,
        timeout=planner_timeout,
        max_tokens=planner_max_tokens,
    )
    plan_source = "model_spec" if raw_spec else "deterministic"
    run_dir = out_dir / f"{now_stamp()}-{task_type}-{slugify(prompt)}"
    run_dir.mkdir(parents=True, exist_ok=True)
    artifact_path: Path | None = None
    project_dir: Path | None = None
    changed_files: list[str] = []
    spec: dict[str, Any] = {}
    response_text = ""
    errors: list[str] = []
    artifact_type = "unknown"

    if task_type == "website":
        if raw_spec:
            website_spec = website.normalize_spec(raw_spec, prompt)
            rendered = website.render_html(website_spec, prompt)
            harness_errors = website.validate_html(rendered, website_spec)
        else:
            website_spec, rendered, harness_errors = website.render_from_prompt(prompt)
        artifact_path = run_dir / "index.html"
        website.write_html(artifact_path, rendered)
        spec = website_spec.__dict__
        artifact_type = "html_website"
        changed_files = ["index.html"]
        response_text = rendered
        errors.extend(harness_errors)
    elif task_type == "business_document":
        if raw_spec:
            doc_spec = business.normalize_spec(raw_spec, prompt)
            rendered = business.render_markdown(doc_spec, prompt)
            harness_errors = business.validate_markdown(rendered, doc_spec)
        else:
            doc_spec, rendered, harness_errors = business.render_from_prompt(prompt)
        artifact_path = run_dir / f"{business.slugify(doc_spec.business_name) if hasattr(business, 'slugify') else 'document'}.md"
        business.write_markdown(artifact_path, rendered)
        spec = doc_spec.__dict__
        artifact_type = "markdown_document"
        changed_files = [artifact_path.name]
        response_text = rendered
        errors.extend(harness_errors)
    elif task_type == "business_suite":
        if raw_spec:
            suite_spec, files = business_suite.render_files(raw_spec, prompt)
            harness_errors = business_suite.validate_files(files, suite_spec)
        else:
            suite_spec, files, harness_errors = business_suite.render_from_prompt(prompt)
        project_dir = run_dir / "business-suite"
        if not harness_errors:
            business_suite.write_project(project_dir, files)
        spec = suite_spec.__dict__
        artifact_type = "business_owner_suite"
        changed_files = sorted(files)
        artifact_path = project_dir / "kaiju-change-summary.md"
        response_text = files.get("kaiju-change-summary.md", "")
        errors.extend(harness_errors)
    elif task_type == "coding":
        if raw_spec:
            coding_spec = coding.normalize_spec(raw_spec, prompt)
            rendered = coding.render_markdown(coding_spec, prompt)
            harness_errors = coding.validate_markdown(rendered, coding_spec)
        else:
            coding_spec, rendered, harness_errors = coding.render_from_prompt(prompt)
        artifact_path = run_dir / f"{coding.slugify(coding_spec.title)}.md"
        coding.write_markdown(artifact_path, rendered)
        spec = coding_spec.__dict__
        artifact_type = "markdown_code"
        changed_files = [artifact_path.name]
        response_text = rendered
        errors.extend(harness_errors)
    elif task_type == "app":
        if raw_spec:
            app_spec = app.normalize_spec(raw_spec, prompt)
            rendered = app.render_html(app_spec, prompt)
            harness_errors = app.validate_html(rendered, app_spec)
        else:
            app_spec, rendered, harness_errors = app.render_from_prompt(prompt)
        artifact_path = run_dir / "index.html"
        app.write_html(artifact_path, rendered)
        spec = app_spec.__dict__
        artifact_type = "html_app"
        changed_files = ["index.html"]
        response_text = rendered
        errors.extend(harness_errors)
    elif task_type == "code_project":
        if raw_spec:
            project_spec, files = code_project.render_files(raw_spec, prompt)
            harness_errors = code_project.validate_files(files, project_spec)
        else:
            project_spec, files, harness_errors = code_project.render_from_prompt(prompt)
        project_dir = run_dir / "project"
        if not harness_errors:
            code_project.write_project(project_dir, files)
        spec = project_spec.__dict__
        artifact_type = "next_project"
        changed_files = sorted(files)
        artifact_path = project_dir / "kaiju-change-summary.md"
        response_text = files.get("kaiju-change-summary.md", "") + "\n\n" + files.get("kaiju.patch", "")
        errors.extend(harness_errors)
    elif task_type == "repo_patch":
        if repo is None:
            errors.append("repo path required for repo_patch tasks")
            project_dir = None
            artifact_type = "repo_patch"
        else:
            patch_result = repo_patch.run_patch(repo, prompt, apply=True, raw_spec=raw_spec)
            project_dir = repo.resolve()
            artifact_path = project_dir / "kaiju-repo-patch-summary.md"
            spec = patch_result.spec.__dict__
            artifact_type = "repo_patch"
            changed_files = patch_result.changed_files
            response_text = patch_result.summary_text + "\n\n" + patch_result.patch_text
            errors.extend(patch_result.errors)
    else:
        errors.append(f"unsupported task type: {task_type}")

    elapsed_s = time.time() - started
    manifest_path = run_dir / "kaiju-manifest.json"
    output_changed_files = changed_files + ["kaiju-manifest.json"]
    verification_results = verify_output(
        task_type=task_type,
        artifact_path=artifact_path,
        project_dir=project_dir,
        changed_files=output_changed_files,
        response_text=response_text,
        spec=spec,
    )
    errors.extend(failed_checks(verification_results))
    errors = list(dict.fromkeys(errors))
    manifest = build_manifest(
        task_type=task_type,
        prompt=prompt,
        spec=spec,
        artifact_type=artifact_type,
        artifact_path=artifact_path,
        project_dir=project_dir,
        changed_files=output_changed_files,
        verification_results=verification_results,
        plan_source=plan_source,
        planner_elapsed_s=planner_elapsed_s,
        planner_warnings=planner_warnings,
        errors=errors,
        elapsed_s=elapsed_s,
    )
    write_manifest(manifest_path, manifest)
    return RouterResult(
        task_type=task_type,
        artifact_type=artifact_type,
        artifact_path=artifact_path,
        project_dir=project_dir,
        changed_files=manifest["changed_files"],
        spec=spec,
        manifest_path=manifest_path,
        manifest=manifest,
        response_text=response_text,
        errors=errors,
    )


def result_to_json(result: RouterResult) -> str:
    return json.dumps(result.manifest, indent=2, ensure_ascii=False)