Structured output and tool calls that do not break production

/** The subset of JSON Schema that model output schemas actually need. */
export type Schema =
  | { type: 'object'; properties: Record<string, Schema>; required?: string[]; additionalProperties?: false; description?: string }
  | { type: 'array'; items: Schema; minItems?: number; maxItems?: number; description?: string }
  | { type: 'string'; enum?: string[]; pattern?: string; maxLength?: number; description?: string }
  | { type: 'number' | 'integer'; minimum?: number; maximum?: number; description?: string }
  | { type: 'boolean'; description?: string }
  | { type: 'null' }
  | { anyOf: Schema[]; description?: string };

export interface ValidationError {
  /** JSON Pointer to the offending value, e.g. /lines/2/amount */
  path: string;
  message: string;
}

/**
 * Returns every error, not just the first: the list is fed back to the model
 * in a repair attempt, and fixing all problems at once saves a round trip.
 */
export function validate(value: unknown, schema: Schema, path = ''): ValidationError[] {
  const err = (message: string): ValidationError[] => [{ path: path || '/', message }];
  if ('anyOf' in schema) {
    const results = schema.anyOf.map((s) => validate(value, s, path));
    return results.some((r) => r.length === 0) ? [] : err(`does not match any allowed shape`);
  }
  switch (schema.type) {
    case 'null':
      return value === null ? [] : err('expected null');
    case 'boolean':
      return typeof value === 'boolean' ? [] : err('expected boolean');
    case 'number':
    case 'integer': {
      if (typeof value !== 'number' || !Number.isFinite(value)) return err(`expected ${schema.type}`);
      if (schema.type === 'integer' && !Number.isInteger(value)) return err('expected integer');
      if (schema.minimum !== undefined && value < schema.minimum) return err(`must be >= ${schema.minimum}`);
      if (schema.maximum !== undefined && value > schema.maximum) return err(`must be <= ${schema.maximum}`);
      return [];
    }
    case 'string': {
      if (typeof value !== 'string') return err('expected string');
      if (schema.enum && !schema.enum.includes(value)) return err(`must be one of ${schema.enum.join(', ')}`);
      if (schema.maxLength !== undefined && value.length > schema.maxLength) return err(`longer than ${schema.maxLength}`);
      if (schema.pattern && !new RegExp(schema.pattern, 'u').test(value)) return err(`does not match ${schema.pattern}`);
      return [];
    }
    case 'array': {
      if (!Array.isArray(value)) return err('expected array');
      const out: ValidationError[] = [];
      if (schema.minItems !== undefined && value.length < schema.minItems) out.push(...err(`fewer than ${schema.minItems} items`));
      if (schema.maxItems !== undefined && value.length > schema.maxItems) out.push(...err(`more than ${schema.maxItems} items`));
      value.forEach((v, i) => out.push(...validate(v, schema.items, `${path}/${i}`)));
      return out;
    }
    case 'object': {
      if (typeof value !== 'object' || value === null || Array.isArray(value)) return err('expected object');
      const obj = value as Record<string, unknown>;
      const out: ValidationError[] = [];
      for (const k of schema.required ?? []) if (!(k in obj)) out.push({ path: `${path}/${k}`, message: 'is required' });
      for (const [k, v] of Object.entries(obj)) {
        const sub = schema.properties[k];
        if (sub) out.push(...validate(v, sub, `${path}/${k}`));
        else if (schema.additionalProperties === false) out.push({ path: `${path}/${k}`, message: 'is not allowed' });
      }
      return out;
    }
  }
}

/**
 * Models wrap JSON in prose or code fences more often than anyone would like.
 * Take the outermost object; anything else is a parse failure, not a guess.
 */
export function extractJson(raw: string): unknown {
  const fenced = raw.match(/```(?:json)?\s*([\s\S]*?)```/);
  const text = (fenced ? fenced[1]! : raw).trim();
  const start = text.indexOf('{');
  const end = text.lastIndexOf('}');
  if (start === -1 || end < start) throw new SyntaxError('no JSON object in output');
  return JSON.parse(text.slice(start, end + 1));
}

/**
 * Generate → parse → validate → repair. On failure the model sees its own
 * output and the exact list of errors, which fixes most problems in one
 * extra call. Business rules that a schema cannot express go in `check`.
 */
export async function generateStructured<T>(o: {
  complete: Complete;
  system: string;
  input: string;
  schema: Schema;
  check?: (value: T) => ValidationError[];
  maxAttempts?: number;
}): Promise<{ value: T; attempts: number }> {
  const max = o.maxAttempts ?? 2;
  const system = `${o.system}\n\nRespond with a single JSON object that matches this JSON Schema, and nothing else:\n${JSON.stringify(o.schema)}`;
  const messages: Message[] = [{ role: 'user', content: o.input }];
  let errors: ValidationError[] = [];

  for (let attempt = 1; attempt <= max; attempt++) {
    const raw = await o.complete(system, messages);
    let value: unknown;
    try {
      value = extractJson(raw);
      errors = validate(value, o.schema);
      if (!errors.length && o.check) errors = o.check(value as T);
    } catch (e) {
      errors = [{ path: '/', message: `invalid JSON: ${(e as Error).message}` }];
    }
    if (!errors.length) return { value: value as T, attempts: attempt };
    messages.push(
      { role: 'assistant', content: raw },
      { role: 'user', content: `Your output failed validation:\n${errors.map((e) => `- ${e.path}: ${e.message}`).join('\n')}\nReturn the corrected JSON object only.` },
    );
  }
  throw new StructuredOutputError(`no valid output after ${max} attempts`, errors, max);
}

