automaton/plugins/pi-automaton-context/index.ts

import type { ExtensionAPI, ExtensionContext, BeforeAgentStartEventResult } from "@earendil-works/pi-coding-agent";
import { existsSync, readFileSync, readdirSync, statSync } from "fs";
import { basename, join } from "path";
import { homedir } from "os";

const AUTOMATON_HOME = join(homedir(), ".automaton");
const STALE_MINUTES = 30;

interface TaskInfo {
  name: string;
  phase: string;
  mtime: Date;
}

function getTasks(autoDir: string): TaskInfo[] {
  const tasksDir = join(autoDir, ".automaton", "tasks");
  if (!existsSync(tasksDir)) return [];
  try {
    const entries = readdirSync(tasksDir);
    const tasks: TaskInfo[] = [];
    for (const entry of entries) {
      const stateFile = join(tasksDir, entry, ".state");
      if (existsSync(stateFile)) {
        try {
          const phase = readFileSync(stateFile, "utf-8").trim();
          if (!phase) continue;
          const stat = statSync(stateFile);
          tasks.push({ name: entry, phase, mtime: stat.mtime });
        } catch {
          // skip unreadable
        }
      }
    }
    tasks.sort((a, b) => b.mtime.getTime() - a.mtime.getTime());
    return tasks;
  } catch {
    return [];
  }
}

function getLoopInfo(autoDir: string): string[] {
  const loopsDir = join(autoDir, ".automaton", "loops");
  if (!existsSync(loopsDir)) return [];
  try {
    const entries = readdirSync(loopsDir);
    const lines: string[] = [];
    for (const entry of entries) {
      const stateFile = join(loopsDir, entry, ".state.loop");
      if (existsSync(stateFile)) {
        try {
          const content = readFileSync(stateFile, "utf-8").trim();
          const state = JSON.parse(content);
          const taskRef = state.current_task ? `, active task: ${state.current_task}` : "";
          lines.push(`Loop "${entry}": ${state.status || "unknown"}${taskRef}`);
        } catch {
          // skip unparseable
        }
      }
    }
    return lines;
  } catch {
    return [];
  }
}

export default function (pi: ExtensionAPI) {
  pi.on("before_agent_start", async (_event, ctx): Promise<BeforeAgentStartEventResult | undefined> => {
    const cwd = ctx.cwd;
    const inFramework = cwd === AUTOMATON_HOME || cwd.startsWith(AUTOMATON_HOME + "/");

    let baseDir: string | null = null;
    let scopeLabel: string;

    if (inFramework) {
      baseDir = AUTOMATON_HOME;
      scopeLabel = "framework (~/.automaton/)";
    } else {
      const projectAuto = join(cwd, ".automaton");
      if (existsSync(projectAuto)) {
        baseDir = cwd;
        const projectName = basename(cwd) || "project";
        scopeLabel = `project (${projectName}/.automaton/)`;
      } else {
        return; // not in automaton context
      }
    }

    const tasks = getTasks(baseDir);
    const editTasks = tasks.filter((t) => t.phase === "implement" || t.phase === "doc_review");
    const currentTask = editTasks.length > 0 ? editTasks[0] : tasks.length > 0 ? tasks[0] : null;

    const lines: string[] = [];
    lines.push(`Automaton scope: ${scopeLabel}`);
    lines.push("");

    lines.push("How Automaton works:");
    lines.push("- Automaton enforces a task phase state machine. File edits are ONLY allowed when a task is in 'implement' or 'doc_review' phase.");
    lines.push("- When the user asks you to build, fix, or change something, you MUST first create a task (use automaton_create_task) and transition it to 'implement' before you can edit any files.");
    lines.push("- Without a task in 'implement', the automaton-guard-pi extension will BLOCK all file edits.");
    lines.push(`- Available tools: automaton_create_task (create+transition), automaton_transition (change phase), automaton_status (check state).`);
    lines.push("");

    if (currentTask) {
      lines.push(`State: Active task "${currentTask.name}" is in "${currentTask.phase}" phase.`);

      if (currentTask.phase === "implement" || currentTask.phase === "doc_review") {
        lines.push(`You CAN edit files under this task.`);
        const ageMinutes = Math.round((Date.now() - currentTask.mtime.getTime()) / 60000);
        if (ageMinutes > STALE_MINUTES) {
          lines.push(
            `Note: task has been in ${currentTask.phase} for ${ageMinutes} min and is considered stale. ` +
            `Run automaton_transition (or touch via status.py) if still active.`,
          );
        }
      } else {
        lines.push(`File edits are BLOCKED. Call automaton_transition to move it to 'implement' before editing.`);
      }
    } else {
      lines.push(`State: No active tasks. When the user makes a work request, call automaton_create_task first.`);
    }

    const byPhase = new Map<string, number>();
    for (const t of tasks) {
      byPhase.set(t.phase, (byPhase.get(t.phase) || 0) + 1);
    }
    const summary = Array.from(byPhase.entries())
      .sort((a, b) => b[1] - a[1])
      .map(([p, c]) => `${p} (${c})`)
      .join(", ");
    lines.push(`All tasks: ${summary || "none"}`);

    const loopInfo = getLoopInfo(baseDir);
    for (const l of loopInfo) {
      lines.push(l);
    }

    return {
      systemPrompt: `${_event.systemPrompt}\n\n## Automaton Context\n\n${lines.join("\n")}`,
    };
  });
}