pi-extended-teams 2.1.16

package/extensions/internal/session-files.ts

@@ -0,0 +1,273 @@
+import * as fs from "node:fs";
+import * as os from "node:os";
+import * as path from "node:path";
+import { getAgentDir } from "@mariozechner/pi-coding-agent";
+import * as paths from "../../src/utils/paths";
+import { hasRecentPrivateAgentSessions } from "./agent-session-files";
+
+export function resolveSkillFile(skillName: string, cwd: string): string {
+  const safeName = paths.sanitizeName(skillName);
+  const candidates = [
+    path.join(cwd, "skills", `${safeName}.md`),
+    path.join(cwd, "skills", safeName, "SKILL.md"),
+    path.join(getAgentDir(), "skills", safeName, "SKILL.md"),
+    path.join(getAgentDir(), "skills", `${safeName}.md`),
+  ];
+  const found = candidates.find(candidate => fs.existsSync(candidate));
+  if (!found) {
+    throw new Error(`Skill ${skillName} not found. Checked project and agent skill directories.`);
+  }
+  return found;
+}
+
+export function getPiSessionId(ctx?: any): string | undefined {
+  const id = ctx?.sessionManager?.getSessionId?.();
+  return typeof id === "string" && id.length > 0 ? id : undefined;
+}
+
+export function findLeadTeamForSession(sessionId?: string): string | null {
+  if (!sessionId) return null;
+  try {
+    const teamsDir = path.dirname(paths.teamDir("__probe__"));
+    if (!fs.existsSync(teamsDir)) return null;
+
+    for (const teamDir of fs.readdirSync(teamsDir)) {
+      try {
+        const sessionFile = paths.leadSessionPath(teamDir);
+        if (!fs.existsSync(sessionFile)) continue;
+        const session = JSON.parse(fs.readFileSync(sessionFile, "utf-8"));
+        if (session.sessionId === sessionId && session.pid === process.pid) {
+          return teamDir;
+        }
+      } catch {
+        // ignore entries that are not valid team names, and invalid session files
+      }
+    }
+  } catch {
+    // ignore errors
+  }
+  return null;
+}
+
+export function registerLeadSession(teamName: string, sessionId?: string) {
+  const sessionFile = paths.leadSessionPath(teamName);
+  const dir = path.dirname(sessionFile);
+  fs.mkdirSync(dir, { recursive: true });
+  fs.writeFileSync(sessionFile, JSON.stringify({
+    pid: process.pid,
+    sessionId,
+    startedAt: Date.now(),
+  }));
+}
+
+function isPidAlive(pid: number): boolean {
+  if (!Number.isFinite(pid) || pid <= 0) return false;
+  try {
+    process.kill(pid, 0);
+    return true;
+  } catch {
+    return false;
+  }
+}
+
+function readJsonFile(filePath: string): any | null {
+  try {
+    if (!fs.existsSync(filePath)) return null;
+    return JSON.parse(fs.readFileSync(filePath, "utf-8"));
+  } catch {
+    return null;
+  }
+}
+
+export interface CleanupPidFileResult {
+  pid?: number;
+  killed: boolean;
+  unlinked: boolean;
+  killError?: unknown;
+}
+
+export function unlinkPidFile(pidFile: string): boolean {
+  if (!fs.existsSync(pidFile)) return false;
+  try {
+    fs.unlinkSync(pidFile);
+    return true;
+  } catch {
+    return false;
+  }
+}
+
+export function cleanupPidFileProcess(
+  pidFile: string,
+  options: { signal?: NodeJS.Signals | number; skipPid?: number } = {}
+): CleanupPidFileResult {
+  const result: CleanupPidFileResult = { killed: false, unlinked: false };
+  if (!fs.existsSync(pidFile)) return result;
+
+  try {
+    const pid = Number.parseInt(fs.readFileSync(pidFile, "utf-8").trim(), 10);
+    if (Number.isFinite(pid) && pid > 0) {
+      result.pid = pid;
+      if (pid !== options.skipPid) {
+        try {
+          process.kill(pid, options.signal ?? "SIGKILL");
+          result.killed = true;
+        } catch (error) {
+          result.killError = error;
+        }
+      }
+    }
+  } catch {
+    // Keep cleanup best-effort: unreadable/corrupt pid files should still be unlinked.
+  } finally {
+    result.unlinked = unlinkPidFile(pidFile);
+  }
+
+  return result;
+}
+
+function killTeamMemberProcesses(teamName: string, config: any, terminal: any): void {
+  for (const member of config?.members || []) {
+    if (member.name === "team-lead") continue;
+
+    const pidFile = path.join(paths.teamDir(teamName), `${member.name}.pid`);
+    cleanupPidFileProcess(pidFile, { skipPid: process.pid });
+
+    if (terminal && member.tmuxPaneId) {
+      try { terminal.kill(member.tmuxPaneId); } catch {}
+    }
+  }
+}
+
+export function forceCleanupTeam(teamName: string, terminal: any): boolean {
+  const config = readJsonFile(paths.configPath(teamName));
+  killTeamMemberProcesses(teamName, config, terminal);
+
+  const teamDirectory = paths.teamDir(teamName);
+  const tasksDirectory = paths.taskDir(teamName);
+  let removed = false;
+  if (fs.existsSync(teamDirectory)) {
+    fs.rmSync(teamDirectory, { recursive: true, force: true });
+    removed = true;
+  }
+  if (fs.existsSync(tasksDirectory)) {
+    fs.rmSync(tasksDirectory, { recursive: true, force: true });
+    removed = true;
+  }
+  return removed;
+}
+
+function teamLastTouchedAt(teamName: string, teamDirectory: string): number {
+  const session = readJsonFile(paths.leadSessionPath(teamName));
+  if (typeof session?.startedAt === "number") return session.startedAt;
+
+  const config = readJsonFile(paths.configPath(teamName));
+  if (typeof config?.createdAt === "number") return config.createdAt;
+
+  try { return fs.statSync(teamDirectory).mtimeMs; } catch { return 0; }
+}
+
+export function cleanupStaleTeam(teamName: string, terminal: any): boolean {
+  const session = readJsonFile(paths.leadSessionPath(teamName));
+  if (!session?.pid) return false;
+  return isPidAlive(Number(session.pid)) ? false : forceCleanupTeam(teamName, terminal);
+}
+
+export interface CleanupOrphanedTeamsOptions {
+  /** How old a team folder without a live lead must be before removal. */
+  maxAgeMs?: number;
+  /** Test seam; defaults to ~/.pi/teams. */
+  teamsRoot?: string;
+  now?: number;
+}
+
+export function cleanupOrphanedTeams(
+  terminal: any,
+  options: CleanupOrphanedTeamsOptions = {}
+): number {
+  const teamsRoot = options.teamsRoot || paths.TEAMS_DIR;
+  if (!fs.existsSync(teamsRoot)) return 0;
+
+  const now = options.now ?? Date.now();
+  const maxAgeMs = options.maxAgeMs ?? 24 * 60 * 60 * 1000;
+  let cleaned = 0;
+
+  for (const dir of fs.readdirSync(teamsRoot)) {
+    const teamDirectory = path.join(teamsRoot, dir);
+    try {
+      if (!fs.statSync(teamDirectory).isDirectory()) continue;
+      paths.sanitizeName(dir);
+
+      const session = readJsonFile(paths.leadSessionPath(dir));
+      if (session?.pid && isPidAlive(Number(session.pid))) continue;
+
+      const missingLiveLead = !session?.pid || !isPidAlive(Number(session.pid));
+      const lastTouchedAt = teamLastTouchedAt(dir, teamDirectory);
+      const oldEnough = maxAgeMs <= 0 || (lastTouchedAt > 0 && (now - lastTouchedAt) > maxAgeMs);
+      const deadLeadPid = !!session?.pid && !isPidAlive(Number(session.pid));
+      const quarantineDir = path.join(teamDirectory, "lifecycle", "quarantine");
+      const hasLifecycleFence = fs.existsSync(quarantineDir)
+        && fs.readdirSync(quarantineDir).some(file => file.endsWith(".json"));
+
+      const shouldCleanOrphan = deadLeadPid || (missingLiveLead && oldEnough);
+
+      // Lifecycle tombstones have no TTL. Even corrupt files are occupied and
+      // must keep the team directory out of orphan cleanup, but they must not
+      // keep child processes alive after their lead is gone.
+      if (hasLifecycleFence) {
+        if (shouldCleanOrphan) {
+          killTeamMemberProcesses(dir, readJsonFile(paths.configPath(dir)), terminal);
+        }
+        continue;
+      }
+
+      if (shouldCleanOrphan) {
+        // Failed and recovery-relevant spawned sessions intentionally outlive
+        // their roster member. Stop orphan processes, but do not let dead-lead
+        // directory cleanup bypass the private transcript TTL.
+        const privateSessionRetention = hasRecentPrivateAgentSessions(teamDirectory, { now, maxAgeMs });
+        if (privateSessionRetention !== "none") {
+          killTeamMemberProcesses(dir, readJsonFile(paths.configPath(dir)), terminal);
+          continue;
+        }
+        if (forceCleanupTeam(dir, terminal)) cleaned++;
+      }
+    } catch {
+      // Ignore malformed team directories; cleanup should never disrupt startup.
+    }
+  }
+
+  return cleaned;
+}
+
+export function cleanupAgentSessionFolders(maxAgeMs: number = 24 * 60 * 60 * 1000): number {
+  const agentTeamsDir = path.join(os.homedir(), ".pi", "agent", "teams");
+  if (!fs.existsSync(agentTeamsDir)) return 0;
+
+  let cleaned = 0;
+  const now = Date.now();
+
+  for (const dir of fs.readdirSync(agentTeamsDir)) {
+    const sessionDir = path.join(agentTeamsDir, dir);
+    const configFile = path.join(sessionDir, "config.json");
+
+    try {
+      // Check if this is a directory with a config.json
+      if (!fs.statSync(sessionDir).isDirectory()) continue;
+      if (!fs.existsSync(configFile)) continue;
+
+      // Read the config to check the creation time
+      const config = JSON.parse(fs.readFileSync(configFile, "utf-8"));
+      const createdAt = config.createdAt ? new Date(config.createdAt).getTime() : 0;
+
+      // If the folder is older than maxAgeMs, delete it
+      if (createdAt > 0 && (now - createdAt) > maxAgeMs) {
+        fs.rmSync(sessionDir, { recursive: true });
+        cleaned++;
+      }
+    } catch {
+      // Ignore errors for individual folders
+    }
+  }
+
+  return cleaned;
+}

