diff --git a/apps/code/src/main/services/agent/schemas.ts b/apps/code/src/main/services/agent/schemas.ts index 5ad875f8f..3ead6cf15 100644 --- a/apps/code/src/main/services/agent/schemas.ts +++ b/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; @@ -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; diff --git a/apps/code/src/main/services/agent/service.ts b/apps/code/src/main/services/agent/service.ts index 839fa2a75..72a4c318d 100644 --- a/apps/code/src/main/services/agent/service.ts +++ b/apps/code/src/main/services/agent/service.ts @@ -197,6 +197,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 | null; } interface ManagedSession { @@ -503,6 +505,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 @@ -554,6 +557,14 @@ When creating pull requests, add the following footer at the end of the PR descr 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( @@ -678,6 +689,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: { ...(additionalDirectories?.length && { @@ -711,6 +723,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: { ...(additionalDirectories?.length && { additionalDirectories }), @@ -1405,6 +1418,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, }; } diff --git a/packages/agent/package.json b/packages/agent/package.json index 07ef18fbf..17a9a3794 100644 --- a/packages/agent/package.json +++ b/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", diff --git a/packages/agent/src/adapters/acp-connection.ts b/packages/agent/src/adapters/acp-connection.ts index 9567b033f..11662129f 100644 --- a/packages/agent/src/adapters/acp-connection.ts +++ b/packages/agent/src/adapters/acp-connection.ts @@ -27,6 +27,8 @@ export type AcpConnectionConfig = { processCallbacks?: ProcessSpawnedCallback; codexOptions?: CodexProcessOptions; allowedModelIds?: Set; + /** Callback invoked when the agent calls the create_output tool for structured output */ + onStructuredOutput?: (output: Record) => Promise; }; 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); diff --git a/packages/agent/src/adapters/claude/claude-agent.ts b/packages/agent/src/adapters/claude/claude-agent.ts index 2aafaaac4..84618e4c6 100644 --- a/packages/agent/src/adapters/claude/claude-agent.ts +++ b/packages/agent/src/adapters/claude/claude-agent.ts @@ -108,6 +108,7 @@ export interface ClaudeAcpAgentOptions { onProcessSpawned?: (info: ProcessSpawnedInfo) => void; onProcessExited?: (pid: number) => void; onMcpServersReady?: (serverNames: string[]) => void; + onStructuredOutput?: (output: Record) => Promise; } export class ClaudeAcpAgent extends BaseAcpAgent { @@ -803,7 +804,44 @@ export class ClaudeAcpAgent extends BaseAcpAgent { const mcpServers = supportsMcpInjection(earlyModelId) ? 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, diff --git a/packages/agent/src/adapters/claude/structured-output/constants.ts b/packages/agent/src/adapters/claude/structured-output/constants.ts new file mode 100644 index 000000000..22656ea16 --- /dev/null +++ b/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}`; diff --git a/packages/agent/src/adapters/claude/structured-output/create-output-server.ts b/packages/agent/src/adapters/claude/structured-output/create-output-server.ts new file mode 100644 index 000000000..fc31d2136 --- /dev/null +++ b/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; + onOutput: (output: Record) => Promise; + 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); + 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], + }); +} diff --git a/packages/agent/src/adapters/claude/tools.ts b/packages/agent/src/adapters/claude/tools.ts index 1f2621aff..9e696dc4d 100644 --- a/packages/agent/src/adapters/claude/tools.ts +++ b/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 = 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> = { diff --git a/packages/agent/src/adapters/claude/types.ts b/packages/agent/src/adapters/claude/types.ts index 8bc8cce35..31e18d235 100644 --- a/packages/agent/src/adapters/claude/types.ts +++ b/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 | null; claudeCode?: { options?: Options; }; diff --git a/packages/agent/src/agent.ts b/packages/agent/src/agent.ts index 6c0081cca..e26c061af 100644 --- a/packages/agent/src/agent.ts +++ b/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 diff --git a/packages/agent/src/server/agent-server.ts b/packages/agent/src/server/agent-server.ts index 8f8fa25fe..c16e4884d 100644 --- a/packages/agent/src/server/agent-server.ts +++ b/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) @@ -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: { diff --git a/packages/agent/src/types.ts b/packages/agent/src/types.ts index b463189e7..7fe150dce 100644 --- a/packages/agent/src/types.ts +++ b/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) => Promise; } export type LogLevel = "debug" | "info" | "warn" | "error"; diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 85471188f..4e3648986 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -621,6 +621,9 @@ importers: '@types/jsonwebtoken': specifier: ^9.0.10 version: 9.0.10 + ajv: + specifier: ^8.17.1 + version: 8.17.1 commander: specifier: ^14.0.2 version: 14.0.3