@aaroncql/pim-agent 0.0.1

package/src/telegram/SessionRegistry.test.ts

@@ -0,0 +1,89 @@
+import type { Api } from "grammy";
+import { mkdtemp, rm } from "node:fs/promises";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import { afterEach, beforeEach, describe, expect, test } from "bun:test";
+
+import { Fs } from "../shared/Fs";
+import { type TelegramConfig } from "./Config";
+import { SessionRegistry } from "./SessionRegistry";
+import type { SessionSettings } from "./Session";
+import { TaskScheduler } from "./TaskScheduler";
+
+let tmp: string;
+let config: TelegramConfig;
+const stubApi = {} as Api;
+const stubScheduler = new TaskScheduler({
+  configDir: "/tmp",
+  runTask: async () => {},
+});
+
+beforeEach(async () => {
+  tmp = await mkdtemp(join(tmpdir(), "pim-session-registry-test-"));
+  config = {
+    token: "token",
+    allow: [],
+    cwd: tmp,
+    configDir: tmp,
+  };
+});
+
+afterEach(async () => {
+  await rm(tmp, { recursive: true, force: true });
+});
+
+async function readState(): Promise<Record<string, SessionSettings>> {
+  return Fs.readJsonOrEmpty<Record<string, SessionSettings>>(
+    join(tmp, "state.json"),
+    {}
+  );
+}
+
+async function writeState(
+  state: Record<string, SessionSettings>
+): Promise<void> {
+  await Fs.writeAtomic(join(tmp, "state.json"), JSON.stringify(state, null, 2));
+}
+
+describe("SessionRegistry state", () => {
+  test("loads persisted state and preserves it when mutating another session", async () => {
+    await writeState({
+      "1-main": {
+        cwd: "/repo",
+        cumulativeCost: 12.5,
+        sessionPath: "/sessions/one.jsonl",
+      },
+    });
+
+    const registry = new SessionRegistry(config, stubApi, stubScheduler);
+    await registry.init();
+    const session = registry.get({ chatId: 2, threadId: undefined });
+    await session.setThinkingLevel("off");
+
+    const loaded = await readState();
+    expect(loaded["1-main"]).toEqual({
+      cwd: "/repo",
+      cumulativeCost: 12.5,
+      sessionPath: "/sessions/one.jsonl",
+    });
+    expect(loaded["2-main"]?.thinkingLevel).toBe("off");
+  });
+
+  test("does not flush state when disposed before init", async () => {
+    await writeState({
+      "1-main": {
+        cwd: "/repo",
+        cumulativeCost: 12.5,
+      },
+    });
+
+    const registry = new SessionRegistry(config, stubApi, stubScheduler);
+    await registry.disposeAll();
+
+    const loaded = await readState();
+    expect(loaded["1-main"]).toEqual({
+      cwd: "/repo",
+      cumulativeCost: 12.5,
+    });
+  });
+});

package/src/telegram/SessionRegistry.ts