package/extensions/internal/session-usage.ts

@@ -0,0 +1,59 @@
+export interface SessionUsageSummary {
+  tokensUsed?: number;
+  costUsd?: number;
+}
+
+function numberOrZero(value: unknown): number {
+  return typeof value === "number" && Number.isFinite(value) ? value : 0;
+}
+
+function assistantMessageKey(message: any): string | undefined {
+  if (!message || message.role !== "assistant") return undefined;
+  const timestamp = numberOrZero(message.timestamp);
+  const provider = typeof message.provider === "string" ? message.provider : "";
+  const model = typeof message.model === "string" ? message.model : "";
+  const stopReason = typeof message.stopReason === "string" ? message.stopReason : "";
+  return `${timestamp}:${provider}:${model}:${stopReason}`;
+}
+
+function usageTokens(usage: any): number {
+  return (
+    numberOrZero(usage?.input) +
+    numberOrZero(usage?.output) +
+    numberOrZero(usage?.cacheRead) +
+    numberOrZero(usage?.cacheWrite)
+  );
+}
+
+function addAssistantUsage(message: any, totals: { tokens: number; cost: number; seen: boolean }): void {
+  if (!message || message.role !== "assistant" || !message.usage) return;
+  totals.tokens += usageTokens(message.usage);
+  totals.cost += numberOrZero(message.usage.cost?.total);
+  totals.seen = true;
+}
+
+export function summarizeSessionUsage(ctx: any, currentAssistantMessage?: any): SessionUsageSummary {
+  const sessionManager = ctx?.sessionManager;
+  const entries = typeof sessionManager?.getBranch === "function"
+    ? sessionManager.getBranch()
+    : typeof sessionManager?.getEntries === "function"
+      ? sessionManager.getEntries()
+      : [];
+
+  const totals = { tokens: 0, cost: 0, seen: false };
+  const currentKey = assistantMessageKey(currentAssistantMessage);
+  let currentAlreadyIncluded = false;
+
+  for (const entry of Array.isArray(entries) ? entries : []) {
+    const message = entry?.type === "message" ? entry.message : entry?.message;
+    if (!message || message.role !== "assistant") continue;
+    if (currentKey && assistantMessageKey(message) === currentKey) currentAlreadyIncluded = true;
+    addAssistantUsage(message, totals);
+  }
+
+  if (currentAssistantMessage && !currentAlreadyIncluded) {
+    addAssistantUsage(currentAssistantMessage, totals);
+  }
+
+  return totals.seen ? { tokensUsed: totals.tokens, costUsd: totals.cost } : {};
+}

