@posthog/agent 2.0.0 → 2.0.2

package/src/sagas/resume-saga.ts

@@ -0,0 +1,332 @@
+import type { ContentBlock } from "@agentclientprotocol/sdk";
+import { Saga } from "@posthog/shared";
+import { POSTHOG_NOTIFICATIONS } from "../acp-extensions.js";
+import type { PostHogAPIClient } from "../posthog-api.js";
+import { TreeTracker } from "../tree-tracker.js";
+import type {
+  DeviceInfo,
+  StoredNotification,
+  TreeSnapshotEvent,
+} from "../types.js";
+import { Logger } from "../utils/logger.js";
+
+export interface ConversationTurn {
+  role: "user" | "assistant";
+  content: ContentBlock[];
+  toolCalls?: ToolCallInfo[];
+}
+
+export interface ToolCallInfo {
+  toolCallId: string;
+  toolName: string;
+  input: unknown;
+  result?: unknown;
+}
+
+export interface ResumeInput {
+  taskId: string;
+  runId: string;
+  repositoryPath: string;
+  apiClient: PostHogAPIClient;
+  logger?: Logger;
+}
+
+export interface ResumeOutput {
+  conversation: ConversationTurn[];
+  latestSnapshot: TreeSnapshotEvent | null;
+  snapshotApplied: boolean;
+  interrupted: boolean;
+  lastDevice?: DeviceInfo;
+  logEntryCount: number;
+}
+
+export class ResumeSaga extends Saga<ResumeInput, ResumeOutput> {
+  protected async execute(input: ResumeInput): Promise<ResumeOutput> {
+    const { taskId, runId, repositoryPath, apiClient } = input;
+    const logger =
+      input.logger || new Logger({ debug: false, prefix: "[Resume]" });
+
+    // Step 1: Fetch task run (read-only)
+    const taskRun = await this.readOnlyStep("fetch_task_run", () =>
+      apiClient.getTaskRun(taskId, runId),
+    );
+
+    if (!taskRun.log_url) {
+      this.log.info("No log URL found, starting fresh");
+      return this.emptyResult();
+    }
+
+    // Step 2: Fetch log entries (read-only)
+    const entries = await this.readOnlyStep("fetch_logs", () =>
+      apiClient.fetchTaskRunLogs(taskRun),
+    );
+
+    if (entries.length === 0) {
+      this.log.info("No log entries found, starting fresh");
+      return this.emptyResult();
+    }
+
+    this.log.info("Fetched log entries", { count: entries.length });
+
+    // Step 3: Find latest snapshot (read-only, pure computation)
+    const latestSnapshot = await this.readOnlyStep("find_snapshot", () =>
+      Promise.resolve(this.findLatestTreeSnapshot(entries)),
+    );
+
+    // Step 4: Apply snapshot if present (wrapped in step for consistent logging)
+    // Note: We use a try/catch inside the step because snapshot failure should NOT fail the saga
+    let snapshotApplied = false;
+    if (latestSnapshot?.archiveUrl) {
+      this.log.info("Found tree snapshot", {
+        treeHash: latestSnapshot.treeHash,
+        hasArchiveUrl: true,
+        changes: latestSnapshot.changes?.length ?? 0,
+        interrupted: latestSnapshot.interrupted,
+      });
+
+      await this.step({
+        name: "apply_snapshot",
+        execute: async () => {
+          const treeTracker = new TreeTracker({
+            repositoryPath,
+            taskId,
+            runId,
+            apiClient,
+            logger: logger.child("TreeTracker"),
+          });
+
+          try {
+            await treeTracker.applyTreeSnapshot(latestSnapshot);
+            treeTracker.setLastTreeHash(latestSnapshot.treeHash);
+            snapshotApplied = true;
+            this.log.info("Tree snapshot applied successfully", {
+              treeHash: latestSnapshot.treeHash,
+            });
+          } catch (error) {
+            // Log but don't fail - continue with conversation rebuild
+            // ApplySnapshotSaga handles its own rollback internally
+            this.log.warn(
+              "Failed to apply tree snapshot, continuing without it",
+              {
+                error: error instanceof Error ? error.message : String(error),
+                treeHash: latestSnapshot.treeHash,
+              },
+            );
+          }
+        },
+        rollback: async () => {
+          // Inner ApplySnapshotSaga handles its own rollback
+        },
+      });
+    } else if (latestSnapshot) {
+      this.log.warn(
+        "Snapshot found but has no archive URL - files cannot be restored",
+        {
+          treeHash: latestSnapshot.treeHash,
+          changes: latestSnapshot.changes?.length ?? 0,
+        },
+      );
+    }
+
+    // Step 5: Rebuild conversation (read-only, pure computation)
+    const conversation = await this.readOnlyStep("rebuild_conversation", () =>
+      Promise.resolve(this.rebuildConversation(entries)),
+    );
+
+    // Step 6: Find device info (read-only, pure computation)
+    const lastDevice = await this.readOnlyStep("find_device", () =>
+      Promise.resolve(this.findLastDeviceInfo(entries)),
+    );
+
+    this.log.info("Resume state rebuilt", {
+      turns: conversation.length,
+      hasSnapshot: !!latestSnapshot,
+      snapshotApplied,
+      interrupted: latestSnapshot?.interrupted ?? false,
+    });
+
+    return {
+      conversation,
+      latestSnapshot,
+      snapshotApplied,
+      interrupted: latestSnapshot?.interrupted ?? false,
+      lastDevice,
+      logEntryCount: entries.length,
+    };
+  }
+
+  private emptyResult(): ResumeOutput {
+    return {
+      conversation: [],
+      latestSnapshot: null,
+      snapshotApplied: false,
+      interrupted: false,
+      logEntryCount: 0,
+    };
+  }
+
+  private findLatestTreeSnapshot(
+    entries: StoredNotification[],
+  ): TreeSnapshotEvent | null {
+    const sdkPrefixedMethod = `_${POSTHOG_NOTIFICATIONS.TREE_SNAPSHOT}`;
+
+    for (let i = entries.length - 1; i >= 0; i--) {
+      const entry = entries[i];
+      const method = entry.notification?.method;
+      if (
+        method === sdkPrefixedMethod ||
+        method === POSTHOG_NOTIFICATIONS.TREE_SNAPSHOT
+      ) {
+        const params = entry.notification.params as
+          | TreeSnapshotEvent
+          | undefined;
+        if (params?.treeHash) {
+          return params;
+        }
+      }
+    }
+    return null;
+  }
+
+  private findLastDeviceInfo(
+    entries: StoredNotification[],
+  ): DeviceInfo | undefined {
+    for (let i = entries.length - 1; i >= 0; i--) {
+      const entry = entries[i];
+      const params = entry.notification?.params as
+        | { device?: DeviceInfo }
+        | undefined;
+      if (params?.device) {
+        return params.device;
+      }
+    }
+    return undefined;
+  }
+
+  private rebuildConversation(
+    entries: StoredNotification[],
+  ): ConversationTurn[] {
+    const turns: ConversationTurn[] = [];
+    let currentAssistantContent: ContentBlock[] = [];
+    let currentToolCalls: ToolCallInfo[] = [];
+
+    for (const entry of entries) {
+      const method = entry.notification?.method;
+      const params = entry.notification?.params as Record<string, unknown>;
+
+      if (method === "session/update" && params?.update) {
+        const update = params.update as Record<string, unknown>;
+        const sessionUpdate = update.sessionUpdate as string;
+
+        switch (sessionUpdate) {
+          case "user_message":
+          case "user_message_chunk": {
+            if (
+              currentAssistantContent.length > 0 ||
+              currentToolCalls.length > 0
+            ) {
+              turns.push({
+                role: "assistant",
+                content: currentAssistantContent,
+                toolCalls:
+                  currentToolCalls.length > 0 ? currentToolCalls : undefined,
+              });
+              currentAssistantContent = [];
+              currentToolCalls = [];
+            }
+
+            const content = update.content as ContentBlock | ContentBlock[];
+            const contentArray = Array.isArray(content) ? content : [content];
+            turns.push({
+              role: "user",
+              content: contentArray,
+            });
+            break;
+          }
+
+          case "agent_message_chunk": {
+            const content = update.content as ContentBlock | undefined;
+            if (content) {
+              if (
+                content.type === "text" &&
+                currentAssistantContent.length > 0 &&
+                currentAssistantContent[currentAssistantContent.length - 1]
+                  .type === "text"
+              ) {
+                const lastBlock = currentAssistantContent[
+                  currentAssistantContent.length - 1
+                ] as { type: "text"; text: string };
+                lastBlock.text += (
+                  content as { type: "text"; text: string }
+                ).text;
+              } else {
+                currentAssistantContent.push(content);
+              }
+            }
+            break;
+          }
+
+          case "tool_call":
+          case "tool_call_update": {
+            const meta = (update._meta as Record<string, unknown>)
+              ?.claudeCode as Record<string, unknown> | undefined;
+            if (meta) {
+              const toolCallId = meta.toolCallId as string | undefined;
+              const toolName = meta.toolName as string | undefined;
+              const toolInput = meta.toolInput;
+              const toolResponse = meta.toolResponse;
+
+              if (toolCallId && toolName) {
+                let toolCall = currentToolCalls.find(
+                  (tc) => tc.toolCallId === toolCallId,
+                );
+                if (!toolCall) {
+                  toolCall = {
+                    toolCallId,
+                    toolName,
+                    input: toolInput,
+                  };
+                  currentToolCalls.push(toolCall);
+                }
+
+                if (toolResponse !== undefined) {
+                  toolCall.result = toolResponse;
+                }
+              }
+            }
+            break;
+          }
+
+          case "tool_result": {
+            const meta = (update._meta as Record<string, unknown>)
+              ?.claudeCode as Record<string, unknown> | undefined;
+            if (meta) {
+              const toolCallId = meta.toolCallId as string | undefined;
+              const toolResponse = meta.toolResponse;
+
+              if (toolCallId) {
+                const toolCall = currentToolCalls.find(
+                  (tc) => tc.toolCallId === toolCallId,
+                );
+                if (toolCall && toolResponse !== undefined) {
+                  toolCall.result = toolResponse;
+                }
+              }
+            }
+            break;
+          }
+        }
+      }
+    }
+
+    if (currentAssistantContent.length > 0 || currentToolCalls.length > 0) {
+      turns.push({
+        role: "assistant",
+        content: currentAssistantContent,
+        toolCalls: currentToolCalls.length > 0 ? currentToolCalls : undefined,
+      });
+    }
+
+    return turns;
+  }
+}

