pi-todo-herdr/src/tools.ts

import { StringEnum } from "@earendil-works/pi-ai";
import type { ExtensionAPI, ExtensionContext, Theme } from "@earendil-works/pi-coding-agent";
import { Text } from "@earendil-works/pi-tui";
import { Type } from "typebox";
import { formatCompactTaskTree, TASK_PROMPT_GUIDELINES } from "./prompt.ts";
import type { TaskState } from "./state.ts";
import type { AddTaskDraft, Task, TaskPatch, TaskToolDetails } from "./types.ts";
import { TaskValidationError } from "./types.ts";
import { renderToolResult } from "./ui.ts";

const TaskStatusSchema = StringEnum(["pending", "in-progress", "blocked", "done"] as const, {
  description: "Task status. Defaults to pending.",
});

const SetTaskModeSchema = StringEnum(["append", "replace"] as const);
const RemovalReasonSchema = StringEnum(["mistaken", "duplicate", "canceled"] as const);

/**
 * TypeBox's Type.Cyclic emits bare local refs ("AddTask") instead of URI
 * references ("#/$defs/AddTask"). Strict JSON Schema gateways (e.g. the
 * moonshot-flavored validator behind Console Go) reject those with a 400.
 * Rewrite every local $ref in place to the #/$defs/ form.
 */
function normalizeLocalRefs<T>(schema: T): T {
  if (Array.isArray(schema)) {
    for (const item of schema) normalizeLocalRefs(item);
    return schema;
  }
  if (schema && typeof schema === "object") {
    for (const [key, value] of Object.entries(schema)) {
      if (key === "$ref" && typeof value === "string" && !value.startsWith("#") && !value.includes("://")) {
        (schema as Record<string, unknown>)[key] = `#/$defs/${value}`;
      } else {
        normalizeLocalRefs(value);
      }
    }
  }
  return schema;
}

const AddTaskSchema = Type.Cyclic(
  {
    AddTask: Type.Object(
      {
        name: Type.String({ minLength: 1 }),
        description: Type.Optional(
          Type.String({
            minLength: 1,
            description: "Detailed scope, constraints, and done criteria. Include when useful.",
          }),
        ),
        status: Type.Optional(TaskStatusSchema),
        parentId: Type.Optional(
          Type.Integer({ minimum: 1, description: "Existing parent ID; append mode and top-level tasks only." }),
        ),
        children: Type.Optional(Type.Array(Type.Ref("AddTask"))),
      },
      { additionalProperties: false },
    ),
  },
  "AddTask",
);

export const SetTaskParams = normalizeLocalRefs(
  Type.Object(
    {
      mode: SetTaskModeSchema,
      tasks: Type.Array(AddTaskSchema, {
        description: "Tasks to append or install. May be empty only in replace mode.",
      }),
      force: Type.Optional(
        Type.Boolean({ description: "Use only when the user explicitly abandons unfinished tasks." }),
      ),
    },
    { additionalProperties: false },
  ),
);

const TaskPatchSchema = Type.Object(
  {
    id: Type.Integer({ minimum: 1 }),
    name: Type.Optional(Type.String({ minLength: 1 })),
    description: Type.Optional(Type.String({ minLength: 1 })),
    status: Type.Optional(TaskStatusSchema),
    parentId: Type.Optional(
      Type.Union([Type.Integer({ minimum: 1 }), Type.Null()], {
        description: "New parent ID, or null to move the task to the root.",
      }),
    ),
    position: Type.Optional(
      Type.Integer({ minimum: 0, description: "Zero-based insertion index among the task's final siblings." }),
    ),
  },
  { additionalProperties: false },
);

export const UpdateTaskParams = Type.Object(
  {
    tasks: Type.Array(TaskPatchSchema, { minItems: 1, description: "Task patches applied atomically in array order." }),
  },
  { additionalProperties: false },
);

const TaskIdsSchema = Type.Array(Type.Integer({ minimum: 1 }), {
  minItems: 1,
  description: "Task IDs to process. Duplicate IDs are deduplicated.",
});

export const RemoveTaskParams = Type.Object(
  {
    ids: TaskIdsSchema,
    reason: RemovalReasonSchema,
  },
  { additionalProperties: false },
);

export const IdsParams = Type.Object({ ids: TaskIdsSchema }, { additionalProperties: false });

export const ListTaskParams = Type.Object({}, { additionalProperties: false });

interface SetTaskInput {
  mode: "append" | "replace";
  tasks: AddTaskDraft[];
  force?: boolean;
}

interface UpdateTaskInput {
  tasks: TaskPatch[];
}

interface IdsInput {
  ids: number[];
}

interface RemoveTaskInput extends IdsInput {
  reason: "mistaken" | "duplicate" | "canceled";
}

type ChangeHandler = (ctx: ExtensionContext) => void;

function taskIds(tasks: Task[]): string {
  return tasks.map((task) => task.id).join(", ");
}

function details(action: TaskToolDetails["action"], state: TaskState, affected: Task[]): TaskToolDetails {
  return { action, snapshot: state.snapshot(), affected };
}

function mutationText(verb: string, affected: Task[], state: TaskState): string {
  const stats = state.stats();
  return `${verb} ${affected.length} task${affected.length === 1 ? "" : "s"}: ${taskIds(affected)}. ${stats.done}/${stats.total} done.`;
}

function setTaskText(mode: SetTaskInput["mode"], affected: Task[], state: TaskState): string {
  if (mode === "append") return mutationText("Added", affected, state);
  if (affected.length === 0) return "Cleared task tree. Task IDs reset to 1.";
  const stats = state.stats();
  return `Replaced task tree with ${affected.length} task${affected.length === 1 ? "" : "s"}. ${stats.done}/${stats.total} done. Task IDs restart at 1.`;
}

