@posthog/agent 2.0.0 → 2.0.2

package/src/test/mocks/claude-sdk.ts

@@ -0,0 +1,251 @@
+import type {
+  Query,
+  SDKMessage,
+  SDKResultError,
+  SDKResultSuccess,
+} from "@anthropic-ai/claude-agent-sdk";
+import type { Mock } from "vitest";
+import { vi } from "vitest";
+
+export interface MockQueryHelpers {
+  sendMessage: (message: SDKMessage) => void;
+  complete: (result?: SDKResultSuccess) => void;
+  sendError: (result: SDKResultError) => void;
+  simulateError: (error: Error) => void;
+  queueError: (error: Error) => void;
+  simulateTimeout: () => void;
+  isAborted: () => boolean;
+}
+
+export interface MockQuery extends Omit<Query, "interrupt"> {
+  _mockHelpers: MockQueryHelpers;
+  _abortController: AbortController;
+  interrupt: Mock<() => Promise<void>>;
+}
+
+export interface CreateMockQueryOptions {
+  abortController?: AbortController;
+}
+
+export function createMockQuery(
+  options: CreateMockQueryOptions = {},
+): MockQuery {
+  const abortController = options.abortController ?? new AbortController();
+  let resolveNext: ((value: IteratorResult<SDKMessage, void>) => void) | null =
+    null;
+  let rejectNext: ((error: Error) => void) | null = null;
+  let isDone = false;
+  let queuedError: Error | null = null;
+  let isTimedOut = false;
+
+  const createNextPromise = (): Promise<IteratorResult<SDKMessage, void>> => {
+    if (isDone) {
+      return Promise.resolve({ value: undefined, done: true as const });
+    }
+    if (queuedError) {
+      const error = queuedError;
+      queuedError = null;
+      return Promise.reject(error);
+    }
+    if (isTimedOut) {
+      return new Promise(() => {});
+    }
+    return new Promise((resolve, reject) => {
+      resolveNext = resolve;
+      rejectNext = reject;
+
+      abortController.signal.addEventListener("abort", () => {
+        isDone = true;
+        resolve({ value: undefined, done: true as const });
+      });
+    });
+  };
+
+  const mockQuery: MockQuery = {
+    next: vi.fn(() => createNextPromise()),
+    return: vi.fn(() => {
+      isDone = true;
+      return Promise.resolve({ value: undefined, done: true as const });
+    }),
+    throw: vi.fn((error: Error) => {
+      isDone = true;
+      return Promise.reject(error);
+    }),
+    [Symbol.asyncIterator]() {
+      return this;
+    },
+    interrupt: vi.fn(async () => {
+      abortController.abort();
+      isDone = true;
+      if (resolveNext) {
+        resolveNext({ value: undefined, done: true as const });
+        resolveNext = null;
+        rejectNext = null;
+      }
+    }),
+    setPermissionMode: vi.fn().mockResolvedValue(undefined),
+    setModel: vi.fn().mockResolvedValue(undefined),
+    setMaxThinkingTokens: vi.fn().mockResolvedValue(undefined),
+    supportedCommands: vi.fn().mockResolvedValue([]),
+    supportedModels: vi.fn().mockResolvedValue([]),
+    mcpServerStatus: vi.fn().mockResolvedValue([]),
+    accountInfo: vi.fn().mockResolvedValue({}),
+    rewindFiles: vi.fn().mockResolvedValue({ canRewind: false }),
+    setMcpServers: vi
+      .fn()
+      .mockResolvedValue({ added: [], removed: [], errors: {} }),
+    streamInput: vi.fn().mockResolvedValue(undefined),
+    [Symbol.asyncDispose]: vi.fn().mockResolvedValue(undefined),
+    _abortController: abortController,
+    _mockHelpers: {
+      sendMessage(message: SDKMessage) {
+        if (resolveNext && !isDone) {
+          resolveNext({ value: message, done: false });
+          resolveNext = null;
+          rejectNext = null;
+        }
+      },
+      complete(result?: SDKResultSuccess) {
+        isDone = true;
+        if (resolveNext) {
+          if (result) {
+            resolveNext({ value: result, done: false });
+          } else {
+            resolveNext({ value: undefined, done: true });
+          }
+          resolveNext = null;
+          rejectNext = null;
+        }
+      },
+      sendError(result: SDKResultError) {
+        isDone = true;
+        if (resolveNext) {
+          resolveNext({ value: result, done: false });
+          resolveNext = null;
+          rejectNext = null;
+        }
+      },
+      simulateError(error: Error) {
+        if (rejectNext) {
+          rejectNext(error);
+          resolveNext = null;
+          rejectNext = null;
+        }
+      },
+      queueError(error: Error) {
+        queuedError = error;
+      },
+      simulateTimeout() {
+        isTimedOut = true;
+      },
+      isAborted() {
+        return abortController.signal.aborted;
+      },
+    },
+  };
+
+  return mockQuery;
+}
+
+export function createSuccessResult(
+  overrides: Partial<SDKResultSuccess> = {},
+): SDKResultSuccess {
+  return {
+    type: "result",
+    subtype: "success",
+    duration_ms: 100,
+    duration_api_ms: 50,
+    is_error: false,
+    num_turns: 1,
+    result: "Done",
+    total_cost_usd: 0.01,
+    usage: {
+      input_tokens: 100,
+      output_tokens: 50,
+      cache_read_input_tokens: 0,
+      cache_creation_input_tokens: 0,
+      cache_creation: {
+        ephemeral_1h_input_tokens: 0,
+        ephemeral_5m_input_tokens: 0,
+      },
+      server_tool_use: { web_search_requests: 0, web_fetch_requests: 0 },
+      service_tier: "standard",
+    },
+    modelUsage: {},
+    permission_denials: [],
+    uuid: crypto.randomUUID() as `${string}-${string}-${string}-${string}-${string}`,
+    session_id: "test-session",
+    ...overrides,
+  };
+}
+
+export function createErrorResult(
+  overrides: Partial<SDKResultError> = {},
+): SDKResultError {
+  return {
+    type: "result",
+    subtype: "error_during_execution",
+    is_error: true,
+    errors: ["Test error"],
+    duration_ms: 100,
+    duration_api_ms: 50,
+    num_turns: 1,
+    total_cost_usd: 0.01,
+    usage: {
+      input_tokens: 100,
+      output_tokens: 50,
+      cache_read_input_tokens: 0,
+      cache_creation_input_tokens: 0,
+      cache_creation: {
+        ephemeral_1h_input_tokens: 0,
+        ephemeral_5m_input_tokens: 0,
+      },
+      server_tool_use: { web_search_requests: 0, web_fetch_requests: 0 },
+      service_tier: "standard",
+    },
+    modelUsage: {},
+    permission_denials: [],
+    uuid: crypto.randomUUID() as `${string}-${string}-${string}-${string}-${string}`,
+    session_id: "test-session",
+    ...overrides,
+  };
+}
+
+export function createInitMessage(sessionId = "test-session"): SDKMessage {
+  return {
+    type: "system",
+    subtype: "init",
+    agents: [],
+    apiKeySource: "user",
+    betas: [],
+    claude_code_version: "1.0.0",
+    cwd: "/tmp",
+    tools: [],
+    mcp_servers: [],
+    model: "claude-sonnet-4-5-20250929",
+    permissionMode: "default",
+    slash_commands: [],
+    output_style: "default",
+    skills: [],
+    plugins: [],
+    uuid: crypto.randomUUID() as `${string}-${string}-${string}-${string}-${string}`,
+    session_id: sessionId,
+  };
+}
+
+export interface MockQueryRef {
+  current: MockQuery | null;
+}
+
+export function createClaudeSdkMock(mockQueryRef: MockQueryRef) {
+  return {
+    query: vi.fn(() => {
+      const mq = createMockQuery();
+      mockQueryRef.current = mq;
+      setTimeout(() => {
+        mq._mockHelpers.sendMessage(createInitMessage());
+      }, 10);
+      return mq;
+    }),
+  };
+}

