use zod for agent events to make sure they conform to schema

Author: joshsny

apps/array/src/renderer/features/task-detail/stores/taskExecutionStore.ts

@@ -1,7 +1,7 @@
 import { useAuthStore } from "@features/auth/stores/authStore";
 import { usePanelLayoutStore } from "@features/panels/store/panelLayoutStore";
 import { useSettingsStore } from "@features/settings/stores/settingsStore";
-import type { AgentEvent } from "@posthog/agent";
+import { type AgentEvent, parseAgentEvents } from "@posthog/agent";
 import { track } from "@renderer/lib/analytics";
 import { logger } from "@renderer/lib/logger";
 import { queryClient } from "@renderer/lib/queryClient";
@@ -94,10 +94,19 @@ async function fetchLogsFromS3Url(logUrl: string): Promise<AgentEvent[]> {
       return [];
     }
 
-    return content
+    const rawEntries = content
       .trim()
       .split("\n")
-      .map((line: string) => JSON.parse(line) as AgentEvent);
+      .map((line: string) => {
+        try {
+          return JSON.parse(line);
+        } catch {
+          return null;
+        }
+      })
+      .filter(Boolean);
+
+    return parseAgentEvents(rawEntries);
   } catch (err) {
     log.warn("Failed to fetch task logs from S3", err);
     return [];

packages/agent/example.ts

@@ -1,8 +1,8 @@
 #!/usr/bin/env bun
 
-import { config } from "dotenv";
 import { execSync } from "node:child_process";
 import * as readline from "node:readline";
+import { config } from "dotenv";
 
 config();
 
@@ -40,7 +40,7 @@ Create a simple hello.txt file in the root of the repository with the text "Hell
 `;
 
 async function testAgent() {
-    let TASK_ID = process.argv[2];
+  let TASK_ID = process.argv[2];
 
   if (!process.env.POSTHOG_API_KEY) {
     console.error("❌ POSTHOG_API_KEY required");
@@ -57,7 +57,7 @@ async function testAgent() {
     process.exit(1);
   }
 
-  const REPO_PATH = process.env.REPO_PATH
+  const REPO_PATH = process.env.REPO_PATH;
 
   if (!REPO_PATH) {
     console.error("❌ REPO_PATH required");
@@ -69,7 +69,9 @@ async function testAgent() {
   // Check for uncommitted changes
   if (hasUncommittedChanges(REPO_PATH)) {
     console.log("⚠️  Warning: There are uncommitted changes in the repository.");
-    const proceed = await promptUser("These changes will be discarded. Continue? (y/n): ");
+    const proceed = await promptUser(
+      "These changes will be discarded. Continue? (y/n): ",
+    );
     if (!proceed) {
       console.log("❌ Aborted.");
       process.exit(0);

packages/agent/index.ts

@@ -61,7 +61,12 @@ export type {
   UserMessageEvent,
   WorktreeInfo,
 } from "./src/types.js";
-export { PermissionMode } from "./src/types.js";
+export {
+  AgentEventSchema,
+  PermissionMode,
+  parseAgentEvent,
+  parseAgentEvents,
+} from "./src/types.js";
 export type { LoggerConfig } from "./src/utils/logger.js";
 export {
   Logger,

packages/agent/package.json

@@ -39,7 +39,6 @@
   },
   "devDependencies": {
     "@changesets/cli": "^2.27.8",
-    "tsx": "^4.20.6",
     "@rollup/plugin-commonjs": "^25.0.7",
     "@rollup/plugin-node-resolve": "^15.2.3",
     "@types/bun": "latest",
@@ -48,12 +47,14 @@
     "rollup": "^4.24.0",
     "rollup-plugin-copy": "^3.5.0",
     "rollup-plugin-typescript2": "^0.36.0",
+    "tsx": "^4.20.6",
     "typescript": "^5.5.0"
   },
   "dependencies": {
     "@anthropic-ai/claude-agent-sdk": "^0.1.47",
     "dotenv": "^17.2.3",
-    "yoga-wasm-web": "^0.3.3"
+    "yoga-wasm-web": "^0.3.3",
+    "zod": "^4.1.12"
   },
   "files": [
     "dist/**/*",

packages/agent/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);
+}