feat: deepen Zod integration in Ax

docs/ZOD_INTEGRATION.md

@@ -0,0 +1,187 @@
+# Zod Integration Blueprint for Ax
+
+## Overview
+
+This document proposes a deep, foundational integration of [Zod](https://zod.dev) into Ax so that Zod schemas become first-class citizens throughout the signature, engine, and assertion layers. The design builds upon the groundwork in PR #388 and aims to unlock seamless schema reuse, rich runtime validation, and best-in-class developer ergonomics for AI workflow authoring.
+
+---
+
+## Architecture Proposal
+
+The integration introduces new adapters and validation flows that allow Zod schemas to travel with Ax signatures from definition to runtime enforcement.
+
+### Component Diagram (textual UML)
+
+```
++-----------------+                 +--------------------+             +---------------------+
+|  Zod Schema     | --(register)--> | AxZodRegistry      | --(build)-> | AxSignatureFactory  |
+| (z.object(...)) |                 | (schema cache)     |             | (existing)          |
++-----------------+                 +--------------------+             +---------------------+
+         |                                      |                                  |
+         |                                      v                                  v
+         |                            +----------------+                 +---------------------+
+         |                            | AxSignature    |<--+             | AxGen / AxFlow      |
+         |                            | (with zodMeta) |   |             | (DSPy-style engine) |
+         |                            +----------------+   |             +---------------------+
+         |                                      |          |                      |
+         |                                      v          |                      v
+         |                            +----------------+   |             +----------------------+
+         |                            | AxAssertion    |<--+--(auto)-----| ZodAssertionAdapter  |
+         |                            | Pipeline       |                 | (.parse/.safeParse) |
+         |                            +----------------+                 +----------------------+
+         |                                      |                                  |
+         |                                      v                                  v
+         |                                +----------------+           +-------------------------+
+         +------------------------------+ | Runtime Output | ---LLM--> | Streaming Validator     |
+                                          | (JSON/prose)   |           | (per-field / final)     |
+                                          +----------------+           +-------------------------+
+                                                      |
+                                                      v
+                                            +-------------------+
+                                            | ValidationResult  |
+                                            |  - success        |
+                                            |  - errors         |
+                                            |  - telemetry      |
+                                            +-------------------+
+
+```
+
+### Key Modules & APIs
+
+- `AxSignature.fromZod(schema, options?: AxZodSignatureOptions): AxSignature`
+
+  - `options.strict`: throw on any downgrade or unsupported feature.
+  - `options.streaming`: enable emission of field-level validators for streaming output enforcement.
+  - `options.mode`: `'parse' | 'safeParse' | 'coerce'` to align with Zod parsing semantics.
+  - `options.assertionLevel`: `'none' | 'final' | 'streaming' | 'both'` to control auto-assertion wiring.
+
+- `AxSignature.toZod(signature, options?: AxToZodOptions): ZodSchema`
+
+  - Round-trip support with metadata preservation.
+
+- `AxZodRegistry`
+
+  - Internal cache keyed by signature ID to store original Zod schema, downgrade notes, and validation options.
+  - Provides `get(schemaId)` for runtime validation modules.
+
+- `ZodAssertionAdapter`
+
+  - Translates Zod `.parse`, `.safeParse`, `.min`, `.max`, `.default`, `.catch`, `.transform`, and `.refine` hooks into Ax assertions.
+  - Emits `AxAssertion` objects with `severity`, `recovery` (fallback result), and `telemetry` payloads.
+
+- `StreamingZodValidator`
+
+  - Wraps Zod schema introspection to produce field-level validators suitable for Ax’s streaming extraction pipeline.
+  - Supports chunk-level validation and progressive parse with buffering.
+
+- `AxValidationTelemetry`
+
+  - Unified event schema capturing downgrade issues, parse failures, defaults applied, and user-facing remediation tips.
+
+- `AxZodCLI`
+  - CLI command (`npx ax zod audit`) to inspect schemas, report downgrades, and generate migration hints.
+
+---
+
+## Code Snippets
+
+### 1. Converting Zod Schema to Ax Signature with Strict Validation
+
+```ts
+import { z } from 'zod'
+import { AxSignature } from 'ax/signature'
+
+const invoiceSchema = z.object({
+  invoiceId: z.string().uuid(),
+  totalCents: z.number().int().min(0),
+  issuedAt: z
+    .string()
+    .datetime()
+    .default(() => new Date().toISOString()),
+})
+
+const signature = AxSignature.fromZod(invoiceSchema, {
+  strict: true,
+  assertionLevel: 'final',
+  mode: 'safeParse',
+})
+```
+
+### 2. Using Auto-Applied Zod Assertions in AxGen
+
+```ts
+import { axgen } from 'ax'
+import { z } from 'zod'
+
+const schema = z.object({
+  customerName: z.string().min(1),
+  preferredContact: z.enum(['email', 'phone']).default('email'),
+})
+
+const gen = axgen()
+  .signature(schema) // equivalent to AxSignature.fromZod(schema)
+  .prompt('Collect the customer preferences from the conversation.')
+  .onFailure((ctx, error) => {
+    ctx.logger.warn('Zod validation failed', { error })
+    return ctx.retry()
+  })
+
+const result = await gen.run({ transcript })
+// result.payload is guaranteed to satisfy schema.parse(...) semantics
+```
+
+### 3. Streaming Validation Hook
+
+```ts
+import { createStreamingValidator } from 'ax/zod/stream'
+
+const validator = createStreamingValidator(schema, { chunkSize: 128 })
+
+for await (const token of llmStream) {
+  const status = validator.ingest(token)
+  if (status.type === 'error') {
+    // Apply jsonrepair, request model correction, or abort
+    await controller.requestFix(status.issues)
+  }
+}
+
+const finalValue = validator.finalize()
+```
+
+### 4. Round-trip Conversion (AxSignature → Zod)
+
+```ts
+const existingSignature = AxSignature.load('customer.profile')
+const zodSchema = existingSignature.toZod()
+
+// Apply additional refinements with Zod API
+const stricterSchema = zodSchema.extend({
+  loyaltyTier: z.enum(['bronze', 'silver', 'gold']).catch('bronze'),
+})
+```
+
+### 5. CLI Audit Example
+
+```bash
+npx ax zod audit ./schemas/customer.ts --strict --report json
+```
+
+---
+
+## Trade-offs & Risks
+
+- **Performance Overhead**: Zod parsing incurs runtime cost, especially for large schemas. Mitigation: allow opt-in streaming validation to catch errors early, cache compiled schemas, and benchmark to ensure <10% throughput regression for typical flows.
+- **Streaming Complexity**: Zod is not inherently streaming-aware. The proposed streaming validator will require custom buffering logic; certain constructs (e.g., regex on entire strings, cross-field refinement) may only be enforceable at finalize-time.
+- **Schema Fidelity Gaps**: Features like `ZodFunction`, complex intersections, and advanced `ZodEffects` may still need JSON fallbacks. Document these limitations and provide telemetry to guide users.
+- **Version Compatibility**: Supporting both Zod v3 and v4 demands careful dependency and type management. Use adapters and peer dependency ranges; run compatibility tests in CI.
+- **Optional Dependency**: Keeping Zod optional reduces bundle size for non-users but requires defensive imports and runtime checks. Provide clear error messages when Zod-specific APIs are invoked without the dependency.
+
+---
+
+## Success Metrics
+
+- **Schema Fidelity**: ≥95% of schemas in the test corpus convert round-trip without downgrade warnings.
+- **Runtime Reliability**: ≥90% reduction in manual signature authoring among beta users migrating from Zod-heavy stacks (survey-based).
+- **DX Satisfaction**: Positive qualitative feedback (≥4/5) from developer surveys on the new CLI, docs, and auto-assertions.
+- **Telemetry Signals**: Monitoring shows decreasing parse failure rates over time due to better defaults and `.catch` handling; streaming validator issues are actionable.
+- **Adoption**: At least two reference integrations (e.g., Mastra migration, OpenAI structured extraction recipe) published using the new APIs.

src/ax/dsp/generate.test.ts

@@ -1,13 +1,17 @@
+// cspell:ignore summarising neutralpositive
 import { ReadableStream } from 'node:stream/web';
 
 import { describe, expect, it } from 'vitest';
+import { z } from 'zod';
 
 import { validateAxMessageArray } from '../ai/base.js';
 import { AxMockAIService } from '../ai/mock/api.js';
 import type { AxChatResponse } from '../ai/types.js';
 import { AxStopFunctionCallException } from './functions.js';
 import { AxGen, type AxGenerateError } from './generate.js';
 import { AxSignature } from './sig.js';
+import { getZodMetadata } from '../zod/metadata.js';
+import { ax } from './template.js';
 import type { AxProgramForwardOptions } from './types.js';
 
 function createStreamingResponse(
@@ -729,6 +733,183 @@ describe('AxGen Message Validation', () => {
   });
 });
 
