@posthog/agent 2.0.0 → 2.0.1

package/src/session-log-writer.ts

@@ -0,0 +1,137 @@
+import {
+  type OtelLogConfig,
+  OtelLogWriter,
+  type SessionContext,
+} from "./otel-log-writer.js";
+import type { PostHogAPIClient } from "./posthog-api.js";
+import type { StoredNotification } from "./types.js";
+import { Logger } from "./utils/logger.js";
+
+export interface SessionLogWriterOptions {
+  /** OTEL config for creating writers per session */
+  otelConfig?: OtelLogConfig;
+  /** PostHog API client for S3 log persistence */
+  posthogAPI?: PostHogAPIClient;
+  /** Logger instance */
+  logger?: Logger;
+}
+
+interface SessionState {
+  context: SessionContext;
+  otelWriter?: OtelLogWriter;
+}
+
+export class SessionLogWriter {
+  private posthogAPI?: PostHogAPIClient;
+  private otelConfig?: OtelLogConfig;
+  private pendingEntries: Map<string, StoredNotification[]> = new Map();
+  private flushTimeouts: Map<string, NodeJS.Timeout> = new Map();
+  private sessions: Map<string, SessionState> = new Map();
+  private logger: Logger;
+
+  constructor(options: SessionLogWriterOptions = {}) {
+    this.posthogAPI = options.posthogAPI;
+    this.otelConfig = options.otelConfig;
+    this.logger =
+      options.logger ??
+      new Logger({ debug: false, prefix: "[SessionLogWriter]" });
+  }
+
+  async flushAll(): Promise<void> {
+    const flushPromises: Promise<void>[] = [];
+    for (const sessionId of this.sessions.keys()) {
+      flushPromises.push(this.flush(sessionId));
+    }
+    await Promise.all(flushPromises);
+  }
+
+  register(sessionId: string, context: SessionContext): void {
+    if (this.sessions.has(sessionId)) {
+      return;
+    }
+
+    let otelWriter: OtelLogWriter | undefined;
+    if (this.otelConfig) {
+      // Create a dedicated OtelLogWriter for this session with resource attributes
+      otelWriter = new OtelLogWriter(
+        this.otelConfig,
+        context,
+        this.logger.child(`OtelWriter:${sessionId}`),
+      );
+    }
+
+    this.sessions.set(sessionId, { context, otelWriter });
+  }
+
+  isRegistered(sessionId: string): boolean {
+    return this.sessions.has(sessionId);
+  }
+
+  appendRawLine(sessionId: string, line: string): void {
+    const session = this.sessions.get(sessionId);
+    if (!session) {
+      return;
+    }
+
+    try {
+      const message = JSON.parse(line);
+      const entry: StoredNotification = {
+        type: "notification",
+        timestamp: new Date().toISOString(),
+        notification: message,
+      };
+
+      if (session.otelWriter) {
+        session.otelWriter.emit({ notification: entry });
+      }
+
+      if (this.posthogAPI) {
+        const pending = this.pendingEntries.get(sessionId) ?? [];
+        pending.push(entry);
+        this.pendingEntries.set(sessionId, pending);
+        this.scheduleFlush(sessionId);
+      }
+    } catch {
+      this.logger.warn("Failed to parse raw line for persistence", {
+        sessionId,
+        lineLength: line.length,
+      });
+    }
+  }
+
+  async flush(sessionId: string): Promise<void> {
+    const session = this.sessions.get(sessionId);
+    if (!session) return;
+
+    if (session.otelWriter) {
+      await session.otelWriter.flush();
+    }
+
+    const pending = this.pendingEntries.get(sessionId);
+    if (!this.posthogAPI || !pending?.length) return;
+
+    this.pendingEntries.delete(sessionId);
+    const timeout = this.flushTimeouts.get(sessionId);
+    if (timeout) {
+      clearTimeout(timeout);
+      this.flushTimeouts.delete(sessionId);
+    }
+
+    try {
+      await this.posthogAPI.appendTaskRunLog(
+        session.context.taskId,
+        session.context.runId,
+        pending,
+      );
+    } catch (error) {
+      this.logger.error("Failed to persist session logs:", error);
+    }
+  }
+
+  private scheduleFlush(sessionId: string): void {
+    const existing = this.flushTimeouts.get(sessionId);
+    if (existing) clearTimeout(existing);
+    const timeout = setTimeout(() => this.flush(sessionId), 500);
+    this.flushTimeouts.set(sessionId, timeout);
+  }
+}

