@od-oneapp/ai-platform 0.1.0

package/src/tools/compliance/compliance-wrapper.ts

@@ -0,0 +1,789 @@
+/**
+ * @fileoverview Compliance Middleware Wrapper
+ * Compliance Middleware Wrapper
+ * Wraps LanguageModelV3 with PHI redaction, ABAC, and audit logging
+ * Implements the Master Playbook's model-boundary guardrails
+ *
+ * Note: This is a simplified wrapper due to AI SDK v6 API changes.
+ * For production, consider using pre/post-processing hooks instead.
+ */
+
+import {
+  logGeneration as auditLogGeneration,
+  getAuditLogger,
+  logRedaction,
+} from '../../governance/audit/audit-logger';
+import { evaluatePolicy } from '../../governance/compliance/abac/policy-engine';
+import { redactPHI } from '../../governance/compliance/validators/redaction-utils';
+import { hashInput } from '../../server/cache/crypto';
+import { wrapLanguageModel } from '../../shared';
+
+import type {
+  ABACContext,
+  ComplianceOptions,
+  PHIDetection,
+  Resource,
+} from '../../governance/compliance/types';
+import type { WrapGenerateArgs, WrapStreamArgs } from '../../sdk/middleware/model-middleware';
+import type {
+  SDKLanguageModelV3CallOptions as LanguageModelV3CallOptions,
+  SDKLanguageModelV3Content as LanguageModelV3Content,
+  SDKLanguageModelV3GenerateResult as LanguageModelV3GenerateResult,
+  SDKLanguageModelV3Message as LanguageModelV3Message,
+  SDKLanguageModelV3Prompt as LanguageModelV3Prompt,
+  SDKLanguageModelV3StreamPart as LanguageModelV3StreamPart,
+  SDKLanguageModelV3ToolResult as LanguageModelV3ToolResult,
+  SDKLanguageModelV3ToolResultOutput as LanguageModelV3ToolResultOutput,
+
+  SDKFilePart,
+  SDKImagePart,
+  SDKLanguageModelV3,
+  SDKLanguageModelV3Middleware,
+  SDKModelMessage,
+  SDKTextPart,
+  SDKToolResultPart} from '../../shared';
+
+
+
+// Define POLICY_VERSION locally since it doesn't exist in config
+const POLICY_VERSION = '1.0.0';
+
+/**
+ * Build a default ABAC context for compliance operations.
+ */
+function createDefaultABACContext(actorId: string): ABACContext {
+  return {
+    actor: {
+      id: actorId,
+      roles: [],
+      attributes: {},
+    },
+    environment: {
+      timestamp: new Date(),
+      ipAddress: 'unknown',
+    },
+    purpose: 'treatment',
+  };
+}
+
+// Re-export for use in other modules
+export { auditLogGeneration as logGeneration };
+
+// Default compliance options (unicorn/no-object-as-default-parameter requires this pattern)
+const DEFAULT_COMPLIANCE_OPTIONS: ComplianceOptions = {
+  enablePHIRedaction: true,
+  enableABAC: false,
+  policyVersion: POLICY_VERSION,
+};
+
+const isRecord = (value: unknown): value is Record<string, unknown> =>
+  typeof value === 'object' && value !== null;
+
+const redactStringValue = (
+  text: string,
+  strategy: ComplianceOptions['redactionStrategy'],
+  detections: PHIDetection[],
+  minConfidence: number,
+): string => {
+  const { redacted, detections: found } = redactPHI(text, strategy, { minConfidence });
+  detections.push(...found);
+  return redacted;
+};
+
+const redactUnknownValue = (
+  value: unknown,
+  strategy: ComplianceOptions['redactionStrategy'],
+  detections: PHIDetection[],
+  minConfidence: number,
+): unknown => {
+  if (typeof value === 'string') {
+    return redactStringValue(value, strategy, detections, minConfidence);
+  }
+
+  if (Array.isArray(value)) {
+    return value.map(entry => redactUnknownValue(entry, strategy, detections, minConfidence));
+  }
+
+  if (isRecord(value)) {
+    return Object.fromEntries(
+      Object.entries(value).map(([key, entry]) => [
+        key,
+        redactUnknownValue(entry, strategy, detections, minConfidence),
+      ]),
+    );
+  }
+
+  return value;
+};
+
+const redactToolResultOutput = (
+  output: LanguageModelV3ToolResultOutput,
+  strategy: ComplianceOptions['redactionStrategy'],
+  detections: PHIDetection[],
+  minConfidence: number,
+): LanguageModelV3ToolResultOutput => {
+  if (output.type === 'text') {
+    return {
+      type: 'text',
+      value: redactStringValue(output.value, strategy, detections, minConfidence),
+      providerOptions: output.providerOptions,
+    };
+  }
+
+  if (output.type === 'json') {
+    return {
+      type: 'json',
+      value: redactUnknownValue(output.value, strategy, detections, minConfidence) as NonNullable<
+        typeof output.value
+      >,
+      providerOptions: output.providerOptions,
+    };
+  }
+
+  return output;
+};
+
+const redactPrompt = (
+  prompt: LanguageModelV3Prompt,
+  options: ComplianceOptions,
+): { prompt: LanguageModelV3Prompt; detections: PHIDetection[] } => {
+  const detections: PHIDetection[] = [];
+  const strategy = options.redactionStrategy ?? 'MASK';
+
+  const processedPrompt = prompt.map(message => {
+    if (message.role === 'system') {
+      return {
+        ...message,
+        content: redactStringValue(message.content, strategy, detections, 0.6),
+      } satisfies LanguageModelV3Message;
+    }
+
+    if (message.role === 'user') {
+      const parts = Array.isArray(message.content) ? message.content : [];
+      const redactedParts = parts.map(part => {
+        if (part.type === 'text') {
+          return {
+            ...part,
+            text: redactStringValue(part.text, strategy, detections, 0.6),
+          };
+        }
+        return part;
+      });
+      return {
+        ...message,
+        content: redactedParts,
+      } satisfies LanguageModelV3Message;
+    }
+
+    if (message.role === 'assistant') {
+      const parts = Array.isArray(message.content) ? message.content : [];
+      const redactedParts = parts.map(part => {
+        if (part.type === 'text' || part.type === 'reasoning') {
+          return {
+            ...part,
+            text: redactStringValue(part.text, strategy, detections, 0.6),
+          };
+        }
+
+        if (part.type === 'tool-result') {
+          return {
+            ...part,
+            output: redactToolResultOutput(part.output, strategy, detections, 0.6),
+          };
+        }
+
+        if (part.type === 'tool-call') {
+          const inputValue = part.input;
+          if (typeof inputValue === 'string') {
+            let redactedInput = inputValue;
+            try {
+              const parsed = JSON.parse(inputValue) as unknown;
+              const redacted = redactUnknownValue(parsed, strategy, detections, 0.6);
+              redactedInput = JSON.stringify(redacted);
+            } catch {
+              redactedInput = redactStringValue(inputValue, strategy, detections, 0.6);
+            }
+            return { ...part, input: redactedInput };
+          }
+
+          const redactedInput = redactUnknownValue(inputValue, strategy, detections, 0.6);
+          return { ...part, input: redactedInput };
+        }
+
+        return part;
+      });
+      return {
+        ...message,
+        content: redactedParts,
+      } satisfies LanguageModelV3Message;
+    }
+
+    if (message.role === 'tool') {
+      const parts = Array.isArray(message.content) ? message.content : [];
+      const redactedParts = parts.map(part => {
+        if (part.type === 'tool-result') {
+          return {
+            ...part,
+            output: redactToolResultOutput(part.output, strategy, detections, 0.6),
+          };
+        }
+        return part;
+      });
+      return {
+        ...message,
+        content: redactedParts,
+      } satisfies LanguageModelV3Message;
+    }
+
+    return message;
+  });
+
+  return { prompt: processedPrompt, detections };
+};
+
+const redactGenerateResult = (
+  result: LanguageModelV3GenerateResult,
+  options: ComplianceOptions,
+): { result: LanguageModelV3GenerateResult; detections: PHIDetection[] } => {
+  const detections: PHIDetection[] = [];
+  const strategy = options.redactionStrategy ?? 'MASK';
+
+  const content = result.content.map(part => {
+    if (part.type === 'text' || part.type === 'reasoning') {
+      return {
+        ...part,
+        text: redactStringValue(part.text, strategy, detections, 0.7),
+      };
+    }
+
+    if (part.type === 'tool-result') {
+      const toolResult = part;
+      const redactedResult = redactUnknownValue(toolResult.result, strategy, detections, 0.7);
+      return { ...toolResult, result: redactedResult as LanguageModelV3ToolResult['result'] };
+    }
+
+    return part;
+  });
+
+  return {
+    result: {
+      ...result,
+      content: content as LanguageModelV3Content[],
+    },
+    detections,
+  };
+};
+
+const createComplianceStreamTransform = (
+  options: ComplianceOptions,
+  onDetections: (detections: PHIDetection[]) => Promise<void>,
+): TransformStream<LanguageModelV3StreamPart, LanguageModelV3StreamPart> => {
+  const detections: PHIDetection[] = [];
+  const strategy = options.redactionStrategy ?? 'MASK';
+  const holdBackSize = 64;
+  const buffers = new Map<string, { pending: string; deltaType: string }>();
+
+  const extractDelta = (part: LanguageModelV3StreamPart): string | null => {
+    if (
+      part.type === 'text-delta' ||
+      part.type === 'reasoning-delta' ||
+      part.type === 'tool-input-delta'
+    ) {
+      return part.delta;
+    }
+    return null;
+  };
+
+  const getKey = (part: LanguageModelV3StreamPart): string | null => {
+    if (
+      part.type === 'text-delta' ||
+      part.type === 'reasoning-delta' ||
+      part.type === 'tool-input-delta'
+    ) {
+      return `${part.type}:${part.id}`;
+    }
+    if (
+      part.type === 'text-end' ||
+      part.type === 'reasoning-end' ||
+      part.type === 'tool-input-end'
+    ) {
+      const deltaType = part.type.replace('-end', '-delta');
+      return `${deltaType}:${part.id}`;
+    }
+    if (
+      part.type === 'text-start' ||
+      part.type === 'reasoning-start' ||
+      part.type === 'tool-input-start'
+    ) {
+      const deltaType = part.type.replace('-start', '-delta');
+      return `${deltaType}:${part.id}`;
+    }
+    return null;
+  };
+
+  const getDeltaType = (part: LanguageModelV3StreamPart): string | null => {
+    if (
+      part.type === 'text-delta' ||
+      part.type === 'reasoning-delta' ||
+      part.type === 'tool-input-delta'
+    ) {
+      return part.type;
+    }
+    if (
+      part.type === 'text-start' ||
+      part.type === 'reasoning-start' ||
+      part.type === 'tool-input-start'
+    ) {
+      return part.type.replace('-start', '-delta');
+    }
+    if (
+      part.type === 'text-end' ||
+      part.type === 'reasoning-end' ||
+      part.type === 'tool-input-end'
+    ) {
+      return part.type.replace('-end', '-delta');
+    }
+    return null;
+  };
+
+  return new TransformStream<LanguageModelV3StreamPart, LanguageModelV3StreamPart>({
+    transform(part, controller) {
+      const delta = extractDelta(part);
+      const key = getKey(part);
+      const deltaType = getDeltaType(part);
+
+      if (delta && key && deltaType) {
+        const state = buffers.get(key) ?? { pending: '', deltaType };
+        const pending = state.pending + delta;
+
+        if (pending.length <= holdBackSize) {
+          buffers.set(key, { ...state, pending });
+          return;
+        }
+
+        const safePrefix = pending.slice(0, pending.length - holdBackSize);
+        const redactedPrefix = redactStringValue(safePrefix, strategy, detections, 0.7);
+        buffers.set(key, { ...state, pending: pending.slice(pending.length - holdBackSize) });
+
+        if (redactedPrefix.length > 0) {
+          controller.enqueue({
+            ...part,
+            delta: redactedPrefix,
+          } as LanguageModelV3StreamPart);
+        }
+        return;
+      }
+
+      if (
+        part.type === 'text-end' ||
+        part.type === 'reasoning-end' ||
+        part.type === 'tool-input-end'
+      ) {
+        if (key) {
+          const state = buffers.get(key);
+          if (state && state.pending.length > 0) {
+            const redactedTail = redactStringValue(state.pending, strategy, detections, 0.7);
+            if (redactedTail.length > 0 && deltaType) {
+              controller.enqueue({
+                type: deltaType as LanguageModelV3StreamPart['type'],
+                id: part.id,
+                delta: redactedTail,
+                providerMetadata: 'providerMetadata' in part ? part.providerMetadata : undefined,
+              } as LanguageModelV3StreamPart);
+            }
+          }
+          buffers.delete(key);
+        }
+        controller.enqueue(part);
+        return;
+      }
+
+      if (
+        part.type === 'text-start' ||
+        part.type === 'reasoning-start' ||
+        part.type === 'tool-input-start'
+      ) {
+        if (key && deltaType && !buffers.has(key)) {
+          buffers.set(key, { pending: '', deltaType });
+        }
+        controller.enqueue(part);
+        return;
+      }
+
+      controller.enqueue(part);
+    },
+    async flush(controller) {
+      for (const [key, state] of buffers.entries()) {
+        if (state.pending.length > 0) {
+          const redactedTail = redactStringValue(state.pending, strategy, detections, 0.7);
+          if (redactedTail.length > 0) {
+            controller.enqueue({
+              type: state.deltaType as LanguageModelV3StreamPart['type'],
+              id: key.split(':').slice(1).join(':'),
+              delta: redactedTail,
+            } as LanguageModelV3StreamPart);
+          }
+        }
+      }
+      buffers.clear();
+      await onDetections(detections);
+      controller.terminate();
+    },
+  });
+};
+
+const buildDefaultResource = (context: ABACContext | undefined): Resource => ({
+  id: 'ai-generation',
+  type: 'ai-generation',
+  sensitivity: 'confidential',
+  ownerId: context?.actor.id,
+});
+
+const enforceABAC = async (options: ComplianceOptions): Promise<void> => {
+  if (!options.enableABAC) {
+    return;
+  }
+
+  const context = options.abacContext;
+  if (!context) {
+    throw new Error('ABAC enforcement requires abacContext');
+  }
+
+  const resource = options.resource ?? buildDefaultResource(context);
+  const evaluation = await evaluatePolicy(context, resource, 'execute');
+
+  if (!evaluation.decision.allowed) {
+    throw new Error(
+      evaluation.decision.reason ?? `ABAC denied by policy ${evaluation.decision.policyId}`,
+    );
+  }
+};
+
+const logGenerationEvent = async (
+  options: ComplianceOptions,
+  params: LanguageModelV3CallOptions,
+  output: unknown,
+): Promise<void> => {
+  const logger = options.auditLogger ?? getAuditLogger();
+  const inputHash = await hashInput(JSON.stringify(params.prompt ?? []));
+  await logger.log({
+    actor: options.abacContext?.actor.id ?? 'system',
+    action: 'generation',
+    input_hash: inputHash,
+    output_snapshot: output,
+    policy_version: options.policyVersion ?? POLICY_VERSION,
+    metadata: {
+      phiRedaction: options.enablePHIRedaction,
+      abacEnabled: options.enableABAC,
+    },
+  });
+};
+
+/**
+ * Wrap a language model with compliance middleware
+ *
+ * @param model - Base LanguageModelV3 to wrap
+ * @param options - Compliance configuration
+ * @returns Wrapped model with compliance guardrails
+ *
+ * @example
+ * ```ts
+ * import { openai } from '@ai-sdk/openai';
+ * import { wrapWithCompliance } from '#/lib/ai-onedigital';
+ *
+ * const complianceModel = wrapWithCompliance(
+ *   openai('gpt-5.1-thinking'),
+ *   {
+ *     enablePHIRedaction: true,
+ *     enableABAC: true,
+ *     policyVersion: 'v2.0',
+ *   }
+ * );
+ * ```
+ */
+export function wrapWithCompliance(
+  model: SDKLanguageModelV3,
+  options: ComplianceOptions,
+): SDKLanguageModelV3 {
+  const complianceOptions: ComplianceOptions = {
+    ...DEFAULT_COMPLIANCE_OPTIONS,
+    ...options,
+  };
+
+  const middleware: SDKLanguageModelV3Middleware = {
+    specificationVersion: 'v3',
+    transformParams: async ({ params }) => {
+      if (!complianceOptions.enablePHIRedaction) {
+        return params;
+      }
+
+      const { prompt, detections } = redactPrompt(params.prompt, complianceOptions);
+      if (detections.length > 0 && complianceOptions.abacContext) {
+        await logRedaction(complianceOptions.abacContext.actor.id, detections, {
+          count: detections.length,
+          location: 'input',
+        });
+      }
+      return {
+        ...params,
+        prompt,
+      };
+    },
+    wrapGenerate: async ({ doGenerate, params }: WrapGenerateArgs) => {
+      await enforceABAC(complianceOptions);
+      const result = await doGenerate();
+
+      if (!complianceOptions.enablePHIRedaction) {
+        await logGenerationEvent(complianceOptions, params, result.content);
+        return result;
+      }
+
+      const { result: redactedResult, detections } = redactGenerateResult(
+        result,
+        complianceOptions,
+      );
+      if (detections.length > 0 && complianceOptions.abacContext) {
+        await logRedaction(complianceOptions.abacContext.actor.id, detections, {
+          count: detections.length,
+          location: 'output',
+        });
+      }
+      await logGenerationEvent(complianceOptions, params, redactedResult.content);
+
+      return redactedResult;
+    },
+    wrapStream: async ({ doStream, params }: WrapStreamArgs) => {
+      await enforceABAC(complianceOptions);
+      const result = await doStream();
+
+      if (!complianceOptions.enablePHIRedaction) {
+        return result;
+      }
+
+      const stream = result.stream.pipeThrough(
+        createComplianceStreamTransform(complianceOptions, async detections => {
+          if (detections.length > 0 && complianceOptions.abacContext) {
+            await logRedaction(complianceOptions.abacContext.actor.id, detections, {
+              count: detections.length,
+              location: 'output',
+            });
+          }
+          await logGenerationEvent(complianceOptions, params, {
+            streamCompleted: true,
+          });
+        }),
+      );
+
+      return {
+        ...result,
+        stream,
+      };
+    },
+  };
+
+  return wrapLanguageModel({
+    model,
+    middleware,
+  });
+}
+
+/**
+ * Quick helper to create a compliance-wrapped model with defaults
+ */
+export function createCompliantModel(
+  model: SDKLanguageModelV3,
+  actorId: string = 'anonymous',
+  options: Partial<ComplianceOptions> = {},
+): SDKLanguageModelV3 {
+  return wrapWithCompliance(model, {
+    enablePHIRedaction: true,
+    enableABAC: false, // Disabled by default for simplicity
+    policyVersion: POLICY_VERSION,
+    abacContext: options.abacContext ?? createDefaultABACContext(actorId),
+    ...options,
+  });
+}
+
+/**
+ * Helper to pre-process messages for PHI redaction
+ * Use this in your route handlers before calling streamText/generateText
+ */
+export async function preprocessMessages(
+  messages: SDKModelMessage[],
+  options: ComplianceOptions = DEFAULT_COMPLIANCE_OPTIONS,
+): Promise<{
+  messages: SDKModelMessage[];
+  detections: PHIDetection[];
+}> {
+  const { enablePHIRedaction, redactionStrategy = 'MASK', abacContext } = options;
+
+  if (!enablePHIRedaction) {
+    return { messages, detections: [] };
+  }
+
+  const allDetections: PHIDetection[] = [];
+
+  const isUserPart = (part: SDKTextPart | SDKFilePart | SDKImagePart): part is SDKTextPart =>
+    part.type === 'text';
+
+  const processedMessages = messages.map(msg => {
+    if (msg.role === 'tool') {
+      if (typeof msg.content === 'string') {
+        const { redacted, detections } = redactPHI(msg.content, redactionStrategy, {
+          minConfidence: 0.6,
+        });
+        allDetections.push(...detections);
+        const toolResult: SDKToolResultPart = {
+          type: 'tool-result',
+          toolCallId: 'unknown',
+          toolName: 'unknown',
+          output: { type: 'text', value: redacted },
+        };
+        return {
+          ...msg,
+          content: [toolResult],
+        };
+      }
+
+      if (Array.isArray(msg.content)) {
+        const updatedToolContent = msg.content.map(part => {
+          if (part.type === 'tool-result') {
+            if (part.output.type === 'text' || part.output.type === 'error-text') {
+              const { redacted, detections } = redactPHI(part.output.value, redactionStrategy, {
+                minConfidence: 0.6,
+              });
+              allDetections.push(...detections);
+              return {
+                ...part,
+                output: { ...part.output, value: redacted },
+              };
+            }
+            return part;
+          }
+          return part;
+        });
+
+        return {
+          ...msg,
+          content: updatedToolContent,
+        };
+      }
+
+      return msg;
+    }
+
+    if (msg.role === 'system') {
+      if (typeof msg.content === 'string') {
+        const { redacted, detections } = redactPHI(msg.content, redactionStrategy, {
+          minConfidence: 0.6,
+        });
+        allDetections.push(...detections);
+        return { ...msg, content: redacted };
+      }
+      return msg;
+    }
+
+    if (msg.role === 'user') {
+      if (typeof msg.content === 'string') {
+        const { redacted, detections } = redactPHI(msg.content, redactionStrategy, {
+          minConfidence: 0.6,
+        });
+        allDetections.push(...detections);
+        return { ...msg, content: redacted };
+      }
+
+      if (Array.isArray(msg.content)) {
+        const updatedParts = msg.content.map(part => {
+          if (isUserPart(part)) {
+            const { redacted, detections } = redactPHI(part.text, redactionStrategy, {
+              minConfidence: 0.6,
+            });
+            allDetections.push(...detections);
+            return { ...part, text: redacted };
+          }
+          return part;
+        });
+
+        return {
+          ...msg,
+          content: updatedParts,
+        };
+      }
+
+      return msg;
+    }
+
+    if (typeof msg.content === 'string') {
+      const { redacted, detections } = redactPHI(msg.content, redactionStrategy, {
+        minConfidence: 0.6,
+      });
+
+      allDetections.push(...detections);
+
+      return {
+        ...msg,
+        content: redacted,
+      };
+    }
+    if (Array.isArray(msg.content)) {
+      const updatedParts = msg.content.map(part => {
+        if (part && typeof part === 'object' && 'type' in part && part.type === 'text') {
+          const textPart = part;
+          const { redacted, detections } = redactPHI(textPart.text, redactionStrategy, {
+            minConfidence: 0.6,
+          });
+          allDetections.push(...detections);
+          return { ...textPart, text: redacted };
+        }
+        return part;
+      });
+
+      return {
+        ...msg,
+        content: updatedParts,
+      };
+    }
+
+    return msg;
+  });
+
+  // Log redactions if any
+  if (allDetections.length > 0 && abacContext) {
+    await logRedaction(abacContext.actor.id, allDetections, {
+      count: allDetections.length,
+    });
+  }
+
+  return { messages: processedMessages, detections: allDetections };
+}
+
+/**
+ * Helper to post-process text output for PHI redaction
+ */
+export async function postprocessText(
+  text: string,
+  options: ComplianceOptions = DEFAULT_COMPLIANCE_OPTIONS,
+): Promise<{
+  text: string;
+  detections: PHIDetection[];
+}> {
+  const { enablePHIRedaction, redactionStrategy = 'MASK', abacContext } = options;
+
+  if (!enablePHIRedaction) {
+    return { text, detections: [] };
+  }
+
+  const { redacted, detections } = redactPHI(text, redactionStrategy, {
+    minConfidence: 0.7,
+  });
+
+  // Log redactions if any
+  if (detections.length > 0 && abacContext) {
+    await logRedaction(abacContext.actor.id, detections, {
+      location: 'output',
+      count: detections.length,
+    });
+  }
+
+  return { text: redacted, detections };
+}