@posthog/agent 1.24.2 → 1.26.0

package/src/schemas.ts

@@ -0,0 +1,241 @@
+import { z } from "zod";
+
+// Base event schema with timestamp
+const BaseEventSchema = z.object({
+  ts: z.number(),
+});
+
+// Streaming content events
+export const TokenEventSchema = BaseEventSchema.extend({
+  type: z.literal("token"),
+  content: z.string(),
+  contentType: z.enum(["text", "thinking", "tool_input"]).optional(),
+});
+
+export const ContentBlockStartEventSchema = BaseEventSchema.extend({
+  type: z.literal("content_block_start"),
+  index: z.number(),
+  contentType: z.enum(["text", "tool_use", "thinking"]),
+  toolName: z.string().optional(),
+  toolId: z.string().optional(),
+});
+
+export const ContentBlockStopEventSchema = BaseEventSchema.extend({
+  type: z.literal("content_block_stop"),
+  index: z.number(),
+});
+
+// Tool events
+export const ToolCallEventSchema = BaseEventSchema.extend({
+  type: z.literal("tool_call"),
+  toolName: z.string(),
+  callId: z.string(),
+  args: z.record(z.string(), z.unknown()),
+  parentToolUseId: z.string().nullable().optional(),
+  tool: z.unknown().optional(),
+  category: z.unknown().optional(),
+});
+
+export const ToolResultEventSchema = BaseEventSchema.extend({
+  type: z.literal("tool_result"),
+  toolName: z.string(),
+  callId: z.string(),
+  result: z.unknown(),
+  isError: z.boolean().optional(),
+  parentToolUseId: z.string().nullable().optional(),
+  tool: z.unknown().optional(),
+  category: z.unknown().optional(),
+});
+
+// Message lifecycle events
+export const MessageStartEventSchema = BaseEventSchema.extend({
+  type: z.literal("message_start"),
+  messageId: z.string().optional(),
+  model: z.string().optional(),
+});
+
+export const MessageDeltaEventSchema = BaseEventSchema.extend({
+  type: z.literal("message_delta"),
+  stopReason: z.string().optional(),
+  stopSequence: z.string().optional(),
+  usage: z
+    .object({
+      outputTokens: z.number(),
+    })
+    .optional(),
+});
+
+export const MessageStopEventSchema = BaseEventSchema.extend({
+  type: z.literal("message_stop"),
+});
+
+// User message events
+export const UserMessageEventSchema = BaseEventSchema.extend({
+  type: z.literal("user_message"),
+  content: z.string(),
+  isSynthetic: z.boolean().optional(),
+});
+
+// System events
+export const StatusEventSchema = BaseEventSchema.extend({
+  type: z.literal("status"),
+  phase: z.string(),
+  kind: z.string().optional(),
+  branch: z.string().optional(),
+  prUrl: z.string().optional(),
+  taskId: z.string().optional(),
+  messageId: z.string().optional(),
+  model: z.string().optional(),
+}).passthrough(); // Allow additional fields
+
+export const InitEventSchema = BaseEventSchema.extend({
+  type: z.literal("init"),
+  model: z.string(),
+  tools: z.array(z.string()),
+  permissionMode: z.string(),
+  cwd: z.string(),
+  apiKeySource: z.string(),
+  agents: z.array(z.string()).optional(),
+  slashCommands: z.array(z.string()).optional(),
+  outputStyle: z.string().optional(),
+  mcpServers: z
+    .array(z.object({ name: z.string(), status: z.string() }))
+    .optional(),
+});
+
+// Console event for log-style output
+export const ConsoleEventSchema = BaseEventSchema.extend({
+  type: z.literal("console"),
+  level: z.enum(["debug", "info", "warn", "error"]),
+  message: z.string(),
+});
+
+export const CompactBoundaryEventSchema = BaseEventSchema.extend({
+  type: z.literal("compact_boundary"),
+  trigger: z.enum(["manual", "auto"]),
+  preTokens: z.number(),
+});
+
+// Result events
+export const DoneEventSchema = BaseEventSchema.extend({
+  type: z.literal("done"),
+  result: z.string().optional(),
+  durationMs: z.number().optional(),
+  durationApiMs: z.number().optional(),
+  numTurns: z.number().optional(),
+  totalCostUsd: z.number().optional(),
+  usage: z.unknown().optional(),
+  modelUsage: z
+    .record(
+      z.string(),
+      z.object({
+        inputTokens: z.number(),
+        outputTokens: z.number(),
+        cacheReadInputTokens: z.number(),
+        cacheCreationInputTokens: z.number(),
+        webSearchRequests: z.number(),
+        costUSD: z.number(),
+        contextWindow: z.number(),
+      }),
+    )
+    .optional(),
+  permissionDenials: z
+    .array(
+      z.object({
+        tool_name: z.string(),
+        tool_use_id: z.string(),
+        tool_input: z.record(z.string(), z.unknown()),
+      }),
+    )
+    .optional(),
+});
+
+export const ErrorEventSchema = BaseEventSchema.extend({
+  type: z.literal("error"),
+  message: z.string(),
+  error: z.unknown().optional(),
+  errorType: z.string().optional(),
+  context: z.record(z.string(), z.unknown()).optional(),
+  sdkError: z.unknown().optional(),
+});
+
+// Metric and artifact events
+export const MetricEventSchema = BaseEventSchema.extend({
+  type: z.literal("metric"),
+  key: z.string(),
+  value: z.number(),
+  unit: z.string().optional(),
+});
+
+export const ArtifactEventSchema = BaseEventSchema.extend({
+  type: z.literal("artifact"),
+  kind: z.string(),
+  content: z.unknown(),
+});
+
+export const RawSDKEventSchema = BaseEventSchema.extend({
+  type: z.literal("raw_sdk_event"),
+  sdkMessage: z.unknown(),
+});
+
+export const AgentEventSchema = z.discriminatedUnion("type", [
+  TokenEventSchema,
+  ContentBlockStartEventSchema,
+  ContentBlockStopEventSchema,
+  ToolCallEventSchema,
+  ToolResultEventSchema,
+  MessageStartEventSchema,
+  MessageDeltaEventSchema,
+  MessageStopEventSchema,
+  UserMessageEventSchema,
+  StatusEventSchema,
+  InitEventSchema,
+  ConsoleEventSchema,
+  CompactBoundaryEventSchema,
+  DoneEventSchema,
+  ErrorEventSchema,
+  MetricEventSchema,
+  ArtifactEventSchema,
+  RawSDKEventSchema,
+]);
+
+export type TokenEvent = z.infer<typeof TokenEventSchema>;
+export type ContentBlockStartEvent = z.infer<
+  typeof ContentBlockStartEventSchema
+>;
+export type ContentBlockStopEvent = z.infer<typeof ContentBlockStopEventSchema>;
+export type ToolCallEvent = z.infer<typeof ToolCallEventSchema>;
+export type ToolResultEvent = z.infer<typeof ToolResultEventSchema>;
+export type MessageStartEvent = z.infer<typeof MessageStartEventSchema>;
+export type MessageDeltaEvent = z.infer<typeof MessageDeltaEventSchema>;
+export type MessageStopEvent = z.infer<typeof MessageStopEventSchema>;
+export type UserMessageEvent = z.infer<typeof UserMessageEventSchema>;
+export type StatusEvent = z.infer<typeof StatusEventSchema>;
+export type InitEvent = z.infer<typeof InitEventSchema>;
+export type ConsoleEvent = z.infer<typeof ConsoleEventSchema>;
+export type CompactBoundaryEvent = z.infer<typeof CompactBoundaryEventSchema>;
+export type DoneEvent = z.infer<typeof DoneEventSchema>;
+export type ErrorEvent = z.infer<typeof ErrorEventSchema>;
+export type MetricEvent = z.infer<typeof MetricEventSchema>;
+export type ArtifactEvent = z.infer<typeof ArtifactEventSchema>;
+export type RawSDKEvent = z.infer<typeof RawSDKEventSchema>;
+export type AgentEvent = z.infer<typeof AgentEventSchema>;
+
+/**
+ * Parse and validate an AgentEvent from unknown input.
+ * Returns the parsed event if valid, or null if invalid.
+ */
+export function parseAgentEvent(input: unknown): AgentEvent | null {
+  const result = AgentEventSchema.safeParse(input);
+  return result.success ? result.data : null;
+}
+
+/**
+ * Parse and validate multiple AgentEvents from an array of unknown inputs.
+ * Invalid entries are discarded.
+ */
+export function parseAgentEvents(inputs: unknown[]): AgentEvent[] {
+  return inputs
+    .map((input) => parseAgentEvent(input))
+    .filter((event): event is AgentEvent => event !== null);
+}