package/src/test/assertions.ts

@@ -0,0 +1,114 @@
+import { expect } from "vitest";
+
+export interface NotificationEntry {
+  notification?: {
+    method?: string;
+    params?: {
+      sessionId?: string;
+      update?: {
+        sessionUpdate?: string;
+        content?: {
+          type?: string;
+          text?: string;
+        };
+      };
+    };
+  };
+}
+
+export interface NotificationMatcher {
+  method?: string;
+  text?: string;
+  sessionId?: string;
+  sessionUpdate?: string;
+}
+
+function entryMatchesNotification(
+  entry: unknown,
+  matcher: NotificationMatcher,
+): boolean {
+  const notification = (entry as NotificationEntry).notification;
+  if (!notification) return false;
+
+  if (matcher.method && notification.method !== matcher.method) {
+    return false;
+  }
+  if (
+    matcher.sessionId &&
+    notification.params?.sessionId !== matcher.sessionId
+  ) {
+    return false;
+  }
+  if (
+    matcher.sessionUpdate &&
+    notification.params?.update?.sessionUpdate !== matcher.sessionUpdate
+  ) {
+    return false;
+  }
+  if (matcher.text) {
+    const text = notification.params?.update?.content?.text;
+    if (!text || !text.includes(matcher.text)) {
+      return false;
+    }
+  }
+  return true;
+}
+
+export function findNotification(
+  appendLogCalls: unknown[][],
+  matcher: NotificationMatcher,
+): NotificationEntry | undefined {
+  for (const entries of appendLogCalls) {
+    for (const entry of entries) {
+      if (entryMatchesNotification(entry, matcher)) {
+        return entry as NotificationEntry;
+      }
+    }
+  }
+  return undefined;
+}
+
+export function hasNotification(
+  appendLogCalls: unknown[][],
+  matcher: NotificationMatcher,
+): boolean {
+  return findNotification(appendLogCalls, matcher) !== undefined;
+}
+
+export function expectNotification(
+  appendLogCalls: unknown[][],
+  matcher: NotificationMatcher,
+): NotificationEntry {
+  const found = findNotification(appendLogCalls, matcher);
+  expect(
+    found,
+    `Expected notification matching ${JSON.stringify(matcher)}`,
+  ).toBeDefined();
+  return found!;
+}
+
+export function expectNoNotification(
+  appendLogCalls: unknown[][],
+  matcher: NotificationMatcher,
+): void {
+  const found = findNotification(appendLogCalls, matcher);
+  expect(
+    found,
+    `Expected no notification matching ${JSON.stringify(matcher)}`,
+  ).toBeUndefined();
+}
+
+export function countNotifications(
+  appendLogCalls: unknown[][],
+  matcher: NotificationMatcher,
+): number {
+  let count = 0;
+  for (const entries of appendLogCalls) {
+    for (const entry of entries) {
+      if (entryMatchesNotification(entry, matcher)) {
+        count++;
+      }
+    }
+  }
+  return count;
+}

package/src/test/controllers/sse-controller.ts

