@posthog/agent 1.30.0 → 2.0.1

package/src/server/bin.ts

@@ -0,0 +1,88 @@
+#!/usr/bin/env node
+import { Command } from "commander";
+import { z } from "zod";
+import { AgentServer } from "./agent-server.js";
+
+const envSchema = z.object({
+  JWT_PUBLIC_KEY: z
+    .string({
+      required_error:
+        "JWT_PUBLIC_KEY is required for authenticating client connections",
+    })
+    .min(1, "JWT_PUBLIC_KEY cannot be empty"),
+  POSTHOG_API_URL: z
+    .string({
+      required_error:
+        "POSTHOG_API_URL is required for LLM gateway communication",
+    })
+    .url("POSTHOG_API_URL must be a valid URL"),
+  POSTHOG_PERSONAL_API_KEY: z
+    .string({
+      required_error:
+        "POSTHOG_PERSONAL_API_KEY is required for authenticating with PostHog services",
+    })
+    .min(1, "POSTHOG_PERSONAL_API_KEY cannot be empty"),
+  POSTHOG_PROJECT_ID: z
+    .string({
+      required_error:
+        "POSTHOG_PROJECT_ID is required for routing requests to the correct project",
+    })
+    .regex(/^\d+$/, "POSTHOG_PROJECT_ID must be a numeric string")
+    .transform((val) => parseInt(val, 10)),
+});
+
+const program = new Command();
+
+program
+  .name("agent-server")
+  .description("PostHog cloud agent server - runs in sandbox environments")
+  .option("--port <port>", "HTTP server port", "3001")
+  .option(
+    "--mode <mode>",
+    "Execution mode: interactive or background",
+    "interactive",
+  )
+  .requiredOption("--repositoryPath <path>", "Path to the repository")
+  .requiredOption("--taskId <id>", "Task ID")
+  .requiredOption("--runId <id>", "Task run ID")
+  .action(async (options) => {
+    const envResult = envSchema.safeParse(process.env);
+
+    if (!envResult.success) {
+      const errors = envResult.error.issues
+        .map((issue) => `  - ${issue.message}`)
+        .join("\n");
+      program.error(`Environment validation failed:\n${errors}`);
+      return;
+    }
+
+    const env = envResult.data;
+
+    const mode = options.mode === "background" ? "background" : "interactive";
+
+    const server = new AgentServer({
+      port: parseInt(options.port, 10),
+      jwtPublicKey: env.JWT_PUBLIC_KEY,
+      repositoryPath: options.repositoryPath,
+      apiUrl: env.POSTHOG_API_URL,
+      apiKey: env.POSTHOG_PERSONAL_API_KEY,
+      projectId: env.POSTHOG_PROJECT_ID,
+      mode,
+      taskId: options.taskId,
+      runId: options.runId,
+    });
+
+    process.on("SIGINT", async () => {
+      await server.stop();
+      process.exit(0);
+    });
+
+    process.on("SIGTERM", async () => {
+      await server.stop();
+      process.exit(0);
+    });
+
+    await server.start();
+  });
+
+program.parse();

package/src/server/jwt.ts

@@ -0,0 +1,65 @@
+import jwt from "jsonwebtoken";
+import { z } from "zod";
+
+export const SANDBOX_CONNECTION_AUDIENCE = "posthog:sandbox_connection";
+
+export const userDataSchema = z.object({
+  run_id: z.string(),
+  task_id: z.string(),
+  team_id: z.number(),
+  user_id: z.number(),
+  distinct_id: z.string(),
+  mode: z.enum(["interactive", "background"]).optional().default("interactive"),
+});
+
+const jwtPayloadSchema = userDataSchema.extend({
+  exp: z.number(),
+  iat: z.number().optional(),
+  aud: z.string().optional(),
+});
+
+export type JwtPayload = z.infer<typeof userDataSchema>;
+
+export class JwtValidationError extends Error {
+  constructor(
+    message: string,
+    public code:
+      | "invalid_token"
+      | "expired"
+      | "invalid_signature"
+      | "server_error",
+  ) {
+    super(message);
+    this.name = "JwtValidationError";
+  }
+}
+
+export function validateJwt(token: string, publicKey: string): JwtPayload {
+  try {
+    const decoded = jwt.verify(token, publicKey, {
+      algorithms: ["RS256"],
+      audience: SANDBOX_CONNECTION_AUDIENCE,
+    });
+
+    const result = jwtPayloadSchema.safeParse(decoded);
+    if (!result.success) {
+      throw new JwtValidationError(
+        `Missing required fields: ${result.error.message}`,
+        "invalid_token",
+      );
+    }
+
+    return result.data;
+  } catch (error) {
+    if (error instanceof JwtValidationError) {
+      throw error;
+    }
+    if (error instanceof jwt.TokenExpiredError) {
+      throw new JwtValidationError("Token expired", "expired");
+    }
+    if (error instanceof jwt.JsonWebTokenError) {
+      throw new JwtValidationError("Invalid signature", "invalid_signature");
+    }
+    throw new JwtValidationError("Invalid token", "invalid_token");
+  }
+}

