@aaroncql/pim-agent 0.0.1

package/src/extensions/todo/todo.ts

@@ -0,0 +1,188 @@
+import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
+import { STATUSES, type TodoItem, type TodoStatus } from "./schema";
+
+export const TODO_STATE_CUSTOM_TYPE = "pim-todo-state";
+
+export type TodoSummary = Record<TodoStatus, number>;
+
+export type TodoDetails = {
+  readonly todos: readonly TodoItem[];
+  readonly summary: TodoSummary;
+};
+
+export type FormatChecklistOptions = {
+  readonly activeOnly?: boolean;
+};
+
+// Identity key for the per-session state slot. Extracted from ExtensionContext
+// because ReadonlySessionManager isn't on the package's public entry point.
+// Only identity is used (no methods called) — WeakMap reclaims the slot when
+// the session is disposed.
+export type TodoSessionKey = ExtensionContext["sessionManager"];
+
+const itemsBySession = new WeakMap<TodoSessionKey, TodoItem[]>();
+
+const markers: Record<TodoStatus, string> = {
+  pending: "[ ]",
+  in_progress: "[>]",
+  completed: "[x]",
+  cancelled: "[~]",
+};
+
+const statusSet: ReadonlySet<TodoStatus> = new Set(STATUSES);
+
+export function getCurrentItems(
+  sessionManager: TodoSessionKey
+): readonly TodoItem[] {
+  return itemsBySession.get(sessionManager) ?? [];
+}
+
+export function replaceItems(
+  sessionManager: TodoSessionKey,
+  items: readonly TodoItem[]
+): readonly TodoItem[] {
+  const normalized = normalizeItems(items);
+  itemsBySession.set(sessionManager, normalized);
+  return normalized;
+}
+
+export function resetItems(sessionManager: TodoSessionKey): void {
+  itemsBySession.set(sessionManager, []);
+}
+
+export function reconstructFromBranch(
+  sessionManager: TodoSessionKey,
+  branch: readonly unknown[]
+): readonly TodoItem[] {
+  const items = findLatestTodoItems(branch);
+  itemsBySession.set(sessionManager, items);
+  return items;
+}
+
+export function normalizeItems(items: readonly TodoItem[]): TodoItem[] {
+  return items.flatMap((item) => {
+    const content = normalizeContent(item.content);
+    return content ? [{ content, status: item.status }] : [];
+  });
+}
+
+export function hasActiveItems(items: readonly TodoItem[]): boolean {
+  return items.some(isActive);
+}
+
+export function summarizeItems(items: readonly TodoItem[]): TodoSummary {
+  const summary: TodoSummary = {
+    pending: 0,
+    in_progress: 0,
+    completed: 0,
+    cancelled: 0,
+  };
+  for (const item of items) {
+    summary[item.status] += 1;
+  }
+  return summary;
+}
+
+export function makeDetails(items: readonly TodoItem[]): TodoDetails {
+  return {
+    todos: structuredClone(items),
+    summary: summarizeItems(items),
+  };
+}
+
+export function formatUpdateSummary(items: readonly TodoItem[]): string {
+  const summary = summarizeItems(items);
+  const segments = [
+    summary.completed > 0 ? `${summary.completed} completed` : undefined,
+    summary.in_progress > 0 ? `${summary.in_progress} in progress` : undefined,
+    summary.pending > 0 ? `${summary.pending} pending` : undefined,
+    summary.cancelled > 0 ? `${summary.cancelled} cancelled` : undefined,
+  ].filter((segment) => segment !== undefined);
+  return segments.length === 0
+    ? "Todos cleared."
+    : `Todos updated: ${segments.join(", ")}.`;
+}
+
+export function formatChecklist(
+  items: readonly TodoItem[],
+  options: FormatChecklistOptions = {}
+): string {
+  return items
+    .filter((item) => !options.activeOnly || isActive(item))
+    .map((item) => `${markers[item.status]} ${item.content}`)
+    .join("\n");
+}
+
+function isActive(item: TodoItem): boolean {
+  return item.status === "pending" || item.status === "in_progress";
+}
+
+function normalizeContent(content: string): string {
+  return content.trim().replaceAll(/\s+/g, " ");
+}
+
+function findLatestTodoItems(branch: readonly unknown[]): TodoItem[] {
+  for (let i = branch.length - 1; i >= 0; i--) {
+    const items = extractTodoItems(branch[i]);
+    if (items) {
+      return items;
+    }
+  }
+  return [];
+}
+
+function extractTodoItems(entry: unknown): TodoItem[] | undefined {
+  if (!isRecord(entry)) {
+    return undefined;
+  }
+  if (entry.type === "message") {
+    const message = entry.message;
+    if (
+      !isRecord(message) ||
+      message.role !== "toolResult" ||
+      message.toolName !== "todo"
+    ) {
+      return undefined;
+    }
+    const details = message.details;
+    if (!isRecord(details) || !Array.isArray(details.todos)) {
+      return undefined;
+    }
+    return normalizeUnknownItems(details.todos);
+  }
+  if (entry.type === "custom" && entry.customType === TODO_STATE_CUSTOM_TYPE) {
+    const data = entry.data;
+    if (!isRecord(data) || !Array.isArray(data.todos)) {
+      return undefined;
+    }
+    return normalizeUnknownItems(data.todos);
+  }
+  return undefined;
+}
+
+function normalizeUnknownItems(items: readonly unknown[]): TodoItem[] {
+  const out: TodoItem[] = [];
+  for (const item of items) {
+    if (!isRecord(item)) {
+      continue;
+    }
+    const { content, status } = item;
+    if (typeof content !== "string" || !isStatus(status)) {
+      continue;
+    }
+    const normalizedContent = normalizeContent(content);
+    if (!normalizedContent) {
+      continue;
+    }
+    out.push({ content: normalizedContent, status });
+  }
+  return out;
+}
+
+function isStatus(value: unknown): value is TodoStatus {
+  return typeof value === "string" && statusSet.has(value as TodoStatus);
+}
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+  return typeof value === "object" && value !== null && !Array.isArray(value);
+}