package/src/test/mocks/msw-handlers.ts

@@ -0,0 +1,48 @@
+import { HttpResponse, http } from "msw";
+import { SseController } from "../controllers/sse-controller.js";
+
+export { SseController };
+
+type AnyHttpResponse = Response | ReturnType<typeof HttpResponse.json>;
+
+export interface PostHogHandlersOptions {
+  baseUrl?: string;
+  onAppendLog?: (entries: unknown[]) => void;
+  getTaskRun?: () => unknown;
+  appendLogResponse?: () => AnyHttpResponse;
+}
+
+export function createPostHogHandlers(options: PostHogHandlersOptions = {}) {
+  const {
+    baseUrl = "http://localhost:8000",
+    onAppendLog,
+    getTaskRun,
+    appendLogResponse,
+  } = options;
+
+  return [
+    // POST /append_log/ - Agent log entries
+    http.post(
+      `${baseUrl}/api/projects/:projectId/tasks/:taskId/runs/:runId/append_log/`,
+      async ({ request }) => {
+        if (appendLogResponse) {
+          return appendLogResponse();
+        }
+        const body = (await request.json()) as { entries: unknown[] };
+        if (body.entries?.length > 0) {
+          onAppendLog?.(body.entries);
+        }
+        return HttpResponse.json({});
+      },
+    ),
+
+    // GET /runs/:runId - Fetch task run details
+    http.get(
+      `${baseUrl}/api/projects/:projectId/tasks/:taskId/runs/:runId`,
+      () => {
+        const taskRun = getTaskRun?.() ?? { log_url: "" };
+        return HttpResponse.json(taskRun);
+      },
+    ),
+  ];
+}