package/src/server/schemas.ts

@@ -0,0 +1,47 @@
+import { z } from "zod";
+
+export const jsonRpcRequestSchema = z.object({
+  jsonrpc: z.literal("2.0"),
+  method: z.string(),
+  params: z.record(z.unknown()).optional(),
+  id: z.union([z.string(), z.number()]).optional(),
+});
+
+export type JsonRpcRequest = z.infer<typeof jsonRpcRequestSchema>;
+
+export const userMessageParamsSchema = z.object({
+  content: z.string().min(1, "Content is required"),
+});
+
+export const commandParamsSchemas = {
+  user_message: userMessageParamsSchema,
+  "posthog/user_message": userMessageParamsSchema,
+  cancel: z.object({}).optional(),
+  "posthog/cancel": z.object({}).optional(),
+  close: z.object({}).optional(),
+  "posthog/close": z.object({}).optional(),
+} as const;
+
+export type CommandMethod = keyof typeof commandParamsSchemas;
+
+export function validateCommandParams(
+  method: string,
+  params: unknown,
+): { success: true; data: unknown } | { success: false; error: string } {
+  const schema =
+    commandParamsSchemas[method as CommandMethod] ??
+    commandParamsSchemas[
+      method.replace("posthog/", "") as keyof typeof commandParamsSchemas
+    ];
+
+  if (!schema) {
+    return { success: false, error: `Unknown method: ${method}` };
+  }
+
+  const result = schema.safeParse(params);
+  if (!result.success) {
+    return { success: false, error: result.error.message };
+  }
+
+  return { success: true, data: result.data };
+}

package/src/server/types.ts

@@ -0,0 +1,13 @@
+import type { AgentMode } from "../types.js";
+
+export interface AgentServerConfig {
+  port: number;
+  repositoryPath: string;
+  apiUrl: string;
+  apiKey: string;
+  projectId: number;
+  jwtPublicKey: string; // RS256 public key for JWT verification
+  mode: AgentMode;
+  taskId: string;
+  runId: string;
+}

package/src/server/utils/retry.test.ts