package/src/session-store.ts

@@ -0,0 +1,259 @@
+import type { PostHogAPIClient, TaskRunUpdate } from "./posthog-api.js";
+import type { StoredNotification, TaskRun } from "./types.js";
+import { Logger } from "./utils/logger.js";
+
+export interface SessionPersistenceConfig {
+  taskId: string;
+  runId: string;
+  logUrl: string;
+  sdkSessionId?: string;
+}
+
+export class SessionStore {
+  private posthogAPI?: PostHogAPIClient;
+  private pendingEntries: Map<string, StoredNotification[]> = new Map();
+  private flushTimeouts: Map<string, NodeJS.Timeout> = new Map();
+  private configs: Map<string, SessionPersistenceConfig> = new Map();
+  private logger: Logger;
+
+  constructor(posthogAPI?: PostHogAPIClient, logger?: Logger) {
+    this.posthogAPI = posthogAPI;
+    this.logger =
+      logger ?? new Logger({ debug: false, prefix: "[SessionStore]" });
+
+    // Flush all pending on process exit
+    const flushAllAndExit = async () => {
+      const flushPromises: Promise<void>[] = [];
+      for (const sessionId of this.configs.keys()) {
+        flushPromises.push(this.flush(sessionId));
+      }
+      await Promise.all(flushPromises);
+      process.exit(0);
+    };
+
+    process.on("beforeExit", () => {
+      flushAllAndExit().catch((e) => this.logger.error("Flush failed:", e));
+    });
+    process.on("SIGINT", () => {
+      flushAllAndExit().catch((e) => this.logger.error("Flush failed:", e));
+    });
+    process.on("SIGTERM", () => {
+      flushAllAndExit().catch((e) => this.logger.error("Flush failed:", e));
+    });
+  }
+
+  /** Register a session for persistence */
+  register(sessionId: string, config: SessionPersistenceConfig): void {
+    this.configs.set(sessionId, config);
+  }
+
+  /** Unregister and flush pending */
+  async unregister(sessionId: string): Promise<void> {
+    await this.flush(sessionId);
+    this.configs.delete(sessionId);
+  }
+
+  /** Check if a session is registered for persistence */
+  isRegistered(sessionId: string): boolean {
+    return this.configs.has(sessionId);
+  }
+
+  /**
+   * Append a raw JSON-RPC line for persistence.
+   * Parses and wraps as StoredNotification for the API.
+   */
+  appendRawLine(sessionId: string, line: string): void {
+    const config = this.configs.get(sessionId);
+    if (!config) {
+      return;
+    }
+
+    try {
+      const message = JSON.parse(line);
+      const entry: StoredNotification = {
+        type: "notification",
+        timestamp: new Date().toISOString(),
+        notification: message,
+      };
+
+      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,
+      });
+    }
+  }
+
+  /** Load raw JSON-RPC messages from S3 */
+  async load(logUrl: string): Promise<StoredNotification[]> {
+    const response = await fetch(logUrl);
+
+    if (!response.ok) {
+      return [];
+    }
+
+    const content = await response.text();
+    if (!content.trim()) return [];
+
+    return content
+      .trim()
+      .split("\n")
+      .map((line) => {
+        try {
+          return JSON.parse(line) as StoredNotification;
+        } catch {
+          return null;
+        }
+      })
+      .filter((entry): entry is StoredNotification => entry !== null);
+  }
+
+  /**
+   * Poll S3 for new entries since last check.
+   * Used for interrupt handling in cloud mode.
+   */
+  async pollForNewEntries(
+    sessionId: string,
+    lastKnownCount: number,
+  ): Promise<StoredNotification[]> {
+    const config = this.configs.get(sessionId);
+    if (!config?.logUrl) return [];
+
+    const entries = await this.load(config.logUrl);
+    return entries.slice(lastKnownCount);
+  }
+
+  /** Force flush pending entries */
+  async flush(sessionId: string): Promise<void> {
+    const config = this.configs.get(sessionId);
+    const pending = this.pendingEntries.get(sessionId);
+
+    if (!config || !pending?.length) return;
+
+    this.pendingEntries.delete(sessionId);
+    const timeout = this.flushTimeouts.get(sessionId);
+    if (timeout) {
+      clearTimeout(timeout);
+      this.flushTimeouts.delete(sessionId);
+    }
+
+    if (!this.posthogAPI) {
+      this.logger.debug("No PostHog API configured, skipping flush");
+      return;
+    }
+
+    try {
+      await this.posthogAPI.appendTaskRunLog(
+        config.taskId,
+        config.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);
+  }
+
+  /** Get the persistence config for a session */
+  getConfig(sessionId: string): SessionPersistenceConfig | undefined {
+    return this.configs.get(sessionId);
+  }
+
+  /**
+   * Start a session for persistence.
+   * Loads the task run and updates status to "in_progress".
+   */
+  async start(
+    sessionId: string,
+    taskId: string,
+    runId: string,
+  ): Promise<TaskRun | undefined> {
+    if (!this.posthogAPI) {
+      this.logger.debug(
+        "No PostHog API configured, registering session without persistence",
+      );
+      this.register(sessionId, {
+        taskId,
+        runId,
+        logUrl: "",
+      });
+      return undefined;
+    }
+
+    const taskRun = await this.posthogAPI.getTaskRun(taskId, runId);
+
+    this.register(sessionId, {
+      taskId,
+      runId,
+      logUrl: taskRun.log_url,
+    });
+
+    await this.updateTaskRun(sessionId, { status: "in_progress" });
+
+    return taskRun;
+  }
+
+  /**
+   * Mark a session as completed.
+   */
+  async complete(sessionId: string): Promise<void> {
+    await this.flush(sessionId);
+    await this.updateTaskRun(sessionId, { status: "completed" });
+  }
+
+  /**
+   * Mark a session as failed.
+   */
+  async fail(sessionId: string, error: Error | string): Promise<void> {
+    await this.flush(sessionId);
+    const message = typeof error === "string" ? error : error.message;
+    await this.updateTaskRun(sessionId, {
+      status: "failed",
+      error_message: message,
+    });
+    this.logger.error("Session failed", { sessionId, error: message });
+  }
+
+  /**
+   * Update the task run associated with a session.
+   */
+  async updateTaskRun(
+    sessionId: string,
+    update: TaskRunUpdate,
+  ): Promise<TaskRun | undefined> {
+    const config = this.configs.get(sessionId);
+    if (!config) {
+      this.logger.error(
+        `Cannot update task run: session ${sessionId} not registered`,
+      );
+      return undefined;
+    }
+
+    if (!this.posthogAPI) {
+      this.logger.debug("No PostHog API configured, skipping task run update");
+      return undefined;
+    }
+
+    try {
+      return await this.posthogAPI.updateTaskRun(
+        config.taskId,
+        config.runId,
+        update,
+      );
+    } catch (error) {
+      this.logger.error("Failed to update task run:", error);
+      return undefined;
+    }
+  }
+}