package/src/sagas/test-fixtures.ts

@@ -0,0 +1,250 @@
+import { execFile } from "node:child_process";
+import { existsSync } from "node:fs";
+import { mkdir, readFile, rm, writeFile } from "node:fs/promises";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import { promisify } from "node:util";
+import type { SagaLogger } from "@posthog/shared";
+import * as tar from "tar";
+import { vi } from "vitest";
+import { POSTHOG_NOTIFICATIONS } from "../acp-extensions.js";
+import type { PostHogAPIClient } from "../posthog-api.js";
+import type { StoredNotification, TaskRun, TreeSnapshot } from "../types.js";
+
+const execFileAsync = promisify(execFile);
+
+export interface TestRepo {
+  path: string;
+  cleanup: () => Promise<void>;
+  git: (args: string[]) => Promise<string>;
+  writeFile: (relativePath: string, content: string) => Promise<void>;
+  readFile: (relativePath: string) => Promise<string>;
+  deleteFile: (relativePath: string) => Promise<void>;
+  exists: (relativePath: string) => boolean;
+}
+
+export async function createTestRepo(prefix = "test-repo"): Promise<TestRepo> {
+  const repoPath = join(
+    tmpdir(),
+    `${prefix}-${Date.now()}-${Math.random().toString(36).slice(2)}`,
+  );
+  await mkdir(repoPath, { recursive: true });
+
+  const git = async (args: string[]): Promise<string> => {
+    const { stdout } = await execFileAsync("git", args, { cwd: repoPath });
+    return stdout.trim();
+  };
+
+  await git(["init"]);
+  await git(["config", "user.email", "test@test.com"]);
+  await git(["config", "user.name", "Test"]);
+
+  await writeFile(join(repoPath, ".gitignore"), ".posthog/\n");
+  await writeFile(join(repoPath, "README.md"), "# Test Repo");
+  await git(["add", "."]);
+  await git(["commit", "-m", "Initial commit"]);
+
+  return {
+    path: repoPath,
+    cleanup: () => rm(repoPath, { recursive: true, force: true }),
+    git,
+    writeFile: async (relativePath: string, content: string) => {
+      const fullPath = join(repoPath, relativePath);
+      const dir = join(fullPath, "..");
+      await mkdir(dir, { recursive: true });
+      await writeFile(fullPath, content);
+    },
+    readFile: async (relativePath: string) => {
+      return readFile(join(repoPath, relativePath), "utf-8");
+    },
+    deleteFile: async (relativePath: string) => {
+      await rm(join(repoPath, relativePath), { force: true });
+    },
+    exists: (relativePath: string) => {
+      return existsSync(join(repoPath, relativePath));
+    },
+  };
+}
+
+export function createMockLogger(): SagaLogger {
+  return {
+    info: vi.fn(),
+    debug: vi.fn(),
+    error: vi.fn(),
+    warn: vi.fn(),
+  };
+}
+
+export function createMockApiClient(
+  overrides: Partial<PostHogAPIClient> = {},
+): PostHogAPIClient {
+  return {
+    uploadTaskArtifacts: vi
+      .fn()
+      .mockResolvedValue([{ storage_path: "gs://bucket/trees/test.tar.gz" }]),
+    downloadArtifact: vi.fn(),
+    getTaskRun: vi.fn(),
+    fetchTaskRunLogs: vi.fn(),
+    ...overrides,
+  } as unknown as PostHogAPIClient;
+}
+
+export interface ArchiveFile {
+  path: string;
+  content: string;
+}
+
+export interface ArchiveSymlink {
+  path: string;
+  target: string;
+}
+
+export async function createArchiveBuffer(
+  files: Array<ArchiveFile>,
+  symlinks: Array<ArchiveSymlink> = [],
+): Promise<Buffer> {
+  const { symlink } = await import("node:fs/promises");
+  const tmpDir = join(
+    tmpdir(),
+    `archive-${Date.now()}-${Math.random().toString(36).slice(2)}`,
+  );
+  await mkdir(tmpDir, { recursive: true });
+
+  const filesToArchive =
+    files.length > 0 ? files : [{ path: ".empty", content: "" }];
+
+  for (const file of filesToArchive) {
+    const fullPath = join(tmpDir, file.path);
+    await mkdir(join(fullPath, ".."), { recursive: true });
+    await writeFile(fullPath, file.content);
+  }
+
+  const symlinkPaths: string[] = [];
+  for (const link of symlinks) {
+    const fullPath = join(tmpDir, link.path);
+    await mkdir(join(fullPath, ".."), { recursive: true });
+    await symlink(link.target, fullPath);
+    symlinkPaths.push(link.path);
+  }
+
+  const archivePath = join(tmpDir, "archive.tar.gz");
+  await tar.create({ gzip: true, file: archivePath, cwd: tmpDir }, [
+    ...filesToArchive.map((f) => f.path),
+    ...symlinkPaths,
+  ]);
+
+  const content = await readFile(archivePath);
+  await rm(tmpDir, { recursive: true, force: true });
+
+  return Buffer.from(content.toString("base64"));
+}
+
+export function createSnapshot(
+  overrides: Partial<TreeSnapshot> = {},
+): TreeSnapshot {
+  return {
+    treeHash: "test-tree-hash",
+    baseCommit: null,
+    archiveUrl: "gs://bucket/trees/test.tar.gz",
+    changes: [],
+    timestamp: new Date().toISOString(),
+    ...overrides,
+  };
+}
+
+export function createTaskRun(overrides: Partial<TaskRun> = {}): TaskRun {
+  return {
+    id: "run-1",
+    task: "task-1",
+    team: 1,
+    branch: null,
+    stage: null,
+    environment: "local",
+    status: "in_progress",
+    log_url: "https://logs.example.com/run-1",
+    error_message: null,
+    output: null,
+    state: {},
+    created_at: new Date().toISOString(),
+    updated_at: new Date().toISOString(),
+    completed_at: null,
+    ...overrides,
+  };
+}
+
+export function createNotification(
+  method: string,
+  params: Record<string, unknown>,
+): StoredNotification {
+  return {
+    type: "notification",
+    timestamp: new Date().toISOString(),
+    notification: {
+      jsonrpc: "2.0",
+      method,
+      params,
+    },
+  };
+}
+
+export function createUserMessage(content: string): StoredNotification {
+  return createNotification("session/update", {
+    update: {
+      sessionUpdate: "user_message",
+      content: { type: "text", text: content },
+    },
+  });
+}
+
+export function createAgentChunk(text: string): StoredNotification {
+  return createNotification("session/update", {
+    update: {
+      sessionUpdate: "agent_message_chunk",
+      content: { type: "text", text },
+    },
+  });
+}
+
+export function createToolCall(
+  toolCallId: string,
+  toolName: string,
+  toolInput: unknown,
+): StoredNotification {
+  return createNotification("session/update", {
+    update: {
+      sessionUpdate: "tool_call",
+      _meta: {
+        claudeCode: { toolCallId, toolName, toolInput },
+      },
+    },
+  });
+}
+
+export function createToolResult(
+  toolCallId: string,
+  toolResponse: unknown,
+): StoredNotification {
+  return createNotification("session/update", {
+    update: {
+      sessionUpdate: "tool_result",
+      _meta: {
+        claudeCode: { toolCallId, toolResponse },
+      },
+    },
+  });
+}
+
+export function createTreeSnapshotNotification(
+  treeHash: string,
+  archiveUrl?: string,
+  options: { interrupted?: boolean; device?: { type: "local" | "cloud" } } = {},
+): StoredNotification {
+  return createNotification(POSTHOG_NOTIFICATIONS.TREE_SNAPSHOT, {
+    treeHash,
+    baseCommit: "abc123",
+    archiveUrl,
+    changes: [{ path: "file.ts", status: "A" }],
+    timestamp: new Date().toISOString(),
+    ...options,
+  });
+}