fix: add validation for outputSchema in StructuredGenerationTask, changes to tool use schema

Enhances the StructuredGenerationTask by adding a check to ensure that outputSchema is a valid JSON Schema. Updates test cases to reflect changes in input handling for better type safety.

Author: sroussey

packages/ai/src/task/ChatMessage.ts

@@ -28,10 +28,22 @@ export type ContentBlockToolUse = {
   readonly input: Record<string, unknown>;
 };
 
+/**
+ * Blocks that may appear in a `tool_result`'s `content` array. Provider payloads
+ * typically use text, image, and tool_use; nested `tool_result` is not modeled
+ * here so the JSON schema can be embedded in parent task schemas without a
+ * recursive `$ref` (which fails to resolve when `ContentBlockSchema` is nested
+ * under a larger document such as `ToolCallingInputSchema`).
+ */
+export type ContentBlockInToolResultBody =
+  | ContentBlockText
+  | ContentBlockImage
+  | ContentBlockToolUse;
+
 export type ContentBlockToolResult = {
   readonly type: "tool_result";
   readonly tool_use_id: string;
-  readonly content: ReadonlyArray<ContentBlock>;
+  readonly content: ReadonlyArray<ContentBlockInToolResultBody>;
   readonly is_error: boolean | undefined;
 };
 
@@ -89,16 +101,19 @@ const ContentBlockToolUseSchema = {
   additionalProperties: false,
 } as const;
 
-// tool_result is recursive — its `content` is an array of ContentBlock.
-// The $ref resolves against ContentBlockSchema.definitions.ContentBlock below.
+/** `tool_result.content` — text, image, and tool_use only (no nested `tool_result`). */
+const ContentBlockInToolResultBodySchema = {
+  oneOf: [ContentBlockTextSchema, ContentBlockImageSchema, ContentBlockToolUseSchema],
+} as const;
+
 const ContentBlockToolResultSchema = {
   type: "object",
   properties: {
     type: { type: "string", enum: ["tool_result"] },
     tool_use_id: { type: "string" },
     content: {
       type: "array",
-      items: { $ref: "#/definitions/ContentBlock" },
+      items: ContentBlockInToolResultBodySchema,
     },
     is_error: { type: "boolean" },
   },
@@ -115,16 +130,6 @@ export const ContentBlockSchema = {
     ContentBlockToolUseSchema,
     ContentBlockToolResultSchema,
   ],
-  definitions: {
-    ContentBlock: {
-      oneOf: [
-        ContentBlockTextSchema,
-        ContentBlockImageSchema,
-        ContentBlockToolUseSchema,
-        ContentBlockToolResultSchema,
-      ],
-    },
-  },
   title: "ContentBlock",
   description: "A single content block within a chat message",
 } as const;
@@ -148,6 +153,28 @@ export const ChatMessageSchema = {
 // Runtime type guards
 // ========================================================================
 
+export function isContentBlockInToolResultBody(
+  value: unknown
+): value is ContentBlockInToolResultBody {
+  if (!value || typeof value !== "object") return false;
+  const v = value as Record<string, unknown>;
+  switch (v.type) {
+    case "text":
+      return typeof v.text === "string";
+    case "image":
+      return typeof v.mimeType === "string" && typeof v.data === "string";
+    case "tool_use":
+      return (
+        typeof v.id === "string" &&
+        typeof v.name === "string" &&
+        v.input !== null &&
+        typeof v.input === "object"
+      );
+    default:
+      return false;
+  }
+}
+
 export function isContentBlock(value: unknown): value is ContentBlock {
   if (!value || typeof value !== "object") return false;
   const v = value as Record<string, unknown>;
@@ -167,7 +194,7 @@ export function isContentBlock(value: unknown): value is ContentBlock {
       return (
         typeof v.tool_use_id === "string" &&
         Array.isArray(v.content) &&
-        v.content.every(isContentBlock)
+        v.content.every(isContentBlockInToolResultBody)
       );
     default:
       return false;

packages/ai/src/task/StructuredGenerationTask.ts

@@ -4,10 +4,10 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { CreateWorkflow, TaskConfigurationError, TaskError, Workflow } from "@workglow/task-graph";
 import type { IExecuteContext, StreamEvent, TaskConfig } from "@workglow/task-graph";
-import { compileSchema } from "@workglow/util/schema";
+import { CreateWorkflow, TaskConfigurationError, TaskError, Workflow } from "@workglow/task-graph";
 import type { DataPortSchema, FromSchema, SchemaNode } from "@workglow/util/schema";
+import { compileSchema } from "@workglow/util/schema";
 import { TypeModel } from "./base/AiTaskSchemas";
 import { StreamingAiTask } from "./base/StreamingAiTask";
 
@@ -163,6 +163,7 @@ export class StructuredGenerationTask extends StreamingAiTask<
     let validator: SchemaNode;
     try {
       validator = compileSchema(input.outputSchema);
+      if (!input.outputSchema) throw new Error("outputSchema is not a valid JSON Schema");
     } catch (err) {
       const msg = err instanceof Error ? err.message : String(err);
       const configErr = new TaskConfigurationError(

packages/test/src/test/task/StructuredGenerationTask.test.ts

@@ -323,10 +323,10 @@ describe("StructuredGenerationTask — schema compile errors", () => {
       outputSchema: null as unknown as Record<string, unknown>,
       maxRetries: 0,
     };
-    const task = new StructuredGenerationTask({ defaults: input } as any);
+    const task = new StructuredGenerationTask({ defaults: input });
     let caught: unknown;
     try {
-      await drain(task.executeStream(input as any, mkContext()));
+      await drain(task.executeStream(input, mkContext()));
     } catch (e) {
       caught = e;
     }