package/src/task-manager.ts

@@ -1,31 +1,31 @@
-import { randomBytes } from 'crypto';
+import { randomBytes } from "node:crypto";
 
 export interface TaskExecutionState {
   taskId: string;
-  status: 'running' | 'completed' | 'failed' | 'canceled' | 'timeout';
-  mode: 'plan_only' | 'plan_and_build' | 'build_only';
-  result?: any;
+  status: "running" | "completed" | "failed" | "canceled" | "timeout";
+  mode: "plan_only" | "plan_and_build" | "build_only";
+  result?: unknown;
   startedAt: number;
   completedAt?: number;
   abortController?: AbortController;
 }
 
 export class TaskManager {
-  private executionStates = new Map<string, TaskExecutionState>();
+  public executionStates = new Map<string, TaskExecutionState>();
   private defaultTimeout = 10 * 60 * 1000; // 10 minutes
 
   generateExecutionId(): string {
-    return randomBytes(16).toString('hex');
+    return randomBytes(16).toString("hex");
   }
 
   startExecution(
-    taskId: string, 
-    mode: 'plan_only' | 'plan_and_build' | 'build_only',
-    executionId: string = this.generateExecutionId()
+    taskId: string,
+    mode: "plan_only" | "plan_and_build" | "build_only",
+    executionId: string = this.generateExecutionId(),
   ): TaskExecutionState {
     const executionState: TaskExecutionState = {
       taskId,
-      status: 'running',
+      status: "running",
       mode,
       startedAt: Date.now(),
       abortController: new AbortController(),
@@ -33,17 +33,17 @@ export class TaskManager {
 
     this.executionStates.set(executionId, executionState);
     this.scheduleTimeout(executionId);
-    
+
     return executionState;
   }
 
-  async waitForCompletion(executionId: string): Promise<any> {
+  async waitForCompletion(executionId: string): Promise<unknown> {
     const execution = this.executionStates.get(executionId);
     if (!execution) {
       throw new Error(`Execution ${executionId} not found`);
     }
 
-    if (execution.result && execution.status === 'completed') {
+    if (execution.result && execution.status === "completed") {
       return execution.result;
     }
 
@@ -56,28 +56,33 @@ export class TaskManager {
           return;
         }
 
-        if (currentExecution.status === 'completed' && currentExecution.result) {
+        if (
+          currentExecution.status === "completed" &&
+          currentExecution.result
+        ) {
           clearInterval(checkInterval);
           resolve(currentExecution.result);
         } else if (
-          currentExecution.status === 'failed' || 
-          currentExecution.status === 'canceled' || 
-          currentExecution.status === 'timeout'
+          currentExecution.status === "failed" ||
+          currentExecution.status === "canceled" ||
+          currentExecution.status === "timeout"
         ) {
           clearInterval(checkInterval);
-          reject(new Error(`Execution ${executionId} ${currentExecution.status}`));
+          reject(
+            new Error(`Execution ${executionId} ${currentExecution.status}`),
+          );
         }
       }, 100);
     });
   }
 
-  completeExecution(executionId: string, result: any): void {
+  completeExecution(executionId: string, result: unknown): void {
     const execution = this.executionStates.get(executionId);
     if (!execution) {
       throw new Error(`Execution ${executionId} not found`);
     }
 
-    execution.status = 'completed';
+    execution.status = "completed";
     execution.result = result;
     execution.completedAt = Date.now();
   }
@@ -88,11 +93,11 @@ export class TaskManager {
       throw new Error(`Execution ${executionId} not found`);
     }
 
-    execution.status = 'failed';
+    execution.status = "failed";
     execution.completedAt = Date.now();
     execution.result = {
       error: error.message,
-      status: 'failed',
+      status: "failed",
     };
   }
 
@@ -102,14 +107,14 @@ export class TaskManager {
       throw new Error(`Execution ${executionId} not found`);
     }
 
-    execution.status = 'canceled';
+    execution.status = "canceled";
     execution.completedAt = Date.now();
     execution.abortController?.abort();
-    
+
     if (!execution.result) {
       execution.result = {
-        status: 'canceled',
-        message: 'Execution was canceled',
+        status: "canceled",
+        message: "Execution was canceled",
       };
     }
   }
@@ -126,18 +131,21 @@ export class TaskManager {
     return this.executionStates.get(executionId)?.abortController;
   }
 
-  private scheduleTimeout(executionId: string, timeout: number = this.defaultTimeout): void {
+  private scheduleTimeout(
+    executionId: string,
+    timeout: number = this.defaultTimeout,
+  ): void {
     setTimeout(() => {
       const execution = this.executionStates.get(executionId);
-      if (execution && execution.status === 'running') {
-        execution.status = 'timeout';
+      if (execution && execution.status === "running") {
+        execution.status = "timeout";
         execution.completedAt = Date.now();
         execution.abortController?.abort();
-        
+
         if (!execution.result) {
           execution.result = {
-            status: 'timeout',
-            message: 'Execution timed out',
+            status: "timeout",
+            message: "Execution timed out",
           };
         }
       }
@@ -152,4 +160,4 @@ export class TaskManager {
       }
     }
   }
-}
+}