function throwToolError(label: string, error: unknown): never {
  if (error instanceof TaskValidationError) {
    const count = error.errors.length;
    throw new Error(
      `${label} failed · ${count} validation error${count === 1 ? "" : "s"}\n${error.errors.map((item) => `- ${item}`).join("\n")}`,
    );
  }
  throw error;
}

function renderCall(title: string, count: number | undefined, theme: Theme): Text {
  const suffix = count === undefined ? "" : theme.fg("dim", ` (${count})`);
  return new Text(theme.fg("toolTitle", theme.bold(title)) + suffix, 0, 0);
}

function notifyChange(onChange: ChangeHandler, ctx: ExtensionContext): void {
  try {
    onChange(ctx);
  } catch (error) {
    console.error("[pi-todo-herdr] failed to refresh task UI:", error);
  }
}

export function registerTaskTools(pi: ExtensionAPI, state: TaskState, onChange: ChangeHandler): void {
  pi.registerTool({
    name: "set_tasks",
    label: "Set Tasks",
    description:
      "Atomically append tasks or replace the task tree. Replace may clear with an empty array; force is required to discard unfinished tasks.",
    promptGuidelines: [...TASK_PROMPT_GUIDELINES],
    parameters: SetTaskParams,
    async execute(_toolCallId, params: SetTaskInput, _signal, _onUpdate, ctx) {
      try {
        const affected = state.set(params.mode, params.tasks, params.force);
        notifyChange(onChange, ctx);
        return {
          content: [{ type: "text" as const, text: setTaskText(params.mode, affected, state) }],
          details: details(params.mode, state, affected),
        };
      } catch (error) {
        throwToolError("Set", error);
      }
    },
    renderCall(args: SetTaskInput, theme: Theme) {
      return renderCall(`Set Tasks · ${args.mode}`, args.tasks?.length, theme);
    },
    renderResult(result, { expanded }, theme) {
      return renderToolResult(result, expanded, theme);
    },
  });

  pi.registerTool({
    name: "update_task",
    label: "Update Tasks",
    description:
      "Atomically patch one or more tasks. Fields are optional except id; parentId null moves to root and position is a zero-based sibling index.",
    parameters: UpdateTaskParams,
    async execute(_toolCallId, params: UpdateTaskInput, _signal, _onUpdate, ctx) {
      try {
        const affected = state.update(params.tasks);
        notifyChange(onChange, ctx);
        return {
          content: [{ type: "text" as const, text: mutationText("Updated", affected, state) }],
          details: details("update", state, affected),
        };
      } catch (error) {
        throwToolError("Update", error);
      }
    },
    renderCall(args: UpdateTaskInput, theme: Theme) {
      return renderCall("Update Tasks", args.tasks?.length, theme);
    },
    renderResult(result, { expanded }, theme) {
      return renderToolResult(result, expanded, theme);
    },
  });

  pi.registerTool({
    name: "rm_task",
    label: "Remove Tasks",
    description:
      "Remove tasks only when mistaken, duplicate, or canceled. Removing a parent also removes its descendant subtree.",
    parameters: RemoveTaskParams,
    async execute(_toolCallId, params: RemoveTaskInput, _signal, _onUpdate, ctx) {
      try {
        const affected = state.remove(params.ids);
        notifyChange(onChange, ctx);
        return {
          content: [
            {
              type: "text" as const,
              text: `${mutationText("Removed", affected, state)} Reason: ${params.reason}.`,
            },
          ],
          details: details("remove", state, affected),
        };
      } catch (error) {
        throwToolError("Remove", error);
      }
    },
    renderCall(args: RemoveTaskInput, theme: Theme) {
      return renderCall(`Remove Tasks · ${args.reason}`, args.ids?.length, theme);
    },
    renderResult(result, { expanded }, theme) {
      return renderToolResult(result, expanded, theme);
    },
  });

  pi.registerTool({
    name: "list_task",
    label: "List Tasks",
    description: "List the full task tree compactly with id, name, and status.",
    parameters: ListTaskParams,
    async execute() {
      const affected = state.list();
      return {
        content: [{ type: "text" as const, text: formatCompactTaskTree(affected) }],
        details: details("list", state, affected),
      };
    },
    renderCall(_args, theme: Theme) {
      return renderCall("List Tasks", undefined, theme);
    },
    renderResult(result, { expanded }, theme) {
      return renderToolResult(result, expanded, theme);
    },
  });

  pi.registerTool({
    name: "get_task",
    label: "Get Tasks",
    description: "Get complete details for one or more task IDs, including parent and child IDs.",
    parameters: IdsParams,
    async execute(_toolCallId, params: IdsInput) {
      const uniqueIds = [...new Set(params.ids)];
      const affected = state.get(uniqueIds);
      const found = new Set(affected.map((task) => task.id));
      const missing = uniqueIds.filter((id) => !found.has(id));
      if (missing.length > 0) {
        throwToolError("Get", new TaskValidationError([`task ids do not exist: ${missing.join(", ")}`]));
      }
      return {
        content: [{ type: "text" as const, text: JSON.stringify(affected) }],
        details: details("get", state, affected),
      };
    },
    renderCall(args: IdsInput, theme: Theme) {
      return renderCall("Get Tasks", args.ids?.length, theme);
    },
    renderResult(result, { expanded }, theme) {
      return renderToolResult(result, expanded, theme);
    },
  });
}