package/src/test/setup.ts

@@ -0,0 +1,114 @@
+import { type SetupServerApi, setupServer } from "msw/node";
+import { AgentServer } from "../server/agent-server.js";
+import { SseController } from "./controllers/sse-controller.js";
+import { createTestRepo, type TestRepo } from "./fixtures/api.js";
+import {
+  type AgentServerConfig,
+  createAgentServerConfig,
+} from "./fixtures/config.js";
+import { createPostHogHandlers } from "./mocks/msw-handlers.js";
+
+export interface TestContext {
+  repo: TestRepo;
+  sseController: SseController;
+  appendLogCalls: unknown[][];
+  server: SetupServerApi;
+  agentServer: AgentServer;
+  config: AgentServerConfig;
+  createAgentServer: (overrides?: Partial<AgentServerConfig>) => AgentServer;
+  resetSseController: () => SseController;
+  cleanup: () => Promise<void>;
+}
+
+export interface CreateTestContextOptions {
+  configOverrides?: Partial<AgentServerConfig>;
+  autoStart?: boolean;
+}
+
+export async function createTestContext(
+  options: CreateTestContextOptions = {},
+): Promise<TestContext> {
+  const repo = await createTestRepo("agent-server");
+  let sseController = new SseController();
+  const appendLogCalls: unknown[][] = [];
+
+  const server = setupServer(
+    ...createPostHogHandlers({
+      baseUrl: "http://localhost:8000",
+      onAppendLog: (entries) => appendLogCalls.push(entries),
+    }),
+  );
+
+  server.listen({ onUnhandledRequest: "bypass" });
+
+  const config = createAgentServerConfig(repo, options.configOverrides);
+
+  const agentServer = new AgentServer(config);
+
+  const createAgentServer = (overrides: Partial<AgentServerConfig> = {}) => {
+    return new AgentServer({
+      ...config,
+      ...overrides,
+    });
+  };
+
+  const resetSseController = () => {
+    sseController.close();
+    sseController = new SseController();
+    return sseController;
+  };
+
+  const cleanup = async () => {
+    sseController.close();
+    server.close();
+    await repo.cleanup();
+  };
+
+  return {
+    repo,
+    sseController,
+    appendLogCalls,
+    server,
+    agentServer,
+    config,
+    createAgentServer,
+    resetSseController,
+    cleanup,
+  };
+}
+
+export {
+  expectNoNotification,
+  expectNotification,
+  findNotification,
+  hasNotification,
+} from "./assertions.js";
+export {
+  createMockApiClient,
+  createTaskRun,
+  createTestRepo,
+  type TestRepo,
+} from "./fixtures/api.js";
+export { createAgentServerConfig } from "./fixtures/config.js";
+export {
+  createAgentChunk,
+  createNotification,
+  createStatusNotification,
+  createToolCall,
+  createToolResult,
+  createTreeSnapshotNotification,
+  createUserMessage,
+} from "./fixtures/notifications.js";
+export {
+  createErrorResult,
+  createInitMessage,
+  createMockQuery,
+  createSuccessResult,
+  type MockQuery,
+} from "./mocks/claude-sdk.js";
+export { createPostHogHandlers, SseController } from "./mocks/msw-handlers.js";
+export {
+  waitForArrayLength,
+  waitForCallCount,
+  waitForCondition,
+} from "./wait.js";