package/src/extensions/tps/index.test.ts

@@ -0,0 +1,254 @@
+import { afterEach, beforeEach, describe, expect, test } from "bun:test";
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import { PimSettings } from "../../shared/PimSettings";
+import registerTps from "./index";
+
+type Handler = (event: unknown, ctx: unknown) => unknown;
+
+type MockPi = {
+  readonly api: ExtensionAPI;
+  readonly handlers: Map<string, Handler[]>;
+};
+
+const originalNow = Date.now;
+const originalGet = PimSettings.get;
+
+let now = 0;
+
+function createPi(): MockPi {
+  const handlers = new Map<string, Handler[]>();
+  const api = {
+    on(event: string, handler: Handler): void {
+      const existing = handlers.get(event) ?? [];
+      existing.push(handler);
+      handlers.set(event, existing);
+    },
+    registerCommand(): void {},
+  } as unknown as ExtensionAPI;
+
+  return { api, handlers };
+}
+
+async function emit(
+  pi: MockPi,
+  event: string,
+  payload: unknown,
+  ctx: unknown
+): Promise<void> {
+  for (const handler of pi.handlers.get(event) ?? []) {
+    await handler(payload, ctx);
+  }
+}
+
+const assistantMessage = {
+  role: "assistant",
+  content: [{ type: "text", text: "hello" }],
+  api: "openai",
+  provider: "openai",
+  model: "test-model",
+  usage: {
+    input: 1000,
+    output: 50,
+    cacheRead: 5000,
+    cacheWrite: 100,
+    totalTokens: 6150,
+    cost: {
+      input: 0,
+      output: 0,
+      cacheRead: 0,
+      cacheWrite: 0,
+      total: 0,
+    },
+  },
+  stopReason: "stop",
+  timestamp: 1000,
+} as const;
+
+describe("tps extension", () => {
+  beforeEach(() => {
+    now = 0;
+    Date.now = () => now;
+    Object.defineProperty(PimSettings, "get", {
+      value: async () => ({ enabled: true }),
+    });
+  });
+
+  afterEach(() => {
+    Date.now = originalNow;
+    Object.defineProperty(PimSettings, "get", { value: originalGet });
+  });
+
+  test("reports metrics when stream updates and final message are different objects", async () => {
+    const pi = createPi();
+    const notifications: string[] = [];
+    const ctx = {
+      hasUI: true,
+      ui: {
+        notify(message: string): void {
+          notifications.push(message);
+        },
+      },
+    };
+
+    registerTps(pi.api);
+
+    await emit(pi, "agent_start", { type: "agent_start" }, ctx);
+
+    now = 1000;
+    await emit(
+      pi,
+      "before_provider_request",
+      { type: "before_provider_request", payload: {} },
+      ctx
+    );
+
+    now = 1150;
+    await emit(
+      pi,
+      "message_update",
+      {
+        type: "message_update",
+        message: { ...assistantMessage },
+        assistantMessageEvent: {
+          type: "text_start",
+          contentIndex: 0,
+          partial: assistantMessage,
+        },
+      },
+      ctx
+    );
+
+    now = 1200;
+    await emit(
+      pi,
+      "message_update",
+      {
+        type: "message_update",
+        message: { ...assistantMessage },
+        assistantMessageEvent: {
+          type: "text_delta",
+          contentIndex: 0,
+          delta: "h",
+          partial: assistantMessage,
+        },
+      },
+      ctx
+    );
+
+    now = 2200;
+    await emit(
+      pi,
+      "message_end",
+      { type: "message_end", message: assistantMessage },
+      ctx
+    );
+    await emit(
+      pi,
+      "agent_end",
+      { type: "agent_end", messages: [assistantMessage] },
+      ctx
+    );
+
+    expect(notifications).toEqual([
+      "Decode: 50.0 tps | Prefill: 5500.0 tps | Cache read: 5,000 | TTFT: 0.20s",
+    ]);
+  });
+
+  test("reports once at the end of a multi-turn agent cycle", async () => {
+    const pi = createPi();
+    const notifications: string[] = [];
+    const ctx = {
+      hasUI: true,
+      ui: {
+        notify(message: string): void {
+          notifications.push(message);
+        },
+      },
+    };
+
+    registerTps(pi.api);
+
+    await emit(pi, "agent_start", { type: "agent_start" }, ctx);
+
+    now = 1000;
+    await emit(
+      pi,
+      "before_provider_request",
+      { type: "before_provider_request", payload: {} },
+      ctx
+    );
+    now = 1200;
+    await emit(
+      pi,
+      "message_update",
+      {
+        type: "message_update",
+        message: { ...assistantMessage },
+        assistantMessageEvent: {
+          type: "thinking_delta",
+          contentIndex: 0,
+          delta: "thinking",
+          partial: assistantMessage,
+        },
+      },
+      ctx
+    );
+    now = 2200;
+    await emit(
+      pi,
+      "message_end",
+      { type: "message_end", message: assistantMessage },
+      ctx
+    );
+    await emit(
+      pi,
+      "turn_end",
+      { type: "turn_end", message: assistantMessage, toolResults: [] },
+      ctx
+    );
+
+    now = 3000;
+    await emit(
+      pi,
+      "before_provider_request",
+      { type: "before_provider_request", payload: {} },
+      ctx
+    );
+    now = 3200;
+    await emit(
+      pi,
+      "message_update",
+      {
+        type: "message_update",
+        message: { ...assistantMessage },
+        assistantMessageEvent: {
+          type: "text_delta",
+          contentIndex: 0,
+          delta: "h",
+          partial: assistantMessage,
+        },
+      },
+      ctx
+    );
+    now = 4200;
+    await emit(
+      pi,
+      "message_end",
+      { type: "message_end", message: assistantMessage },
+      ctx
+    );
+
+    expect(notifications).toEqual([]);
+
+    await emit(
+      pi,
+      "agent_end",
+      { type: "agent_end", messages: [assistantMessage, assistantMessage] },
+      ctx
+    );
+
+    expect(notifications).toEqual([
+      "Decode: 50.0 tps | Prefill: 5500.0 tps | Cache read: 10,000 | TTFT: 0.20s",
+    ]);
+  });
+});