@@ -0,0 +1,170 @@
+import {
+  AuthStorage,
+  getAgentDir,
+  ModelRegistry,
+  SettingsManager,
+} from "@earendil-works/pi-coding-agent";
+import type { Api } from "grammy";
+import { mkdir } from "node:fs/promises";
+import { join } from "node:path";
+
+import { Fs } from "../shared/Fs";
+import { type TelegramConfig } from "./Config";
+import { Session, type SessionId, type SessionSettings } from "./Session";
+import type { TaskScheduler } from "./TaskScheduler";
+
+const LRU_CAP = 16;
+
+export class SessionRegistry {
+  private readonly config: TelegramConfig;
+  private readonly api: Api;
+  private readonly scheduler: TaskScheduler;
+  private readonly cache = new Map<string, Session>();
+  private readonly authStorage: AuthStorage;
+  private readonly modelRegistry: ModelRegistry;
+  private readonly settingsManagers = new Map<string, SettingsManager>();
+  private readonly agentDir: string;
+  private settings: Map<string, SessionSettings> = new Map();
+  private initialized = false;
+  private initPromise: Promise<void> | undefined;
+  private botUsername: string | undefined;
+
+  public constructor(
+    config: TelegramConfig,
+    api: Api,
+    scheduler: TaskScheduler
+  ) {
+    this.config = config;
+    this.api = api;
+    this.scheduler = scheduler;
+    this.agentDir = getAgentDir();
+    this.authStorage = AuthStorage.create(join(this.agentDir, "auth.json"));
+    this.modelRegistry = ModelRegistry.create(
+      this.authStorage,
+      join(this.agentDir, "models.json")
+    );
+  }
+
+  public setBotUsername(username: string): void {
+    this.botUsername = username;
+  }
+
+  public async init(): Promise<void> {
+    if (this.initialized) {
+      return;
+    }
+    this.initPromise ??= this.bootstrap().catch((err: unknown) => {
+      this.initPromise = undefined;
+      throw err;
+    });
+    await this.initPromise;
+  }
+
+  public get(sessionId: SessionId): Session {
+    this.requireInitialized();
+    const key = Session.encodeId(sessionId);
+    const cached = this.cache.get(key);
+    if (cached) {
+      cached.lastUsed = Date.now();
+      return cached;
+    }
+    this.evictIfNeeded();
+    const session = new Session({
+      id: sessionId,
+      settings: this.settings.get(key) ?? {},
+      config: this.config,
+      api: this.api,
+      agentDir: this.agentDir,
+      authStorage: this.authStorage,
+      modelRegistry: this.modelRegistry,
+      scheduler: this.scheduler,
+      settingsManagerFor: (cwd) => this.settingsManagerFor(cwd),
+      persistSettings: (patch) => this.persistSettings(key, patch),
+      getBotUsername: () => this.botUsername,
+    });
+    this.cache.set(key, session);
+    return session;
+  }
+
+  public async disposeAll(): Promise<void> {
+    for (const session of this.cache.values()) {
+      session.dispose();
+    }
+    this.cache.clear();
+    if (this.initialized) {
+      await this.flushSettings();
+    }
+  }
+
+  private async bootstrap(): Promise<void> {
+    const loaded = await Fs.readJsonOrEmpty<Record<string, SessionSettings>>(
+      join(this.config.configDir, "state.json"),
+      {}
+    );
+    this.settings = new Map(Object.entries(loaded));
+    await mkdir(join(this.config.configDir, "sessions"), { recursive: true });
+    await mkdir(join(this.config.configDir, "isolated-sessions"), {
+      recursive: true,
+    });
+    await mkdir(join(this.config.configDir, "instructions"), {
+      recursive: true,
+    });
+    this.initialized = true;
+  }
+
+  private requireInitialized(): void {
+    if (!this.initialized) {
+      throw new Error("SessionRegistry.init() must complete before use");
+    }
+  }
+
+  private async persistSettings(
+    key: string,
+    patch: Partial<SessionSettings>
+  ): Promise<void> {
+    const prev = this.settings.get(key) ?? {};
+    this.settings.set(key, { ...prev, ...patch });
+    await this.flushSettings();
+  }
+
+  private async flushSettings(): Promise<void> {
+    try {
+      await Fs.writeAtomic(
+        join(this.config.configDir, "state.json"),
+        JSON.stringify(Object.fromEntries(this.settings), null, 2)
+      );
+    } catch (err) {
+      console.warn(`[registry] state save failed:`, err);
+    }
+  }
+
+  private settingsManagerFor(cwd: string): SettingsManager {
+    const existing = this.settingsManagers.get(cwd);
+    if (existing) {
+      return existing;
+    }
+    const settingsManager = SettingsManager.create(cwd, this.agentDir);
+    this.settingsManagers.set(cwd, settingsManager);
+    return settingsManager;
+  }
+
+  private evictIfNeeded(): void {
+    if (this.cache.size < LRU_CAP) {
+      return;
+    }
+    let oldestKey: string | undefined;
+    let oldestTime = Infinity;
+    for (const [k, v] of this.cache) {
+      if (v.lastUsed < oldestTime) {
+        oldestTime = v.lastUsed;
+        oldestKey = k;
+      }
+    }
+    if (!oldestKey) {
+      return;
+    }
+    const entry = this.cache.get(oldestKey)!;
+    entry.dispose();
+    this.cache.delete(oldestKey);
+  }
+}

