Implement support for output schemas (#411)

* tweak examples * tweak prompt * Implement support for output schemas * fix: add example for structured output with zod schema * tweak * remove redundant cast * fix input name * strip $schema * hack around vendor requirement * clarify required result output --------- Co-authored-by: pullfrog[bot] <226033991+pullfrog[bot]@users.noreply.github.com>
2026-03-05 23:07:03 +00:00
parent fafe930c77
commit 5684cbef77
11 changed files with 5981 additions and 89 deletions
@@ -1,3 +1,5 @@
+import type { StandardSchemaV1 } from "@standard-schema/spec";
+import { Ajv } from "ajv";
 import { type } from "arktype";
 import { log } from "../utils/cli.ts";
 import type { ToolContext } from "./server.ts";
@@ -7,29 +9,76 @@ export const SetOutputParams = type({
  value: type.string.describe("the output value to expose as a GitHub Action output"),
 });

-export function SetOutputTool(ctx: ToolContext) {
+type JsonSchema = Record<string, unknown>;
+
+function jsonSchemaToStandardSchema({
+  $schema: _,
+  ...jsonSchema
+}: JsonSchema): StandardSchemaV1<any> & { toJsonSchema(): JsonSchema } {
+  const ajv = new Ajv();
+  const validate = ajv.compile(jsonSchema);
+
+  return {
+    "~standard": {
+      version: 1,
+      // xsschema (used by fastmcp) hardcodes supported vendors instead of duck-typing toJsonSchema().
+      // "arktype" works because its converter is just `(schema) => schema.toJsonSchema()`.
+      vendor: "arktype",
+      validate(input: unknown) {
+        if (validate(input)) {
+          return { value: input };
+        }
+        return {
+          issues: (validate.errors ?? []).map((err) => ({
+            message: `${err.instancePath || "/"}: ${err.message ?? "validation error"}`,
+            path: err.instancePath ? err.instancePath.split("/").filter(Boolean) : [],
+          })),
+        };
+      },
+    },
+    toJsonSchema() {
+      return jsonSchema;
+    },
+  };
+}
+
+function storeOutput(ctx: ToolContext, value: string) {
+  const selfId = ctx.toolState.selfSubagentId;
+  if (selfId) {
+    const subagent = ctx.toolState.subagents.get(selfId);
+    if (subagent) {
+      subagent.output = value;
+      log.debug(`set_output: routed to subagent ${selfId} (value=${value.slice(0, 80)})`);
+      return { success: true, routed: "subagent" as const };
+    }
+    log.warning(
+      `set_output: selfSubagentId=${selfId} but subagent not found in map — routing to action output`
+    );
+  }
+  ctx.toolState.output = value;
+  return { success: true, routed: "action_output" as const };
+}
+
+export function SetOutputTool(ctx: ToolContext, outputSchema?: JsonSchema) {
+  if (outputSchema) {
+    return tool({
+      name: "set_output",
+      description:
+        "Set the structured action output. You MUST call this tool before finishing — the output is required. Pass the output object directly as the tool arguments (no wrapping needed).",
+      parameters: jsonSchemaToStandardSchema(outputSchema),
+      execute: execute(async (params) => {
+        return storeOutput(ctx, JSON.stringify(params));
+      }),
+    });
+  }
+
  return tool({
    name: "set_output",
    description:
      "Set the action output. When called by a subagent, returns a summary result to the orchestrator — this is the ONLY way to pass results back. When called by the orchestrator in standalone mode (trigger: unknown), exposes the value as the 'result' GitHub Action output for downstream workflow steps. Do NOT use this for progress reporting — use report_progress instead.",
    parameters: SetOutputParams,
    execute: execute(async (params) => {
-      const selfId = ctx.toolState.selfSubagentId;
-      if (selfId) {
-        const subagent = ctx.toolState.subagents.get(selfId);
-        if (subagent) {
-          subagent.output = params.value;
-          log.debug(
-            `set_output: routed to subagent ${selfId} (value=${params.value.slice(0, 80)})`
-          );
-          return { success: true, routed: "subagent" };
-        }
-        log.warning(
-          `set_output: selfSubagentId=${selfId} but subagent not found in map — routing to action output`
-        );
-      }
-      ctx.toolState.output = params.value;
-      return { success: true, routed: "action_output" };
+      return storeOutput(ctx, params.value);
    }),
  });
 }
@@ -189,8 +189,10 @@ function isAddressInUse(error: unknown): boolean {
  return message.includes("eaddrinuse") || message.includes("address already in use");
 }

+type JsonSchema = Record<string, unknown>;
+
 // tools shared by both orchestrator and subagent servers
-function buildCommonTools(ctx: ToolContext): Tool<any, any>[] {
+function buildCommonTools(ctx: ToolContext, outputSchema?: JsonSchema): Tool<any, any>[] {
  const tools: Tool<any, any>[] = [
    StartDependencyInstallationTool(ctx),
    AwaitDependencyInstallationTool(ctx),
@@ -213,7 +215,7 @@ function buildCommonTools(ctx: ToolContext): Tool<any, any>[] {
    GitTool(ctx),
    GitFetchTool(ctx),
    UploadFileTool(ctx),
-    SetOutputTool(ctx),
+    SetOutputTool(ctx, outputSchema),
    FileReadTool(ctx),
    FileWriteTool(ctx),
    FileEditTool(ctx),
@@ -234,9 +236,9 @@ function buildCommonTools(ctx: ToolContext): Tool<any, any>[] {
 }

 // orchestrator gets common tools + delegation + remote-mutating tools
-function buildOrchestratorTools(ctx: ToolContext): Tool<any, any>[] {
+function buildOrchestratorTools(ctx: ToolContext, outputSchema?: JsonSchema): Tool<any, any>[] {
  return [
-    ...buildCommonTools(ctx),
+    ...buildCommonTools(ctx, outputSchema),
    ReportProgressTool(ctx),
    SelectModeTool(ctx),
    DelegateTool(ctx),
@@ -353,13 +355,18 @@ async function killBackgroundProcesses(toolState: ToolState): Promise<void> {
  backgroundProcesses.clear();
 }

+type McpHttpServerOptions = {
+  outputSchema?: JsonSchema | undefined;
+};
+
 /**
 * Start the orchestrator MCP HTTP server (has all tools including push/PR/delegation).
 */
 export async function startMcpHttpServer(
-  ctx: ToolContext
+  ctx: ToolContext,
+  options?: McpHttpServerOptions
 ): Promise<{ url: string; [Symbol.asyncDispose]: () => Promise<void> }> {
-  const tools = buildOrchestratorTools(ctx);
+  const tools = buildOrchestratorTools(ctx, options?.outputSchema);
  const startResult = await selectMcpPort(ctx, tools);

  return {