package/src/extensions/tps/index.ts

@@ -0,0 +1,136 @@
+import type { AssistantMessageEvent } from "@earendil-works/pi-ai";
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import { PimSettings } from "../../shared/PimSettings";
+
+type RequestTiming = {
+  readonly sentMs: number;
+  firstOutputMs: number | null;
+};
+
+function isOutputEvent(event: AssistantMessageEvent): boolean {
+  switch (event.type) {
+    case "text_delta":
+    case "thinking_delta":
+    case "toolcall_delta":
+      return event.delta.length > 0;
+    case "text_end":
+    case "thinking_end":
+      return event.content.length > 0;
+    case "toolcall_end":
+      return true;
+    default:
+      return false;
+  }
+}
+
+export default function (pi: ExtensionAPI): void {
+  pi.registerCommand("tps", {
+    description: "Toggle per-cycle decode/prefill tps reporting",
+    handler: async (_args, ctx) => {
+      const current = await PimSettings.get("tps");
+      const next = { ...current, enabled: !current.enabled };
+      await PimSettings.set("tps", next);
+      ctx.ui.notify(
+        `TPS reporting ${next.enabled ? "enabled" : "disabled"}`,
+        "info"
+      );
+    },
+  });
+
+  let requestTiming: RequestTiming | null = null;
+
+  let promptTokens = 0;
+  let prefillMs = 0;
+  let outputTokens = 0;
+  let decodeMs = 0;
+  let cacheReadTokens = 0;
+  let firstTtftMs: number | null = null;
+
+  pi.on("agent_start", () => {
+    promptTokens = 0;
+    prefillMs = 0;
+    outputTokens = 0;
+    decodeMs = 0;
+    cacheReadTokens = 0;
+    firstTtftMs = null;
+  });
+
+  pi.on("before_provider_request", () => {
+    requestTiming = {
+      sentMs: Date.now(),
+      firstOutputMs: null,
+    };
+  });
+
+  pi.on("message_update", (event) => {
+    if (
+      event.message.role === "assistant" &&
+      requestTiming !== null &&
+      requestTiming.firstOutputMs === null &&
+      isOutputEvent(event.assistantMessageEvent)
+    ) {
+      requestTiming.firstOutputMs = Date.now();
+    }
+  });
+
+  pi.on("message_end", (event) => {
+    if (event.message.role !== "assistant") {
+      return;
+    }
+
+    const timing = requestTiming;
+    const endedMs = Date.now();
+    requestTiming = null;
+    if (timing === null) {
+      return;
+    }
+
+    const usage = event.message.usage;
+    const responseMs = timing.firstOutputMs ?? endedMs;
+    const ttft = responseMs - timing.sentMs;
+    const decode = endedMs - responseMs;
+
+    if (firstTtftMs === null && ttft > 0) {
+      firstTtftMs = ttft;
+    }
+
+    const prefillCounted = (usage.input ?? 0) + (usage.cacheWrite ?? 0);
+    if (prefillCounted > 0 && ttft > 0) {
+      promptTokens += prefillCounted;
+      prefillMs += ttft;
+    }
+    if ((usage.output ?? 0) > 0 && decode > 0) {
+      outputTokens += usage.output;
+      decodeMs += decode;
+    }
+    cacheReadTokens += usage.cacheRead ?? 0;
+  });
+
+  pi.on("agent_end", async (_event, ctx) => {
+    if (!ctx.hasUI) {
+      return;
+    }
+    if (decodeMs <= 0 && prefillMs <= 0) {
+      return;
+    }
+    const { enabled } = await PimSettings.get("tps");
+    if (!enabled) {
+      return;
+    }
+
+    const decodeTps = decodeMs > 0 ? outputTokens / (decodeMs / 1000) : 0;
+    const prefillTps = prefillMs > 0 ? promptTokens / (prefillMs / 1000) : 0;
+    const ttftSec = firstTtftMs !== null ? firstTtftMs / 1000 : 0;
+
+    const parts = [
+      `Decode: ${decodeTps.toFixed(1)} tps`,
+      `Prefill: ${prefillTps.toFixed(1)} tps`,
+    ];
+    if (cacheReadTokens > 0) {
+      parts.push(`Cache read: ${cacheReadTokens.toLocaleString()}`);
+    }
+    parts.push(`TTFT: ${ttftSec.toFixed(2)}s`);
+
+    ctx.ui.notify(parts.join(" | "), "info");
+  });
+}