@@ -0,0 +1,107 @@
+type SseState = "idle" | "streaming" | "closing" | "closed";
+
+export interface EventOptions {
+  id?: string;
+  event?: string;
+}
+
+export class SseController {
+  private encoder = new TextEncoder();
+  private controller: ReadableStreamDefaultController<Uint8Array> | null = null;
+  private state: SseState = "idle";
+
+  createStream(): ReadableStream<Uint8Array> {
+    return new ReadableStream({
+      start: (controller) => {
+        this.controller = controller;
+        this.state = "streaming";
+      },
+      cancel: () => {
+        this.state = "closed";
+      },
+    });
+  }
+
+  sendEvent(data: unknown, options: EventOptions = {}): boolean {
+    if (
+      this.state === "closed" ||
+      this.state === "closing" ||
+      !this.controller
+    ) {
+      return false;
+    }
+
+    const lines: string[] = [];
+    if (options.id) {
+      lines.push(`id: ${options.id}`);
+    }
+    if (options.event) {
+      lines.push(`event: ${options.event}`);
+    }
+    lines.push(`data: ${JSON.stringify(data)}`);
+    lines.push("");
+    lines.push("");
+
+    try {
+      this.controller.enqueue(this.encoder.encode(lines.join("\n")));
+      return true;
+    } catch {
+      this.state = "closed";
+      return false;
+    }
+  }
+
+  sendRaw(rawData: string): boolean {
+    if (
+      this.state === "closed" ||
+      this.state === "closing" ||
+      !this.controller
+    ) {
+      return false;
+    }
+    try {
+      this.controller.enqueue(this.encoder.encode(rawData));
+      return true;
+    } catch {
+      this.state = "closed";
+      return false;
+    }
+  }
+
+  sendPartial(partialData: string): boolean {
+    return this.sendRaw(partialData);
+  }
+
+  error(err: Error): void {
+    if (this.state !== "streaming" || !this.controller) return;
+
+    try {
+      this.state = "closed";
+      this.controller.error(err);
+    } catch {
+      // Already errored or closed
+    }
+  }
+
+  close(): void {
+    if (this.state === "closed") return;
+
+    this.state = "closing";
+    if (this.controller) {
+      try {
+        this.controller.close();
+      } catch {
+        // Already closed
+      }
+    }
+    this.state = "closed";
+  }
+
+  get closed(): boolean {
+    return this.state === "closed";
+  }
+
+  get currentState(): SseState {
+    return this.state;
+  }
+}

package/src/test/fixtures/api.ts

@@ -0,0 +1,111 @@
+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 { vi } from "vitest";
+import type { PostHogAPIClient } from "../../posthog-api.js";
+import type { 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 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 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 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,
+  };
+}

package/src/test/fixtures/config.ts

@@ -0,0 +1,33 @@
+import type { AgentServerConfig } from "../../server/types.js";
+import type { TestRepo } from "./api.js";
+
+export type { AgentServerConfig };
+
+// Test RSA public key (for testing only - matches TEST_PRIVATE_KEY in agent-server.test.ts)
+export const TEST_PUBLIC_KEY = `-----BEGIN PUBLIC KEY-----
+MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEA6ofeEmDBbLxuAgqPQUho
+7T69vzsc7jRq+NTuPgDJA0EXfaSjuPQ4UwOFc8Jzr6x/MuHiPTfDkJ3uwcKXaKLk
+p+A6AwEv290lH4o/0aBVEsmYk0KFs9B+qNlbrn4s9B3/gc5WRFZ4UkNa7r6kn/uJ
+fHFoHjF2hV4HQ+ZEPBo70ebqisbzthJ79YTCUSnjjhBoAnqf9HOkpDFE10ngvlY8
+qVYPfvMj8bSKTkO1yr/u3vzwNIpanoUUIeH6PQQFo1Ftfh527bIyQI43754MyI6W
+o7kFcjIuxu/b/Dr4o4SzCYyQYd03W1SH4vkZFY/x/eFFHylkXyQNHi8pAFb04hX9
+JwIDAQAB
+-----END PUBLIC KEY-----`;
+
+export function createAgentServerConfig(
+  repo: TestRepo,
+  overrides: Partial<AgentServerConfig> = {},
+): AgentServerConfig {
+  return {
+    port: 3001,
+    repositoryPath: repo.path,
+    apiUrl: "http://localhost:8000",
+    apiKey: "test-api-key",
+    projectId: 1,
+    jwtPublicKey: TEST_PUBLIC_KEY,
+    mode: "interactive",
+    taskId: "test-task-id",
+    runId: "test-run-id",
+    ...overrides,
+  };
+}

package/src/test/fixtures/notifications.ts

@@ -0,0 +1,92 @@
+import { POSTHOG_NOTIFICATIONS } from "../../acp-extensions.js";
+import type { StoredNotification } from "../../types.js";
+
+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,
+  });
+}
+
+export function createStatusNotification(
+  status: "connected" | "disconnected" | "error",
+  message?: string,
+): StoredNotification {
+  return createNotification("session/update", {
+    update: {
+      sessionUpdate: "status",
+      status,
+      message,
+    },
+  });
+}