CedarCopilot/cedar-mail

/**
 * Zustand capture middleware for the deep debugger.
 *
 * Wraps the store's `set` so every `setState` lands in the timeline buffer as
 * a `zustand` entry — the before/after diff of changed top-level keys plus a
 * best-effort caller (so you can see which action fired the change).
 *
 * Replaces the old external `store.subscribe` approach: a middleware sees
 * every discrete `set` call (not coalesced snapshots) and runs inside the
 * caller's stack, so it can attribute the change. Apply it as the OUTERMOST
 * middleware in a store's `create()` chain.
 */

import type { StateCreator, StoreMutatorIdentifier } from 'zustand';
import { push, isEnabled, isCaptureBodies } from './buffer';
import { redact } from './redact';

// Store-internal keys that are noisy or risk feedback loops — never diffed.
const SKIP_KEYS = new Set<string>([
  'agentConnectionLogs',
  'messages',
  'voiceState',
  'voiceSettings',
]);

const PREVIEW_MAX_DEPTH = 3;
const PREVIEW_MAX_KEYS = 40;
const PREVIEW_MAX_STRING = 300;

function changedTopLevelKeys(
  next: Record<string, unknown>,
  prev: Record<string, unknown>,
): string[] {
  const keys: string[] = [];
  const all = new Set([...Object.keys(next), ...Object.keys(prev)]);
  for (const k of all) {
    if (SKIP_KEYS.has(k)) continue;
    // Functions are stable references — they never "change" and aren't data.
    if (typeof next[k] === 'function') continue;
    if (next[k] !== prev[k]) keys.push(k);
  }
  return keys;
}

/**
 * Depth- and breadth-bounded preview. Crucially this does NOT traverse the
 * whole value — capture runs on every setState, and store slices can hold
 * hundreds of KB (threadData, aopsById, …). Long strings are sliced (O(cap),
 * not O(length)) so a 350 KB email body never gets fully serialized.
 */
function boundedPreview(value: unknown, depth: number): unknown {
  if (typeof value === 'string') {
    return value.length > PREVIEW_MAX_STRING
      ? `${value.slice(0, PREVIEW_MAX_STRING)}…[+${value.length - PREVIEW_MAX_STRING} chars]`
      : value;
  }
  if (value === null || typeof value !== 'object') return value;
  if (depth <= 0) {
    return Array.isArray(value) ? `[Array(${value.length})]` : '[Object]';
  }
  if (Array.isArray(value)) {
    const head = value.slice(0, PREVIEW_MAX_KEYS).map((v) => boundedPreview(v, depth - 1));
    return value.length > PREVIEW_MAX_KEYS
      ? [...head, `…+${value.length - PREVIEW_MAX_KEYS} more`]
      : head;
  }
  const entries = Object.entries(value as Record<string, unknown>);
  const out: Record<string, unknown> = {};
  for (const [k, v] of entries.slice(0, PREVIEW_MAX_KEYS)) {
    out[k] = boundedPreview(v, depth - 1);
  }
  if (entries.length > PREVIEW_MAX_KEYS) {
    out['…'] = `+${entries.length - PREVIEW_MAX_KEYS} more keys`;
  }
  return out;
}

function summarize(v: unknown): unknown {
  try {
    // Bound first (cheap, no full traversal), then redact the small result.
    return redact(boundedPreview(v, PREVIEW_MAX_DEPTH), { captureBodies: isCaptureBodies() });
  } catch {
    return '[unserializable]';
  }
}

/**
 * Best-effort caller attribution: the first stack frame outside zustand,
 * immer, and this file — i.e. the action (or component) that called `set`.
 */
function callerFrame(): string | undefined {
  const stack = new Error().stack;
  if (!stack) return undefined;
  for (const line of stack.split('\n').slice(2)) {
    if (
      line.includes('zustandCapture') ||
      line.includes('/zustand/') ||
      line.includes('/immer/')
    ) {
      continue;
    }
    return line.trim().replace(/^at\s+/, '');
  }
  return undefined;
}

// Middleware signature that leaves the store type (and mutator chain) unchanged.
type CaptureMiddleware = <
  T,
  Mps extends [StoreMutatorIdentifier, unknown][] = [],
  Mcs extends [StoreMutatorIdentifier, unknown][] = [],
>(
  creator: StateCreator<T, Mps, Mcs>,
) => StateCreator<T, Mps, Mcs>;

type CaptureImpl = <T>(creator: StateCreator<T, [], []>) => StateCreator<T, [], []>;

const captureImpl: CaptureImpl = (creator) => (set, get, api) => {
  const loggedSet: typeof set = (...args) => {
    if (!isEnabled()) {
      set(...(args as Parameters<typeof set>));
      return;
    }
    const prev = get() as Record<string, unknown>;
    set(...(args as Parameters<typeof set>));
    const next = get() as Record<string, unknown>;
    const changedKeys = changedTopLevelKeys(next, prev);
    if (changedKeys.length === 0) return;
    const diff: Record<string, { from?: unknown; to?: unknown }> = {};
    for (const k of changedKeys) {
      diff[k] = { from: summarize(prev[k]), to: summarize(next[k]) };
    }
    // devtools forwards an action name as the third arg (often 'anonymous').
    const actionName = (args as unknown[])[2];
    const named =
      typeof actionName === 'string' && actionName !== 'anonymous'
        ? actionName
        : undefined;
    push({
      source: 'zustand',
      changedKeys,
      diff,
      caller: named ?? callerFrame(),
    });
  };
  return creator(loggedSet, get, api);
};

export const captureMiddleware = captureImpl as CaptureMiddleware;