package/extensions/resources/spawn-resource-plan.test.ts

@@ -0,0 +1,233 @@
+import fs from "node:fs";
+import os from "node:os";
+import path from "node:path";
+import { afterEach, describe, expect, it } from "vitest";
+import { DEFAULT_SETTINGS } from "../../src/utils/settings";
+import {
+  createSpawnResourcePlan,
+  EXTENSIONS_COMMAND_DESCRIPTION,
+  parentProjectTrustForSpawn,
+  snapshotLeadExtensions,
+  type ExtensionSourceInfo,
+} from "./spawn-resource-plan";
+
+const roots: string[] = [];
+
+function tempRoot(): string {
+  const root = fs.mkdtempSync(path.join(os.tmpdir(), "pet-resource-plan-"));
+  roots.push(root);
+  return root;
+}
+
+function touch(filePath: string): string {
+  fs.mkdirSync(path.dirname(filePath), { recursive: true });
+  fs.writeFileSync(filePath, "export default function () {}\n");
+  return filePath;
+}
+
+function sourceInfo(filePath: string, overrides: Partial<ExtensionSourceInfo> = {}): ExtensionSourceInfo {
+  return {
+    path: filePath,
+    source: "local",
+    scope: "user",
+    origin: "top-level",
+    ...overrides,
+  };
+}
+
+afterEach(() => {
+  for (const root of roots.splice(0)) fs.rmSync(root, { recursive: true, force: true });
+});
+
+describe("lead extension snapshots", () => {
+  it("uses only public command/tool sourceInfo, dedupes canonical paths, and identifies self", () => {
+    const root = tempRoot();
+    const self = touch(path.join(root, "pi-extended-teams", "extensions", "index.ts"));
+    const selected = touch(path.join(root, "selected", "index.ts"));
+    const selectedLink = path.join(root, "selected-link.ts");
+    fs.symlinkSync(selected, selectedLink);
+
+    const snapshot = snapshotLeadExtensions({
+      getCommands: () => [
+        { name: "agents-extensions", source: "extension", sourceInfo: sourceInfo(self) },
+        { name: "selected-command", source: "extension", sourceInfo: sourceInfo(selected) },
+        { name: "skill:anything", source: "skill", sourceInfo: sourceInfo("/skills/anything.md") },
+      ],
+      getAllTools: () => [
+        { name: "read", sourceInfo: sourceInfo("<builtin:read>", { source: "builtin", scope: "temporary" }) },
+        { name: "selected-tool", sourceInfo: sourceInfo(selectedLink) },
+        { name: "other", sourceInfo: sourceInfo(path.join(root, "other.ts"), { source: "npm:other", origin: "package" }) },
+      ],
+    });
+
+    expect(snapshot.map((extension) => extension.path)).toEqual([self, selected, path.join(root, "other.ts")]);
+    expect(snapshot.map((extension) => extension.name)).toEqual(["pi-extended-teams", "selected", "other"]);
+    expect(snapshot.map((extension) => extension.isSelf)).toEqual([true, false, false]);
+    expect(Object.isFrozen(snapshot)).toBe(true);
+    expect(Object.isFrozen(snapshot[0]?.sourceInfo)).toBe(true);
+  });
+
+  it("identifies self by its public command description when invocation names collide", () => {
+    const root = tempRoot();
+    const other = touch(path.join(root, "other.ts"));
+    const self = touch(path.join(root, "pi-extended-teams", "extensions", "index.ts"));
+    const snapshot = snapshotLeadExtensions({
+      getCommands: () => [
+        { name: "agents-extensions", description: "Another command", source: "extension", sourceInfo: sourceInfo(other) },
+        { name: "agents-extensions:2", description: EXTENSIONS_COMMAND_DESCRIPTION, source: "extension", sourceInfo: sourceInfo(self) },
+      ],
+      getAllTools: () => [],
+    });
+
+    expect(snapshot.map((extension) => extension.isSelf)).toEqual([false, true]);
+  });
+
+  it("cannot observe event-only or synthetic inline extensions through the public registration surface", () => {
+    const snapshot = snapshotLeadExtensions({
+      getCommands: () => [],
+      getAllTools: () => [
+        { name: "sdk-tool", sourceInfo: sourceInfo("<sdk:sdk-tool>", { source: "sdk", scope: "temporary" }) },
+      ],
+    });
+
+    expect(snapshot).toEqual([]);
+  });
+});
+
+describe("spawn resource planning", () => {
+  it("preserves null, empty, explicit, and block semantics over the lead snapshot", () => {
+    const root = tempRoot();
+    const self = touch(path.join(root, "pi-extended-teams", "extensions", "index.ts"));
+    const first = touch(path.join(root, "first", "index.ts"));
+    const second = touch(path.join(root, "second.ts"));
+    const blocked = touch(path.join(root, "blocked.ts"));
+    const leadExtensions = snapshotLeadExtensions({
+      getCommands: () => [
+        { name: "agents-extensions", source: "extension", sourceInfo: sourceInfo(self) },
+        { name: "first-command", source: "extension", sourceInfo: sourceInfo(first) },
+      ],
+      getAllTools: () => [
+        { name: "second-tool", sourceInfo: sourceInfo(second) },
+        { name: "blocked-tool", sourceInfo: sourceInfo(blocked) },
+      ],
+    });
+
+    const defaultSettings = structuredClone(DEFAULT_SETTINGS);
+    defaultSettings.extensions.block = ["blocked"];
+    const defaultPlan = createSpawnResourcePlan({
+      cwd: root,
+      projectTrusted: true,
+      settings: defaultSettings,
+      leadExtensions,
+    });
+    expect(defaultPlan.extensionPaths).toEqual([first, second]);
+    expect(defaultPlan.selfExtensionPath).toBe(self);
+    expect(defaultPlan.extensions.map((extension) => extension.state)).toEqual([
+      "self",
+      "selected",
+      "selected",
+      "blocked",
+    ]);
+
+    const noneSettings = structuredClone(DEFAULT_SETTINGS);
+    noneSettings.extensions.allow = [];
+    const none = createSpawnResourcePlan({ cwd: root, projectTrusted: true, settings: noneSettings, leadExtensions });
+    expect(none.selectionMode).toBe("explicit");
+    expect(none.extensionPaths).toEqual([]);
+
+    const explicitSettings = structuredClone(DEFAULT_SETTINGS);
+    explicitSettings.extensions.allow = ["first", second, "blocked", "missing-extension", "pi-extended-teams"];
+    explicitSettings.extensions.block = ["blocked"];
+    const explicit = createSpawnResourcePlan({ cwd: root, projectTrusted: true, settings: explicitSettings, leadExtensions });
+    expect(explicit.extensionPaths).toEqual([first, second]);
+    expect(explicit.diagnostics).toEqual(expect.arrayContaining([
+      expect.objectContaining({ code: "blocked-selection", configuredEntry: "blocked" }),
+      expect.objectContaining({ code: "stale-selection", configuredEntry: "missing-extension" }),
+      expect.objectContaining({ code: "self-excluded", configuredEntry: "pi-extended-teams" }),
+    ]));
+    expect(Object.isFrozen(explicit)).toBe(true);
+    expect(Object.isFrozen(explicit.extensionPaths)).toBe(true);
+  });
+
+  it("uses names when unique and paths when names collide", () => {
+    const root = tempRoot();
+    const first = touch(path.join(root, "a", "same", "index.ts"));
+    const second = touch(path.join(root, "b", "same", "index.ts"));
+    const unique = touch(path.join(root, "unique.ts"));
+    const leadExtensions = snapshotLeadExtensions({
+      getCommands: () => [
+        { name: "one", source: "extension", sourceInfo: sourceInfo(first) },
+        { name: "two", source: "extension", sourceInfo: sourceInfo(second) },
+        { name: "three", source: "extension", sourceInfo: sourceInfo(unique) },
+      ],
+      getAllTools: () => [],
+    });
+
+    const plan = createSpawnResourcePlan({
+      cwd: root,
+      projectTrusted: false,
+      settings: structuredClone(DEFAULT_SETTINGS),
+      leadExtensions,
+    });
+
+    expect(plan.extensions.map((extension) => extension.selector)).toEqual([first, second, "unique"]);
+  });
+
+  it("keeps canonical selections stable when a same-name extension is later observed", () => {
+    const root = tempRoot();
+    const original = touch(path.join(root, "first", "same.ts"));
+    const second = touch(path.join(root, "second", "same.ts"));
+    const oneExtension = snapshotLeadExtensions({
+      getCommands: () => [{ name: "first", source: "extension", sourceInfo: sourceInfo(original) }],
+      getAllTools: () => [],
+    });
+    const settings = structuredClone(DEFAULT_SETTINGS);
+    settings.extensions.allow = [oneExtension[0]!.identity];
+
+    const rebuilt = snapshotLeadExtensions({
+      getCommands: () => [
+        { name: "first", source: "extension", sourceInfo: sourceInfo(original) },
+        { name: "second", source: "extension", sourceInfo: sourceInfo(second) },
+      ],
+      getAllTools: () => [],
+    });
+    const canonical = createSpawnResourcePlan({ cwd: root, projectTrusted: true, settings, leadExtensions: rebuilt });
+    expect(canonical.extensionPaths).toEqual([original]);
+
+    const legacySettings = structuredClone(DEFAULT_SETTINGS);
+    legacySettings.extensions.allow = ["same"];
+    const legacy = createSpawnResourcePlan({ cwd: root, projectTrusted: true, settings: legacySettings, leadExtensions: rebuilt });
+    expect(legacy.extensionPaths).toEqual([original, second]);
+  });
+
+  it("matches canonical and symlink entries while blocks always override allow", () => {
+    const root = tempRoot();
+    const extension = touch(path.join(root, "extension.ts"));
+    const extensionLink = path.join(root, "extension-link.ts");
+    fs.symlinkSync(extension, extensionLink);
+    const leadExtensions = snapshotLeadExtensions({
+      getCommands: () => [{ name: "extension", source: "extension", sourceInfo: sourceInfo(extension) }],
+      getAllTools: () => [],
+    });
+    const settings = structuredClone(DEFAULT_SETTINGS);
+    settings.extensions.allow = [leadExtensions[0]!.identity];
+    settings.extensions.block = [extensionLink];
+
+    const plan = createSpawnResourcePlan({ cwd: root, projectTrusted: true, settings, leadExtensions });
+    expect(plan.extensionPaths).toEqual([]);
+    expect(plan.extensions[0]?.state).toBe("blocked");
+    expect(plan.diagnostics).toEqual([
+      expect.objectContaining({ code: "blocked-selection", configuredEntry: leadExtensions[0]!.identity }),
+    ]);
+  });
+
+  it("never transfers project trust across cwd", () => {
+    const root = tempRoot();
+    const project = path.join(root, "project");
+    const ctx = { cwd: project, isProjectTrusted: () => true };
+
+    expect(parentProjectTrustForSpawn(ctx, project)).toBe(true);
+    expect(parentProjectTrustForSpawn(ctx, path.join(root, "other-project"))).toBe(false);
+    expect(parentProjectTrustForSpawn({ cwd: project }, project)).toBe(false);
+  });
+});