@@ -0,0 +1,122 @@
+import { describe, expect, it, vi } from "vitest";
+import { retry } from "./retry.js";
+
+describe("retry", () => {
+  it("returns result on first success", async () => {
+    const fn = vi.fn().mockResolvedValue("success");
+
+    const result = await retry(fn);
+
+    expect(result).toBe("success");
+    expect(fn).toHaveBeenCalledTimes(1);
+  });
+
+  it("retries on transient error and succeeds", async () => {
+    const fn = vi
+      .fn()
+      .mockRejectedValueOnce(new Error("network error"))
+      .mockResolvedValue("success");
+
+    const result = await retry(fn, { baseDelayMs: 1 });
+
+    expect(result).toBe("success");
+    expect(fn).toHaveBeenCalledTimes(2);
+  });
+
+  it("throws after max attempts", async () => {
+    const fn = vi.fn().mockRejectedValue(new Error("network error"));
+
+    await expect(retry(fn, { maxAttempts: 3, baseDelayMs: 1 })).rejects.toThrow(
+      "network error",
+    );
+
+    expect(fn).toHaveBeenCalledTimes(3);
+  });
+
+  it("does not retry on non-transient error", async () => {
+    const fn = vi.fn().mockRejectedValue(new Error("validation failed"));
+
+    await expect(retry(fn, { baseDelayMs: 1 })).rejects.toThrow(
+      "validation failed",
+    );
+
+    expect(fn).toHaveBeenCalledTimes(1);
+  });
+
+  it("respects custom shouldRetry", async () => {
+    const fn = vi
+      .fn()
+      .mockRejectedValueOnce(new Error("custom error"))
+      .mockResolvedValue("success");
+
+    const result = await retry(fn, {
+      baseDelayMs: 1,
+      shouldRetry: (err) => err.message === "custom error",
+    });
+
+    expect(result).toBe("success");
+    expect(fn).toHaveBeenCalledTimes(2);
+  });
+
+  it("uses exponential backoff", async () => {
+    const fn = vi
+      .fn()
+      .mockRejectedValueOnce(new Error("timeout"))
+      .mockRejectedValueOnce(new Error("timeout"))
+      .mockResolvedValue("success");
+
+    const start = Date.now();
+    await retry(fn, { baseDelayMs: 50, maxDelayMs: 200 });
+    const elapsed = Date.now() - start;
+
+    // First retry: 50ms, second retry: 100ms = 150ms minimum
+    expect(elapsed).toBeGreaterThanOrEqual(100);
+  });
+
+  it("caps delay at maxDelayMs", async () => {
+    const fn = vi
+      .fn()
+      .mockRejectedValueOnce(new Error("timeout"))
+      .mockRejectedValueOnce(new Error("timeout"))
+      .mockRejectedValueOnce(new Error("timeout"))
+      .mockResolvedValue("success");
+
+    const start = Date.now();
+    await retry(fn, { maxAttempts: 4, baseDelayMs: 100, maxDelayMs: 150 });
+    const elapsed = Date.now() - start;
+
+    // With cap: 100 + 150 + 150 = 400ms max (not 100 + 200 + 400 = 700ms)
+    expect(elapsed).toBeLessThan(600);
+  });
+
+  it("retries on 429 rate limit", async () => {
+    const fn = vi
+      .fn()
+      .mockRejectedValueOnce(new Error("429 Too Many Requests"))
+      .mockResolvedValue("success");
+
+    const result = await retry(fn, { baseDelayMs: 1 });
+
+    expect(result).toBe("success");
+    expect(fn).toHaveBeenCalledTimes(2);
+  });
+
+  it("retries on 502/503 server errors", async () => {
+    const fn = vi
+      .fn()
+      .mockRejectedValueOnce(new Error("502 Bad Gateway"))
+      .mockRejectedValueOnce(new Error("503 Service Unavailable"))
+      .mockResolvedValue("success");
+
+    const result = await retry(fn, { baseDelayMs: 1 });
+
+    expect(result).toBe("success");
+    expect(fn).toHaveBeenCalledTimes(3);
+  });
+
+  it("converts non-Error throws to Error", async () => {
+    const fn = vi.fn().mockRejectedValue("string error");
+
+    await expect(retry(fn, { maxAttempts: 1 })).rejects.toThrow("string error");
+  });
+});

package/src/server/utils/retry.ts

@@ -0,0 +1,61 @@
+export interface RetryOptions {
+  maxAttempts?: number;
+  baseDelayMs?: number;
+  maxDelayMs?: number;
+  shouldRetry?: (error: Error) => boolean;
+}
+
+const DEFAULT_OPTIONS: Required<Omit<RetryOptions, "shouldRetry">> = {
+  maxAttempts: 3,
+  baseDelayMs: 1000,
+  maxDelayMs: 10000,
+};
+
+function isTransientError(error: Error): boolean {
+  const message = error.message.toLowerCase();
+  return (
+    message.includes("network") ||
+    message.includes("timeout") ||
+    message.includes("econnreset") ||
+    message.includes("econnrefused") ||
+    message.includes("socket") ||
+    message.includes("503") ||
+    message.includes("502") ||
+    message.includes("429")
+  );
+}
+
+export async function retry<T>(
+  fn: () => Promise<T>,
+  options: RetryOptions = {},
+): Promise<T> {
+  const opts = { ...DEFAULT_OPTIONS, ...options };
+  const shouldRetry = opts.shouldRetry ?? isTransientError;
+
+  let lastError: Error | null = null;
+
+  for (let attempt = 1; attempt <= opts.maxAttempts; attempt++) {
+    try {
+      return await fn();
+    } catch (error) {
+      lastError = error instanceof Error ? error : new Error(String(error));
+
+      const isLastAttempt = attempt === opts.maxAttempts;
+      if (isLastAttempt || !shouldRetry(lastError)) {
+        throw lastError;
+      }
+
+      const delay = Math.min(
+        opts.baseDelayMs * 2 ** (attempt - 1),
+        opts.maxDelayMs,
+      );
+      await sleep(delay);
+    }
+  }
+
+  throw lastError ?? new Error("Retry failed with no error");
+}
+
+function sleep(ms: number): Promise<void> {
+  return new Promise((resolve) => setTimeout(resolve, ms));
+}

package/src/server/utils/sse-parser.test.ts