+describe('AxGen with Zod signatures', () => {
+  it('constructs directly from a Zod schema with custom options', () => {
+    const schema = z.object({
+      title: z.string().min(3),
+      tags: z.array(z.string()).default([]),
+    });
+
+    const gen = new AxGen(schema, {
+      zod: {
+        assertionLevel: 'none',
+        mode: 'parse',
+      },
+    });
+
+    const signature = gen.getSignature();
+    const metadata = getZodMetadata(signature);
+    expect(metadata?.options.mode).toBe('parse');
+    expect(metadata?.options.assertionLevel).toBe('none');
+    const asserts = (gen as any).asserts as unknown[];
+    expect(asserts?.length ?? 0).toBe(0);
+  });
+
+  it('creates a generator via ax() helper when given a Zod schema', () => {
+    const schema = z.object({
+      topic: z.string(),
+      summary: z.string().catch('summary unavailable'),
+    });
+
+    const gen = ax(schema, {
+      zod: {
+        assertionLevel: 'final',
+        mode: 'safeParse',
+      },
+    });
+
+    expect(gen).toBeInstanceOf(AxGen);
+    const asserts = (gen as any).asserts as unknown[];
+    expect(asserts.length).toBeGreaterThan(0);
+    const metadata = getZodMetadata(gen.getSignature());
+    expect(metadata?.schema).toBe(schema);
+  });
+
+  it('applies final Zod assertions and default values', async () => {
+    const schema = z.object({
+      name: z.string().min(1),
+      count: z.number().int().default(1),
+    });
+
+    const signature = AxSignature.fromZod(schema);
+    const gen = new AxGen(signature);
+    const asserts = (gen as any).asserts as Array<{
+      fn: (values: Record<string, unknown>) => Promise<unknown>;
+    }>;
+
+    expect(asserts.length).toBeGreaterThan(0);
+    const zodAssert = asserts[asserts.length - 1];
+
+    const values: Record<string, unknown> = { name: 'Ada' };
+    await expect(zodAssert.fn(values)).resolves.toBe(true);
+    expect(values.count).toBe(1);
+
+    const failure = await zodAssert.fn({ name: '' });
+    expect(typeof failure).toBe('string');
+    expect(failure).toContain('Zod validation failed');
+  });
+
+  it('honors catch, default, min, and max constraints', async () => {
+    const schema = z.object({
+      status: z.enum(['ok', 'warn', 'error']).catch('error'),
+      attempts: z.number().int().min(1).max(5).default(3),
+    });
+
+    const signature = AxSignature.fromZod(schema);
+    const gen = new AxGen(signature);
+    const asserts = (gen as any).asserts as Array<{
+      fn: (values: Record<string, unknown>) => Promise<unknown>;
+    }>;
+
+    const zodAssert = asserts[asserts.length - 1];
+
+    const withFallback: Record<string, unknown> = {
+      status: 'unexpected',
+    };
+    await expect(zodAssert.fn(withFallback)).resolves.toBe(true);
+    expect(withFallback.status).toBe('error');
+    expect(withFallback.attempts).toBe(3);
+
+    await expect(zodAssert.fn({ status: 'ok', attempts: 0 })).resolves.toMatch(
+      /greater than or equal to 1/i
+    );
+
+    await expect(zodAssert.fn({ status: 'ok', attempts: 10 })).resolves.toMatch(
+      /less than or equal to 5/i
+    );
+  });
+
+  it('emits repaired values after streaming assertions mutate the payload', async () => {
+    const schema = z.object({
+      summary: z
+        .string()
+        .min(
+          40,
+          'Provide a short paragraph summarising the topic in at least 40 characters.'
+        )
+        .catch(
+          'Summary unavailable from the model output — manual review recommended to craft a compliant paragraph.'
+        ),
+      sentiment: z.enum(['positive', 'neutral', 'negative']).catch('neutral'),
+      confidence: z
+        .number()
+        .min(0, 'Confidence must be between 0 and 1.')
+        .max(1, 'Confidence must be between 0 and 1.')
+        .catch(0.65),
+      highlights: z
+        .array(z.string().min(10, 'Each highlight should be a full thought.'))
+        .min(2, 'Provide at least two highlights.')
+        .max(4, 'Keep the highlight list focused.')
+        .catch([
+          'The model response did not supply valid highlights. Highlight that human review is required.',
+          'Flag the output for review because automatic repair was triggered.',
+        ]),
+      actionPlan: z
+        .string()
+        .min(10, 'Provide concrete guidance.')
+        .default('No immediate action required.'),
+    });
+
+    const signature = AxSignature.fromZod(schema, { mode: 'safeParse' });
+    const gen = new AxGen(signature);
+
+    const malformedPayload = [
+      'Summary: Too short.',
+      'Sentiment: neutralpositive',
+      'Confidence: 2.0',
+      'Highlights:',
+      '- tiny',
+      '- also tiny',
+    ].join('\n');
+
+    const chunks: AxChatResponse['results'] = [
+      {
+        index: 0,
+        content: malformedPayload,
+      },
+      {
+        index: 0,
+        content: '',
+        finishReason: 'stop',
+      },
+    ];
+
+    const streamingResponse = createStreamingResponse(chunks);
+    const ai = new AxMockAIService({
+      features: { functions: false, streaming: true },
+      chatResponse: streamingResponse as any,
+    });
+
+    const result = await gen.forward(
+      ai,
+      { prompt: 'Provide a structured analysis.' },
+      { stream: true }
+    );
+
+    expect(result).toEqual({
+      summary:
+        'Summary unavailable from the model output — manual review recommended to craft a compliant paragraph.',
+      sentiment: 'neutral',
+      confidence: 0.65,
+      highlights: [
+        'The model response did not supply valid highlights. Highlight that human review is required.',
+        'Flag the output for review because automatic repair was triggered.',
+      ],
+      actionPlan: 'No immediate action required.',
+    });
+  });
+});
+
 describe('AxGen Signature Validation', () => {
   it('should validate signature on construction and fail for incomplete signature', () => {
     // This should throw when trying to create AxGen with a signature that has only input fields