export interface Tool<A = unknown, R = unknown> {
  name: string;
  /** Read by the model. Treat it as part of the prompt: versioned and reviewed. */
  description: string;
  input: Schema;
  /** Writes change the world and get stricter handling than reads. */
  effect: 'read' | 'write';
  run(args: A, ctx: ToolContext): Promise<R>;
}

export interface ToolContext {
  /** From the authenticated session, never from the model's arguments. */
  tenantId: string;
  signal: AbortSignal;
}

/** What the model gets back. Errors are data the model can act on, not exceptions. */
export type ToolResult =
  | { ok: true; content: unknown; replayed?: boolean }
  | { ok: false; error: 'unknown_tool' | 'invalid_arguments' | 'not_permitted' | 'in_progress' | 'conflict' | 'failed'; detail: string };

export class ToolRegistry {
  private readonly tools = new Map<string, Tool<never, unknown>>();

  constructor(private readonly idempotency: IdempotencyStore, private readonly timeoutMs = 10_000) {}

  register<A, R>(tool: Tool<A, R>): void {
    this.tools.set(tool.name, tool as unknown as Tool<never, unknown>);
  }

  /** Tool definitions for the provider API, filtered by what this caller may use. */
  definitions(allowWrites: boolean) {
    return [...this.tools.values()]
      .filter((t) => allowWrites || t.effect === 'read')
      .map((t) => ({ name: t.name, description: t.description, input_schema: t.input }));
  }

  async execute(call: ToolCall, ctx: ToolContext & { allowWrites: boolean }): Promise<ToolResult> {
    const tool = this.tools.get(call.name);
    if (!tool) return { ok: false, error: 'unknown_tool', detail: `no tool named ${call.name}` };
    // Enforced here, not only by hiding the definition: models can and do call tools they were not offered.
    if (tool.effect === 'write' && !ctx.allowWrites) return { ok: false, error: 'not_permitted', detail: `${call.name} is not allowed here` };

    const errors = validate(call.arguments, tool.input);
    if (errors.length) return { ok: false, error: 'invalid_arguments', detail: errors.map((e) => `${e.path}: ${e.message}`).join('; ') };

    // Writes run at most once per tool call id, even if the turn is retried.
    if (tool.effect === 'write') {
      const begin = this.idempotency.begin(ctx.tenantId, `${call.name}:${call.id}`, call.arguments);
      if (begin.kind === 'replay') return { ok: true, content: begin.result, replayed: true };
      if (begin.kind === 'in-progress') return { ok: false, error: 'in_progress', detail: 'this call is already running' };
      if (begin.kind === 'conflict') return { ok: false, error: 'conflict', detail: 'call id reused with different arguments' };
    }

    const signal = AbortSignal.any([ctx.signal, AbortSignal.timeout(this.timeoutMs)]);
    try {
      const content = await tool.run(call.arguments as never, { tenantId: ctx.tenantId, signal });
      if (tool.effect === 'write') this.idempotency.complete(ctx.tenantId, `${call.name}:${call.id}`, content);
      return { ok: true, content };
    } catch (e) {
      if (tool.effect === 'write') this.idempotency.abandon(ctx.tenantId, `${call.name}:${call.id}`);
      return { ok: false, error: 'failed', detail: e instanceof Error ? e.message : String(e) };
    }
  }
}

/**
 * At-most-once execution for side effects. The key is scoped to the tenant,
 * and the request body is hashed: the same key with a different body is a
 * client bug and must fail loudly instead of returning someone else's result.
 * In production this lives in Redis or a database with a unique constraint;
 * the protocol is the same.
 */
export class IdempotencyStore {
  private readonly entries = new Map<string, Entry>();

  constructor(private readonly ttlMs: number, private readonly now: () => number = Date.now) {}

  static hash(body: unknown): string {
    return createHash('sha256').update(JSON.stringify(body)).digest('hex');
  }

  begin(tenantId: string, key: string, body: unknown): Begin {
    const id = `${tenantId}:${key}`;
    const requestHash = IdempotencyStore.hash(body);
    const e = this.entries.get(id);
    if (e && e.expiresAt > this.now()) {
      if (e.requestHash !== requestHash) return { kind: 'conflict' };
      return e.state === 'done' ? { kind: 'replay', result: e.result } : { kind: 'in-progress' };
    }
    this.entries.set(id, { requestHash, state: 'running', expiresAt: this.now() + this.ttlMs });
    return { kind: 'new' };
  }

  complete(tenantId: string, key: string, result: unknown): void {
    const e = this.entries.get(`${tenantId}:${key}`);
    if (e) Object.assign(e, { state: 'done', result });
  }

  /** On failure, release the key so a retry can run the operation again. */
  abandon(tenantId: string, key: string): void {
    this.entries.delete(`${tenantId}:${key}`);
  }
}