feat: Add initial structured output for agent (#1466)

Author: ryans-posthog

apps/code/src/main/services/agent/schemas.ts

@@ -50,6 +50,7 @@ export const startSessionInput = z.object({
   customInstructions: z.string().max(2000).optional(),
   effort: effortLevelSchema.optional(),
   model: z.string().optional(),
+  jsonSchema: z.record(z.string(), z.unknown()).nullish(),
 });
 
 export type StartSessionInput = z.infer<typeof startSessionInput>;
@@ -173,6 +174,7 @@ export const reconnectSessionInput = z.object({
   permissionMode: z.string().optional(),
   customInstructions: z.string().max(2000).optional(),
   effort: effortLevelSchema.optional(),
+  jsonSchema: z.record(z.string(), z.unknown()).nullish(),
 });
 
 export type ReconnectSessionInput = z.infer<typeof reconnectSessionInput>;

apps/code/src/main/services/agent/service.ts

@@ -255,6 +255,8 @@ interface SessionConfig {
   effort?: EffortLevel;
   /** Model to use for the session (e.g. "claude-sonnet-4-6") */
   model?: string;
+  /** JSON Schema for structured task output — when set, the agent gets a create_output tool */
+  jsonSchema?: Record<string, unknown> | null;
 }
 
 interface ManagedSession {
@@ -561,6 +563,7 @@ When creating pull requests, add the following footer at the end of the PR descr
       customInstructions,
       effort,
       model,
+      jsonSchema,
     } = config;
 
     // Preview config doesn't need a real repo — use a temp directory
@@ -623,6 +626,14 @@ When creating pull requests, add the following footer at the end of the PR descr
         codexBinaryPath: adapter === "codex" ? getCodexBinaryPath() : undefined,
         model,
         instructions: adapter === "codex" ? systemPrompt.append : undefined,
+        onStructuredOutput: jsonSchema
+          ? async (output) => {
+              const posthogAPI = agent.getPosthogAPI();
+              if (posthogAPI) {
+                await posthogAPI.updateTaskRun(taskId, taskRunId, { output });
+              }
+            }
+          : undefined,
         processCallbacks: {
           onProcessSpawned: (info) => {
             this.processTracking.register(
@@ -758,6 +769,7 @@ When creating pull requests, add the following footer at the end of the PR descr
             systemPrompt,
             ...(permissionMode && { permissionMode }),
             ...(model != null && { model }),
+            ...(jsonSchema && { jsonSchema }),
             claudeCode: {
               options: claudeCodeOptions,
             },
@@ -780,6 +792,7 @@ When creating pull requests, add the following footer at the end of the PR descr
             systemPrompt,
             ...(permissionMode && { permissionMode }),
             ...(model != null && { model }),
+            ...(jsonSchema && { jsonSchema }),
             claudeCode: {
               options: claudeCodeOptions,
             },
@@ -1470,6 +1483,7 @@ For git operations while detached:
         "customInstructions" in params ? params.customInstructions : undefined,
       effort: "effort" in params ? params.effort : undefined,
       model: "model" in params ? params.model : undefined,
+      jsonSchema: "jsonSchema" in params ? params.jsonSchema : undefined,
     };
   }
 

packages/agent/package.json

@@ -96,6 +96,7 @@
   },
   "dependencies": {
     "@agentclientprotocol/sdk": "0.16.1",
+    "ajv": "^8.17.1",
     "@anthropic-ai/claude-agent-sdk": "0.2.76",
     "@anthropic-ai/sdk": "^0.78.0",
     "@hono/node-server": "^1.19.9",

packages/agent/src/adapters/acp-connection.ts

@@ -24,6 +24,8 @@ export type AcpConnectionConfig = {
   processCallbacks?: ProcessSpawnedCallback;
   codexOptions?: CodexProcessOptions;
   allowedModelIds?: Set<string>;
+  /** Callback invoked when the agent calls the create_output tool for structured output */
+  onStructuredOutput?: (output: Record<string, unknown>) => Promise<void>;
 };
 
 export type AcpConnection = {
@@ -97,7 +99,10 @@ function createClaudeConnection(config: AcpConnectionConfig): AcpConnection {
 
   let agent: ClaudeAcpAgent | null = null;
   const agentConnection = new AgentSideConnection((client) => {
-    agent = new ClaudeAcpAgent(client, config.processCallbacks);
+    agent = new ClaudeAcpAgent(client, {
+      ...config.processCallbacks,
+      onStructuredOutput: config.onStructuredOutput,
+    });
     logger.info(`Created ${agent.adapterName} agent`);
     return agent;
   }, agentStream);

packages/agent/src/adapters/claude/claude-agent.ts

@@ -109,6 +109,7 @@ export interface ClaudeAcpAgentOptions {
   onProcessSpawned?: (info: ProcessSpawnedInfo) => void;
   onProcessExited?: (pid: number) => void;
   onMcpServersReady?: (serverNames: string[]) => void;
+  onStructuredOutput?: (output: Record<string, unknown>) => Promise<void>;
 }
 
 export class ClaudeAcpAgent extends BaseAcpAgent {
@@ -483,6 +484,17 @@ export class ClaudeAcpAgent extends BaseAcpAgent {
             const result = handleResultMessage(message);
             if (result.error) throw result.error;
 
+            // Deliver structured output from SDK's native outputFormat
+            if (
+              message.subtype === "success" &&
+              message.structured_output != null &&
+              this.options?.onStructuredOutput
+            ) {
+              await this.options.onStructuredOutput(
+                message.structured_output as Record<string, unknown>,
+              );
+            }
+
             // For local-only commands, forward the result text to the client
             if (
               isLocalOnlyCommand &&
@@ -825,6 +837,12 @@ export class ClaudeAcpAgent extends BaseAcpAgent {
       : {};
     const systemPrompt = buildSystemPrompt(meta?.systemPrompt);
 
+    // Configure structured output via SDK's native outputFormat
+    const outputFormat =
+      meta?.jsonSchema && this.options?.onStructuredOutput
+        ? { type: "json_schema" as const, schema: meta.jsonSchema }
+        : undefined;
+
     this.logger.info(isResume ? "Resuming session" : "Creating new session", {
       sessionId,
       taskId,
@@ -854,6 +872,7 @@ export class ClaudeAcpAgent extends BaseAcpAgent {
         ...(meta?.additionalRoots ?? []),
       ],
       disableBuiltInTools: meta?.disableBuiltInTools,
+      outputFormat,
       settingsManager,
       onModeChange: this.createOnModeChange(),
       onProcessSpawned: this.options?.onProcessSpawned,

packages/agent/src/adapters/claude/session/options.ts

@@ -6,6 +6,7 @@ import type {
   CanUseTool,
   McpServerConfig,
   Options,
+  OutputFormat,
   SpawnedProcess,
   SpawnOptions,
 } from "@anthropic-ai/claude-agent-sdk";
@@ -42,6 +43,7 @@ export interface BuildOptionsParams {
   forkSession?: boolean;
   additionalDirectories?: string[];
   disableBuiltInTools?: boolean;
+  outputFormat?: OutputFormat;
   settingsManager: SettingsManager;
   onModeChange?: OnModeChange;
   onProcessSpawned?: (info: ProcessSpawnedInfo) => void;
@@ -268,6 +270,7 @@ export function buildSessionOptions(params: BuildOptionsParams): Options {
       params.settingsManager,
       params.logger,
     ),
+    outputFormat: params.outputFormat,
     abortController: getAbortController(
       params.userProvidedOptions?.abortController,
     ),

packages/agent/src/adapters/claude/types.ts

@@ -110,6 +110,7 @@ export type NewSessionMeta = {
   allowedDomains?: string[];
   /** Model ID to use for this session (e.g. "claude-sonnet-4-6") */
   model?: string;
+  jsonSchema?: Record<string, unknown> | null;
   claudeCode?: {
     options?: Options;
   };

packages/agent/src/agent.ts

@@ -122,6 +122,7 @@ export class Agent {
       deviceType: "local",
       logger: this.logger,
       processCallbacks: options.processCallbacks,
+      onStructuredOutput: options.onStructuredOutput,
       allowedModelIds,
       codexOptions:
         options.adapter === "codex" && gatewayConfig

packages/agent/src/posthog-api.ts

@@ -158,6 +158,20 @@ export class PostHogAPIClient {
     );
   }
 
+  async setTaskRunOutput(
+    taskId: string,
+    runId: string,
+    output: Record<string, unknown>,
+  ): Promise<TaskRun> {
+    return this.apiRequest(
+      `/api/projects/${this.getTeamId()}/tasks/${taskId}/runs/${runId}/set_output/`,
+      {
+        method: "PATCH",
+        body: JSON.stringify(output),
+      },
+    );
+  }
+
   async appendTaskRunLog(
     taskId: string,
     runId: string,

packages/agent/src/server/agent-server.ts

@@ -665,6 +665,15 @@ export class AgentServer {
       taskId: payload.task_id,
       deviceType: deviceInfo.type,
       logWriter,
+      onStructuredOutput: async (output) => {
+        await this.posthogAPI.setTaskRunOutput(
+          payload.task_id,
+          payload.run_id,
+          {
+            output,
+          },
+        );
+      },
     });
 
     // Tap both streams to broadcast all ACP messages via SSE (mimics local transport)
@@ -700,18 +709,25 @@ export class AgentServer {
       clientCapabilities: {},
     });
 
-    let preTaskRun: TaskRun | null = null;
-    try {
-      preTaskRun = await this.posthogAPI.getTaskRun(
-        payload.task_id,
-        payload.run_id,
-      );
-    } catch {
-      this.logger.warn("Failed to fetch task run for session context", {
-        taskId: payload.task_id,
-        runId: payload.run_id,
-      });
-    }
+    const [preTaskRun, preTask] = await Promise.all([
+      this.posthogAPI
+        .getTaskRun(payload.task_id, payload.run_id)
+        .catch((err) => {
+          this.logger.warn("Failed to fetch task run for session context", {
+            taskId: payload.task_id,
+            runId: payload.run_id,
+            error: err,
+          });
+          return null;
+        }),
+      this.posthogAPI.getTask(payload.task_id).catch((err) => {
+        this.logger.warn("Failed to fetch task for session context", {
+          taskId: payload.task_id,
+          error: err,
+        });
+        return null;
+      }),
+    ]);
 
     const prUrl =
       typeof (preTaskRun?.state as Record<string, unknown>)
@@ -732,6 +748,7 @@ export class AgentServer {
         taskRunId: payload.run_id,
         systemPrompt: this.buildSessionSystemPrompt(prUrl),
         allowedDomains: this.config.allowedDomains,
+        jsonSchema: preTask?.json_schema ?? null,
         ...(this.config.claudeCode?.plugins?.length && {
           claudeCode: {
             options: {