package/src/test/wait.ts

@@ -0,0 +1,41 @@
+export interface WaitOptions {
+  timeout?: number;
+  interval?: number;
+}
+
+export async function waitForCondition(
+  condition: () => boolean | Promise<boolean>,
+  options: WaitOptions = {},
+): Promise<void> {
+  const { timeout = 1000, interval = 20 } = options;
+  const start = Date.now();
+
+  while (true) {
+    const result = await condition();
+    if (result) {
+      return;
+    }
+
+    if (Date.now() - start > timeout) {
+      throw new Error(`Condition not met within ${timeout}ms`);
+    }
+
+    await new Promise((r) => setTimeout(r, interval));
+  }
+}
+
+export async function waitForArrayLength<T>(
+  getArray: () => T[],
+  minLength: number,
+  options: WaitOptions = {},
+): Promise<void> {
+  await waitForCondition(() => getArray().length >= minLength, options);
+}
+
+export async function waitForCallCount(
+  mockFn: { mock: { calls: unknown[][] } },
+  minCalls: number,
+  options: WaitOptions = {},
+): Promise<void> {
+  await waitForCondition(() => mockFn.mock.calls.length >= minCalls, options);
+}

package/src/tree-tracker.ts

@@ -0,0 +1,173 @@
+/**
+ * TreeTracker - Git tree-based state capture for cloud/local sync
+ *
+ * Captures the entire working state as a git tree hash + archive:
+ * - Atomic state snapshots (no partial syncs)
+ * - Efficient delta detection using git's diffing
+ * - Simpler resume logic (restore tree, continue)
+ *
+ * Uses Saga pattern for atomic operations with automatic rollback on failure.
+ * Uses a temporary git index to avoid modifying the user's staging area.
+ */
+
+import { isCommitOnRemote as gitIsCommitOnRemote } from "@twig/git/queries";
+import type { PostHogAPIClient } from "./posthog-api.js";
+import { ApplySnapshotSaga } from "./sagas/apply-snapshot-saga.js";
+import { CaptureTreeSaga } from "./sagas/capture-tree-saga.js";
+import type { TreeSnapshot } from "./types.js";
+import { Logger } from "./utils/logger.js";
+
+export type { TreeSnapshot };
+
+export interface TreeTrackerConfig {
+  repositoryPath: string;
+  taskId: string;
+  runId: string;
+  apiClient?: PostHogAPIClient;
+  logger?: Logger;
+}
+
+export class TreeTracker {
+  private repositoryPath: string;
+  private taskId: string;
+  private runId: string;
+  private apiClient?: PostHogAPIClient;
+  private logger: Logger;
+  private lastTreeHash: string | null = null;
+
+  constructor(config: TreeTrackerConfig) {
+    this.repositoryPath = config.repositoryPath;
+    this.taskId = config.taskId;
+    this.runId = config.runId;
+    this.apiClient = config.apiClient;
+    this.logger =
+      config.logger || new Logger({ debug: false, prefix: "[TreeTracker]" });
+  }
+
+  /**
+   * Capture current working tree state as a snapshot.
+   * Uses a temporary index to avoid modifying user's staging area.
+   * Uses Saga pattern for atomic operation with automatic cleanup on failure.
+   */
+  async captureTree(options?: {
+    interrupted?: boolean;
+  }): Promise<TreeSnapshot | null> {
+    const saga = new CaptureTreeSaga(this.logger);
+
+    const result = await saga.run({
+      repositoryPath: this.repositoryPath,
+      taskId: this.taskId,
+      runId: this.runId,
+      apiClient: this.apiClient,
+      lastTreeHash: this.lastTreeHash,
+      interrupted: options?.interrupted,
+    });
+
+    if (!result.success) {
+      this.logger.error("Failed to capture tree", {
+        error: result.error,
+        failedStep: result.failedStep,
+      });
+      throw new Error(
+        `Failed to capture tree at step '${result.failedStep}': ${result.error}`,
+      );
+    }
+
+    // Only update lastTreeHash on success
+    if (result.data.newTreeHash !== null) {
+      this.lastTreeHash = result.data.newTreeHash;
+    }
+
+    return result.data.snapshot;
+  }
+
+  /**
+   * Download and apply a tree snapshot.
+   * Uses Saga pattern for atomic operation with rollback on failure.
+   */
+  async applyTreeSnapshot(snapshot: TreeSnapshot): Promise<void> {
+    if (!this.apiClient) {
+      throw new Error("Cannot apply snapshot: API client not configured");
+    }
+
+    if (!snapshot.archiveUrl) {
+      this.logger.warn("Cannot apply snapshot: no archive URL", {
+        treeHash: snapshot.treeHash,
+        changes: snapshot.changes.length,
+      });
+      throw new Error("Cannot apply snapshot: no archive URL");
+    }
+
+    const saga = new ApplySnapshotSaga(this.logger);
+
+    const result = await saga.run({
+      snapshot,
+      repositoryPath: this.repositoryPath,
+      apiClient: this.apiClient,
+      taskId: this.taskId,
+      runId: this.runId,
+    });
+
+    if (!result.success) {
+      this.logger.error("Failed to apply tree snapshot", {
+        error: result.error,
+        failedStep: result.failedStep,
+        treeHash: snapshot.treeHash,
+      });
+      throw new Error(
+        `Failed to apply snapshot at step '${result.failedStep}': ${result.error}`,
+      );
+    }
+
+    // Only update lastTreeHash on success
+    this.lastTreeHash = result.data.treeHash;
+  }
+
+  /**
+   * Get the last captured tree hash.
+   */
+  getLastTreeHash(): string | null {
+    return this.lastTreeHash;
+  }
+
+  /**
+   * Set the last tree hash (used when resuming).
+   */
+  setLastTreeHash(hash: string | null): void {
+    this.lastTreeHash = hash;
+  }
+}
+
+/**
+ * Check if a commit is available on any remote branch.
+ * Used to validate that cloud can fetch the base commit during handoff.
+ */
+export async function isCommitOnRemote(
+  commit: string,
+  cwd: string,
+): Promise<boolean> {
+  return gitIsCommitOnRemote(cwd, commit);
+}
+
+/**
+ * Validate that a snapshot can be handed off to cloud execution.
+ * Cloud needs to be able to fetch the baseCommit from a remote.
+ *
+ * @throws Error if the snapshot cannot be restored on cloud
+ */
+export async function validateForCloudHandoff(
+  snapshot: TreeSnapshot,
+  repositoryPath: string,
+): Promise<void> {
+  if (!snapshot.baseCommit) {
+    throw new Error("Cannot hand off to cloud: no base commit");
+  }
+
+  const onRemote = await isCommitOnRemote(snapshot.baseCommit, repositoryPath);
+  if (!onRemote) {
+    throw new Error(
+      `Cannot hand off to cloud: commit ${snapshot.baseCommit.slice(0, 7)} is not pushed. ` +
+        `Run 'git push' to push your branch first.`,
+    );
+  }
+}