@@ -0,0 +1,93 @@
+import { describe, expect, it } from "vitest";
+import { SseEventParser } from "./sse-parser.js";
+
+describe("SseEventParser", () => {
+  it("parses complete SSE event", () => {
+    const parser = new SseEventParser();
+    const events = parser.parse('data: {"message":"hello"}\n\n');
+
+    expect(events).toHaveLength(1);
+    expect(events[0].data).toEqual({ message: "hello" });
+  });
+
+  it("parses event with id", () => {
+    const parser = new SseEventParser();
+    const events = parser.parse('id: 123\ndata: {"test":true}\n\n');
+
+    expect(events).toHaveLength(1);
+    expect(events[0].id).toBe("123");
+    expect(events[0].data).toEqual({ test: true });
+  });
+
+  it("handles chunked data", () => {
+    const parser = new SseEventParser();
+
+    const events1 = parser.parse("id: 1\n");
+    expect(events1).toHaveLength(0);
+
+    const events2 = parser.parse('data: {"part":"one"}');
+    expect(events2).toHaveLength(0);
+
+    const events3 = parser.parse("\n\n");
+    expect(events3).toHaveLength(1);
+    expect(events3[0].id).toBe("1");
+    expect(events3[0].data).toEqual({ part: "one" });
+  });
+
+  it("parses multiple events in one chunk", () => {
+    const parser = new SseEventParser();
+    const events = parser.parse('data: {"first":1}\n\ndata: {"second":2}\n\n');
+
+    expect(events).toHaveLength(2);
+    expect(events[0].data).toEqual({ first: 1 });
+    expect(events[1].data).toEqual({ second: 2 });
+  });
+
+  it("skips malformed JSON", () => {
+    const parser = new SseEventParser();
+    const events = parser.parse('data: not json\n\ndata: {"valid":true}\n\n');
+
+    expect(events).toHaveLength(1);
+    expect(events[0].data).toEqual({ valid: true });
+  });
+
+  it("handles empty data lines", () => {
+    const parser = new SseEventParser();
+    const events = parser.parse("data: \n\n");
+
+    expect(events).toHaveLength(0);
+  });
+
+  it("resets state correctly", () => {
+    const parser = new SseEventParser();
+
+    parser.parse("id: 1\ndata: {");
+    parser.reset();
+
+    const events = parser.parse('data: {"fresh":true}\n\n');
+
+    expect(events).toHaveLength(1);
+    expect(events[0].id).toBeUndefined();
+    expect(events[0].data).toEqual({ fresh: true });
+  });
+
+  it("handles events with whitespace in id", () => {
+    const parser = new SseEventParser();
+    const events = parser.parse('id:   abc123  \ndata: {"test":1}\n\n');
+
+    expect(events).toHaveLength(1);
+    expect(events[0].id).toBe("abc123");
+  });
+
+  it("preserves incomplete line in buffer", () => {
+    const parser = new SseEventParser();
+
+    const events1 = parser.parse('data: {"complete":tr');
+    expect(events1).toHaveLength(0);
+
+    const events2 = parser.parse('ue}\n\ndata: {"next":1}\n\n');
+    expect(events2).toHaveLength(2);
+    expect(events2[0].data).toEqual({ complete: true });
+    expect(events2[1].data).toEqual({ next: 1 });
+  });
+});

package/src/server/utils/sse-parser.ts

@@ -0,0 +1,46 @@
+export interface SseEvent {
+  id?: string;
+  data: unknown;
+}
+
+export class SseEventParser {
+  private buffer = "";
+  private currentEventId: string | null = null;
+  private currentData: string | null = null;
+
+  parse(chunk: string): SseEvent[] {
+    this.buffer += chunk;
+    const lines = this.buffer.split("\n");
+    this.buffer = lines.pop() || "";
+
+    const events: SseEvent[] = [];
+
+    for (const line of lines) {
+      if (line.startsWith("id: ")) {
+        this.currentEventId = line.slice(4).trim();
+      } else if (line.startsWith("data: ")) {
+        this.currentData = line.slice(6);
+      } else if (line === "" && this.currentData !== null) {
+        try {
+          const data = JSON.parse(this.currentData);
+          events.push({
+            id: this.currentEventId ?? undefined,
+            data,
+          });
+        } catch {
+          // Skip malformed data
+        }
+        this.currentData = null;
+        this.currentEventId = null;
+      }
+    }
+
+    return events;
+  }
+
+  reset(): void {
+    this.buffer = "";
+    this.currentEventId = null;
+    this.currentData = null;
+  }
+}

package/src/session-log-writer.test.ts

