feat: Add initial structured output for agent

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>;
@@ -183,6 +184,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

@@ -196,6 +196,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 {
@@ -473,6 +475,7 @@ export class AgentService extends TypedEventEmitter<AgentServiceEvents> {
       customInstructions,
       effort,
       model,
+      jsonSchema,
     } = config;
 
     // Preview config doesn't need a real repo — use a temp directory
@@ -524,6 +527,14 @@ export class AgentService extends TypedEventEmitter<AgentServiceEvents> {
         adapter,
         gatewayUrl: proxyUrl,
         codexBinaryPath: adapter === "codex" ? getCodexBinaryPath() : undefined,
+        onStructuredOutput: jsonSchema
+          ? async (output) => {
+              const posthogAPI = agent.getPosthogAPI();
+              if (posthogAPI) {
+                await posthogAPI.updateTaskRun(taskId, taskRunId, { output });
+              }
+            }
+          : undefined,
         processCallbacks: {
           onProcessSpawned: (info) => {
             this.processTracking.register(
@@ -647,6 +658,7 @@ export class AgentService extends TypedEventEmitter<AgentServiceEvents> {
             systemPrompt,
             ...(permissionMode && { permissionMode }),
             ...(model != null && { model }),
+            ...(jsonSchema && { jsonSchema }),
             claudeCode: {
               options: {
                 ...(additionalDirectories?.length && {
@@ -679,6 +691,7 @@ export class AgentService extends TypedEventEmitter<AgentServiceEvents> {
             systemPrompt,
             ...(permissionMode && { permissionMode }),
             ...(model != null && { model }),
+            ...(jsonSchema && { jsonSchema }),
             claudeCode: {
               options: {
                 ...(additionalDirectories?.length && { additionalDirectories }),
@@ -1373,6 +1386,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

@@ -27,6 +27,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 = {
@@ -202,7 +204,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

@@ -107,6 +107,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 {
@@ -798,7 +799,44 @@ export class ClaudeAcpAgent extends BaseAcpAgent {
     await settingsManager.initialize();
 
     const mcpServers = parseMcpServers(params);
-    const systemPrompt = buildSystemPrompt(meta?.systemPrompt);
+    let systemPrompt = buildSystemPrompt(meta?.systemPrompt);
+
+    // Inject structured output tool if the task defines a JSON schema
+    if (meta?.jsonSchema && this.options?.onStructuredOutput) {
+      const { createOutputMcpServer, OUTPUT_SERVER_NAME } = await import(
+        "./structured-output/create-output-server"
+      );
+      mcpServers[OUTPUT_SERVER_NAME] = createOutputMcpServer({
+        jsonSchema: meta.jsonSchema,
+        onOutput: this.options.onStructuredOutput,
+        logger: this.logger,
+      });
+
+      const schemaStr = JSON.stringify(meta.jsonSchema, null, 2);
+      const outputInstruction =
+        "\n\n# Structured Output\n\n" +
+        "This task requires structured output. You MUST use the `create_output` tool " +
+        "(available as `mcp__posthog_output__create_output`) to deliver your final result " +
+        "before ending the task. The output must conform to the following JSON Schema:\n\n" +
+        `\`\`\`json\n${schemaStr}\n\`\`\`\n\n` +
+        "Call the create_output tool with the required fields as arguments once you have " +
+        "gathered all necessary information. Do not end the task without calling create_output.";
+
+      if (typeof systemPrompt === "string") {
+        systemPrompt = systemPrompt + outputInstruction;
+      } else if (
+        systemPrompt &&
+        typeof systemPrompt === "object" &&
+        "append" in systemPrompt
+      ) {
+        systemPrompt = {
+          ...systemPrompt,
+          append:
+            ((systemPrompt as { append?: string }).append ?? "") +
+            outputInstruction,
+        };
+      }
+    }
 
     this.logger.info(isResume ? "Resuming session" : "Creating new session", {
       sessionId,

packages/agent/src/adapters/claude/structured-output/constants.ts

@@ -0,0 +1,3 @@
+export const OUTPUT_SERVER_NAME = "posthog_output";
+export const OUTPUT_TOOL_NAME = "create_output";
+export const OUTPUT_TOOL_FULL_NAME = `mcp__${OUTPUT_SERVER_NAME}__${OUTPUT_TOOL_NAME}`;

packages/agent/src/adapters/claude/structured-output/create-output-server.ts

@@ -0,0 +1,89 @@
+import {
+  createSdkMcpServer,
+  type McpSdkServerConfigWithInstance,
+  tool,
+} from "@anthropic-ai/claude-agent-sdk";
+import Ajv from "ajv";
+import * as z from "zod";
+import type { Logger } from "../../../utils/logger";
+import { OUTPUT_SERVER_NAME, OUTPUT_TOOL_NAME } from "./constants";
+
+export {
+  OUTPUT_SERVER_NAME,
+  OUTPUT_TOOL_FULL_NAME,
+  OUTPUT_TOOL_NAME,
+} from "./constants";
+
+export interface CreateOutputServerOptions {
+  jsonSchema: Record<string, unknown>;
+  onOutput: (output: Record<string, unknown>) => Promise<void>;
+  logger: Logger;
+}
+
+export function createOutputMcpServer(
+  options: CreateOutputServerOptions,
+): McpSdkServerConfigWithInstance {
+  const { jsonSchema, onOutput, logger } = options;
+
+  const ajv = new Ajv({ allErrors: true });
+  const validate = ajv.compile(jsonSchema);
+  const zodType: z.ZodType = z.fromJSONSchema(jsonSchema); // Validate that the JSON schema can be converted to Zod schema, will throw if invalid
+  if (!(zodType instanceof z.ZodObject)) {
+    throw new Error(
+      "Only JSON schemas that correspond to Zod objects are supported",
+    );
+  }
+  const outputTool = tool(
+    OUTPUT_TOOL_NAME,
+    "Submit the structured output for this task. Call this tool with the required fields to deliver your final result. The output must conform to the task's JSON schema.",
+    zodType.shape,
+    async (args) => {
+      const valid = validate(args);
+      if (!valid) {
+        const errors = validate.errors
+          ?.map((e) => `${e.instancePath || "/"}: ${e.message}`)
+          .join("; ");
+        logger.warn("Structured output validation failed", { errors });
+        return {
+          content: [
+            {
+              type: "text" as const,
+              text: `Validation failed: ${errors}. Please fix the output and try again.`,
+            },
+          ],
+          isError: true,
+        };
+      }
+
+      try {
+        await onOutput(args as Record<string, unknown>);
+        logger.info("Structured output persisted successfully");
+        return {
+          content: [
+            {
+              type: "text" as const,
+              text: "Output submitted successfully.",
+            },
+          ],
+        };
+      } catch (error) {
+        const message = error instanceof Error ? error.message : String(error);
+        logger.error("Failed to persist structured output", { error: message });
+        return {
+          content: [
+            {
+              type: "text" as const,
+              text: `Failed to submit output: ${message}`,
+            },
+          ],
+          isError: true,
+        };
+      }
+    },
+  );
+
+  return createSdkMcpServer({
+    name: OUTPUT_SERVER_NAME,
+    tools: [outputTool],
+  });
+}

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

@@ -7,6 +7,7 @@ export {
 
 import type { CodeExecutionMode } from "../../execution-mode";
 import { isMcpToolReadOnly } from "./mcp/tool-metadata";
+import { OUTPUT_TOOL_FULL_NAME } from "./structured-output/constants";
 
 export const READ_TOOLS: Set<string> = new Set(["Read", "NotebookRead"]);
 
@@ -38,6 +39,7 @@ const BASE_ALLOWED_TOOLS = [
   ...SEARCH_TOOLS,
   ...WEB_TOOLS,
   ...AGENT_TOOLS,
+  OUTPUT_TOOL_FULL_NAME,
 ];
 
 const AUTO_ALLOWED_TOOLS: Record<string, Set<string>> = {

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

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

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

@@ -650,6 +650,11 @@ export class AgentServer {
       taskId: payload.task_id,
       deviceType: deviceInfo.type,
       logWriter,
+      onStructuredOutput: async (output) => {
+        await this.posthogAPI.updateTaskRun(payload.task_id, payload.run_id, {
+          output,
+        });
+      },
     });
 
     // Tap both streams to broadcast all ACP messages via SSE (mimics local transport)
@@ -685,18 +690,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>)
@@ -717,6 +729,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: {

packages/agent/src/types.ts

@@ -113,6 +113,8 @@ export interface TaskExecutionOptions {
   gatewayUrl?: string;
   codexBinaryPath?: string;
   processCallbacks?: ProcessSpawnedCallback;
+  /** Callback invoked when the agent calls the create_output tool for structured output */
+  onStructuredOutput?: (output: Record<string, unknown>) => Promise<void>;
 }
 
 export type LogLevel = "debug" | "info" | "warn" | "error";