package/src/telegram/Supervisor.ts

@@ -0,0 +1,357 @@
+import { mkdir, realpath, rm, stat, unlink } from "node:fs/promises";
+import { homedir } from "node:os";
+import { dirname, join } from "node:path";
+
+import { Fs } from "../shared/Fs";
+
+const UNIT_NAME = "pim-telegram";
+const LAUNCHD_LABEL = "com.aaroncql.pim-telegram";
+const NPM_PACKAGE = "@aaroncql/pim-agent";
+const CONFIRM_FILE = "update-confirm.json";
+
+export type UpdateConfirmEntry = {
+  readonly chatId: number;
+  readonly threadId: number | undefined;
+  readonly messageId: number;
+};
+
+export type Mode = {
+  readonly kind: "dev" | "prod";
+  readonly packageRoot: string;
+  readonly pimEntry: string;
+  readonly bunPath: string;
+};
+
+export type UpdateResult =
+  | { readonly ok: true }
+  | { readonly ok: false; readonly error: string };
+
+export class Supervisor {
+  public static async install(): Promise<void> {
+    const mode = await Supervisor.detectMode();
+    console.log(`[install] ${mode.kind} mode, root=${mode.packageRoot}`);
+    if (process.platform === "linux") {
+      const path = Supervisor.systemdUnitPath();
+      await Fs.writeAtomic(path, Supervisor.systemdUnit(mode));
+      console.log(`[install] wrote ${path}`);
+      await Supervisor.runOrThrow(["systemctl", "--user", "daemon-reload"]);
+      await Supervisor.runOrThrow([
+        "systemctl",
+        "--user",
+        "enable",
+        "--now",
+        UNIT_NAME,
+      ]);
+      console.log(`[install] enabled and started ${UNIT_NAME}.service`);
+      if (!(await Supervisor.lingerEnabled())) {
+        console.log(
+          `[install] hint: run 'loginctl enable-linger' so the service starts at boot without an active login`
+        );
+      }
+      return;
+    }
+    if (process.platform === "darwin") {
+      await mkdir(join(homedir(), "Library", "Logs"), { recursive: true });
+      const path = Supervisor.launchdPlistPath();
+      await Fs.writeAtomic(path, Supervisor.launchdPlist(mode));
+      console.log(`[install] wrote ${path}`);
+      const uid = process.getuid?.() ?? 0;
+      try {
+        await Supervisor.runOrThrow([
+          "launchctl",
+          "bootout",
+          `gui/${uid}/${LAUNCHD_LABEL}`,
+        ]);
+      } catch {
+        // bootout fails when the service isn't currently loaded; safe to ignore before bootstrap
+      }
+      await Supervisor.runOrThrow([
+        "launchctl",
+        "bootstrap",
+        `gui/${uid}`,
+        path,
+      ]);
+      console.log(`[install] bootstrapped ${LAUNCHD_LABEL}`);
+      return;
+    }
+    throw new Error(`Unsupported platform: ${process.platform}`);
+  }
+
+  public static async uninstall(): Promise<void> {
+    if (process.platform === "linux") {
+      const path = Supervisor.systemdUnitPath();
+      try {
+        await Supervisor.runOrThrow([
+          "systemctl",
+          "--user",
+          "disable",
+          "--now",
+          UNIT_NAME,
+        ]);
+      } catch (err) {
+        console.warn(`[uninstall] disable failed:`, (err as Error).message);
+      }
+      try {
+        await rm(path);
+        console.log(`[uninstall] removed ${path}`);
+      } catch (err) {
+        if ((err as NodeJS.ErrnoException).code !== "ENOENT") {
+          throw err;
+        }
+      }
+      try {
+        await Supervisor.runOrThrow(["systemctl", "--user", "daemon-reload"]);
+      } catch (err) {
+        console.warn(
+          `[uninstall] daemon-reload failed:`,
+          (err as Error).message
+        );
+      }
+      return;
+    }
+    if (process.platform === "darwin") {
+      const path = Supervisor.launchdPlistPath();
+      const uid = process.getuid?.() ?? 0;
+      try {
+        await Supervisor.runOrThrow([
+          "launchctl",
+          "bootout",
+          `gui/${uid}/${LAUNCHD_LABEL}`,
+        ]);
+      } catch (err) {
+        console.warn(`[uninstall] bootout failed:`, (err as Error).message);
+      }
+      try {
+        await rm(path);
+        console.log(`[uninstall] removed ${path}`);
+      } catch (err) {
+        if ((err as NodeJS.ErrnoException).code !== "ENOENT") {
+          throw err;
+        }
+      }
+      return;
+    }
+    throw new Error(`Unsupported platform: ${process.platform}`);
+  }
+
+  public static async update(): Promise<UpdateResult> {
+    const mode = await Supervisor.detectMode();
+    const cmd =
+      mode.kind === "dev"
+        ? ["bun", "install"]
+        : ["bun", "install", "-g", `${NPM_PACKAGE}@latest`];
+    const proc = Bun.spawn([...cmd], {
+      cwd: mode.kind === "dev" ? mode.packageRoot : undefined,
+      stdout: "inherit",
+      stderr: "pipe",
+    });
+    const stderr = await new Response(proc.stderr).text();
+    const code = await proc.exited;
+    if (code !== 0) {
+      return {
+        ok: false,
+        error: `${cmd.join(" ")} exit ${code}: ${stderr.trim() || "(no stderr)"}`,
+      };
+    }
+    return { ok: true };
+  }
+
+  public static restart(): never {
+    process.exit(0);
+  }
+
+  public static async readVersion(): Promise<string> {
+    const mode = await Supervisor.detectMode();
+    const pkg = await Bun.file(join(mode.packageRoot, "package.json")).json();
+    return typeof pkg?.version === "string" ? pkg.version : "?";
+  }
+
+  public static async detectMode(): Promise<Mode> {
+    if (Supervisor.cachedMode) {
+      return Supervisor.cachedMode;
+    }
+    const here = await realpath(Bun.fileURLToPath(import.meta.url));
+    const packageRoot = await Supervisor.findPackageRoot(dirname(here));
+    const hasGit = await Supervisor.pathExists(join(packageRoot, ".git"));
+    Supervisor.cachedMode = {
+      kind: hasGit ? "dev" : "prod",
+      packageRoot,
+      pimEntry: join(packageRoot, "bin", "pim.ts"),
+      bunPath: process.execPath,
+    };
+    return Supervisor.cachedMode;
+  }
+
+  public static async appendUpdateConfirm(
+    configDir: string,
+    entry: UpdateConfirmEntry
+  ): Promise<void> {
+    const merged = [...(await Supervisor.readUpdateConfirm(configDir)), entry];
+    await Fs.writeAtomic(
+      Supervisor.updateConfirmPath(configDir),
+      JSON.stringify(merged, null, 2)
+    );
+  }
+
+  public static async readUpdateConfirm(
+    configDir: string
+  ): Promise<ReadonlyArray<UpdateConfirmEntry>> {
+    const data = await Fs.readJsonOrEmpty<unknown[]>(
+      Supervisor.updateConfirmPath(configDir),
+      []
+    );
+    if (!Array.isArray(data)) {
+      return [];
+    }
+    return data.filter(
+      (e): e is UpdateConfirmEntry =>
+        !!e &&
+        typeof e === "object" &&
+        typeof (e as UpdateConfirmEntry).chatId === "number" &&
+        ((e as UpdateConfirmEntry).threadId === undefined ||
+          typeof (e as UpdateConfirmEntry).threadId === "number") &&
+        typeof (e as UpdateConfirmEntry).messageId === "number"
+    );
+  }
+
+  public static async clearUpdateConfirm(configDir: string): Promise<void> {
+    try {
+      await unlink(Supervisor.updateConfirmPath(configDir));
+    } catch (err) {
+      if ((err as NodeJS.ErrnoException).code !== "ENOENT") {
+        console.warn(`[update-confirm] unlink failed:`, err);
+      }
+    }
+  }
+
+  private static cachedMode: Mode | undefined;
+
+  private static updateConfirmPath(configDir: string): string {
+    return join(configDir, CONFIRM_FILE);
+  }
+
+  private static async pathExists(p: string): Promise<boolean> {
+    try {
+      await stat(p);
+      return true;
+    } catch {
+      return false;
+    }
+  }
+
+  private static async findPackageRoot(start: string): Promise<string> {
+    let dir = start;
+    for (let i = 0; i < 32; i++) {
+      if (await Supervisor.pathExists(join(dir, "package.json"))) {
+        return dir;
+      }
+      const parent = dirname(dir);
+      if (parent === dir) {
+        break;
+      }
+      dir = parent;
+    }
+    throw new Error(`Could not locate package root from ${start}`);
+  }
+
+  private static systemdUnitPath(): string {
+    return join(
+      homedir(),
+      ".config",
+      "systemd",
+      "user",
+      `${UNIT_NAME}.service`
+    );
+  }
+
+  private static launchdPlistPath(): string {
+    return join(homedir(), "Library", "LaunchAgents", `${LAUNCHD_LABEL}.plist`);
+  }
+
+  private static launchdLogPath(): string {
+    return join(homedir(), "Library", "Logs", `${UNIT_NAME}.log`);
+  }
+
+  private static unitPath(mode: Mode): string {
+    return `${dirname(mode.bunPath)}:/usr/local/bin:/usr/bin:/bin`;
+  }
+
+  private static systemdUnit(mode: Mode): string {
+    return [
+      "[Unit]",
+      "Description=Pim Telegram daemon",
+      "After=network-online.target",
+      "Wants=network-online.target",
+      "",
+      "[Service]",
+      "Type=simple",
+      `Environment=PATH=${Supervisor.unitPath(mode)}`,
+      `ExecStart=${mode.bunPath} ${mode.pimEntry} --mode telegram`,
+      "Restart=always",
+      "RestartSec=2",
+      "",
+      "[Install]",
+      "WantedBy=default.target",
+      "",
+    ].join("\n");
+  }
+
+  private static launchdPlist(mode: Mode): string {
+    const logPath = Supervisor.launchdLogPath();
+    return [
+      `<?xml version="1.0" encoding="UTF-8"?>`,
+      `<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">`,
+      `<plist version="1.0">`,
+      `<dict>`,
+      `  <key>Label</key>`,
+      `  <string>${LAUNCHD_LABEL}</string>`,
+      `  <key>ProgramArguments</key>`,
+      `  <array>`,
+      `    <string>${mode.bunPath}</string>`,
+      `    <string>${mode.pimEntry}</string>`,
+      `    <string>--mode</string>`,
+      `    <string>telegram</string>`,
+      `  </array>`,
+      `  <key>EnvironmentVariables</key>`,
+      `  <dict>`,
+      `    <key>PATH</key>`,
+      `    <string>${Supervisor.unitPath(mode)}</string>`,
+      `  </dict>`,
+      `  <key>RunAtLoad</key>`,
+      `  <true/>`,
+      `  <key>KeepAlive</key>`,
+      `  <true/>`,
+      `  <key>StandardOutPath</key>`,
+      `  <string>${logPath}</string>`,
+      `  <key>StandardErrorPath</key>`,
+      `  <string>${logPath}</string>`,
+      `</dict>`,
+      `</plist>`,
+      ``,
+    ].join("\n");
+  }
+
+  private static async lingerEnabled(): Promise<boolean> {
+    const proc = Bun.spawn(["loginctl", "show-user", "--property=Linger"], {
+      stdout: "pipe",
+      stderr: "pipe",
+    });
+    const code = await proc.exited;
+    if (code !== 0) {
+      return false;
+    }
+    const out = await new Response(proc.stdout).text();
+    return out.includes("Linger=yes");
+  }
+
+  private static async runOrThrow(cmd: ReadonlyArray<string>): Promise<void> {
+    const proc = Bun.spawn([...cmd], { stdout: "inherit", stderr: "pipe" });
+    const stderr = await new Response(proc.stderr).text();
+    const code = await proc.exited;
+    if (code !== 0) {
+      throw new Error(
+        `${cmd.join(" ")} exit ${code}: ${stderr.trim() || "(no stderr)"}`
+      );
+    }
+  }
+}