@@ -0,0 +1,140 @@
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import { OtelLogWriter } from "./otel-log-writer.js";
+import { SessionLogWriter } from "./session-log-writer.js";
+
+// Mock the OtelLogWriter
+vi.mock("./otel-log-writer.js", () => ({
+  OtelLogWriter: vi.fn(),
+}));
+
+describe("SessionLogWriter", () => {
+  let logWriter: SessionLogWriter;
+  let mockEmit: ReturnType<typeof vi.fn>;
+  let mockFlush: ReturnType<typeof vi.fn>;
+  let mockShutdown: ReturnType<typeof vi.fn>;
+
+  beforeEach(() => {
+    mockEmit = vi.fn();
+    mockFlush = vi.fn().mockResolvedValue(undefined);
+    mockShutdown = vi.fn().mockResolvedValue(undefined);
+
+    vi.mocked(OtelLogWriter).mockImplementation(
+      () =>
+        ({
+          emit: mockEmit,
+          flush: mockFlush,
+          shutdown: mockShutdown,
+        }) as unknown as OtelLogWriter,
+    );
+
+    logWriter = new SessionLogWriter({
+      otelConfig: {
+        posthogHost: "http://localhost:8000",
+        apiKey: "test-api-key",
+        logsPath: "/i/v1/agent-logs",
+      },
+    });
+  });
+
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
+  describe("appendRawLine", () => {
+    it("emits entries immediately via OtelLogWriter", () => {
+      const sessionId = "test-session";
+      logWriter.register(sessionId, { taskId: "task-1", runId: sessionId });
+
+      logWriter.appendRawLine(
+        sessionId,
+        JSON.stringify({ method: "test", params: {} }),
+      );
+      logWriter.appendRawLine(
+        sessionId,
+        JSON.stringify({ method: "test2", params: {} }),
+      );
+
+      expect(mockEmit).toHaveBeenCalledTimes(2);
+    });
+
+    it("wraps raw messages in StoredNotification format", () => {
+      const sessionId = "test-session";
+      logWriter.register(sessionId, { taskId: "task-1", runId: sessionId });
+
+      const message = {
+        jsonrpc: "2.0",
+        method: "session/update",
+        params: { foo: "bar" },
+      };
+      logWriter.appendRawLine(sessionId, JSON.stringify(message));
+
+      expect(mockEmit).toHaveBeenCalledTimes(1);
+      const emitArg = mockEmit.mock.calls[0][0];
+      expect(emitArg.notification.type).toBe("notification");
+      expect(emitArg.notification.timestamp).toBeDefined();
+      expect(emitArg.notification.notification).toEqual(message);
+    });
+
+    it("ignores unregistered sessions", () => {
+      logWriter.appendRawLine(
+        "unknown-session",
+        JSON.stringify({ method: "test" }),
+      );
+
+      expect(mockEmit).not.toHaveBeenCalled();
+    });
+
+    it("ignores invalid JSON", () => {
+      const sessionId = "test-session";
+      logWriter.register(sessionId, { taskId: "task-1", runId: sessionId });
+
+      logWriter.appendRawLine(sessionId, "not valid json {{{");
+
+      expect(mockEmit).not.toHaveBeenCalled();
+    });
+  });
+
+  describe("flush", () => {
+    it("calls flush on OtelLogWriter", async () => {
+      const sessionId = "test-session";
+      logWriter.register(sessionId, { taskId: "task-1", runId: sessionId });
+
+      logWriter.appendRawLine(sessionId, JSON.stringify({ method: "test" }));
+      await logWriter.flush(sessionId);
+
+      expect(mockFlush).toHaveBeenCalledTimes(1);
+    });
+
+    it("does nothing for unregistered sessions", async () => {
+      await logWriter.flush("unknown-session");
+
+      expect(mockFlush).not.toHaveBeenCalled();
+    });
+  });
+
+  describe("register", () => {
+    it("creates OtelLogWriter with session context", () => {
+      const sessionId = "test-session";
+      const context = { taskId: "task-1", runId: sessionId };
+
+      logWriter.register(sessionId, context);
+
+      expect(OtelLogWriter).toHaveBeenCalledWith(
+        expect.objectContaining({
+          posthogHost: "http://localhost:8000",
+          apiKey: "test-api-key",
+        }),
+        context,
+        expect.anything(),
+      );
+    });
+
+    it("does not re-register existing sessions", () => {
+      const sessionId = "test-session";
+      logWriter.register(sessionId, { taskId: "task-1", runId: sessionId });
+      logWriter.register(sessionId, { taskId: "task-2", runId: sessionId });
+
+      expect(OtelLogWriter).toHaveBeenCalledTimes(1);
+    });
+  });
+});