pig/apps/api/test/helpers/piggy-store.ts

/**
 * An in-memory transcript store, for driving the relay without a database.
 *
 * The relay's job is now half persistence, and the properties worth asserting
 * about it are about ORDER and ABOUT FAILURE: the question is filed before the
 * answer, tool evidence lands as it streams, and a store that throws must not
 * be able to reach the stream the user is reading. None of that needs SQL, and
 * a real Postgres would make it harder to assert — `failOn` here fails a
 * specific method on demand, which is the case that matters most and the one a
 * live database will not perform to order.
 *
 * What it is NOT is a second implementation of the store's semantics. Ownership
 * predicates, `seq` under concurrency and the capability gate are asserted
 * against a real database in piggy-conversations.test.ts, because that is where
 * they are either true or not.
 */
import { randomUUID } from 'node:crypto';
import type { ReadCapability } from '@pig/core';
import type { Principal } from '../../src/lib/auth';
import type {
  PiggyConversationCreateInput,
  PiggyConversationDetail,
  PiggyConversationOwner,
  PiggyMessageInput,
  PiggyTranscriptMessage,
  PiggyTranscriptStore,
} from '../../src/services/piggy-conversations';

export interface RecordedAppend {
  conversationId: string;
  message: PiggyMessageInput;
}

export interface RecordedConversation {
  id: string;
  userId: string;
  title: string;
  readCapability: ReadCapability;
}

export type PiggyStoreMethod = keyof PiggyTranscriptStore;

export interface RecordingTranscriptStore {
  store: PiggyTranscriptStore;
  /** Every append, in the order the store received it. */
  appends: RecordedAppend[];
  conversations: Map<string, RecordedConversation>;
  /** Conversation ids `linkAgentRuns` was called for. */
  linked: string[];
  /** Seed a conversation that already exists — a thread being resumed. */
  seed(conversation: {
    userId: string;
    title?: string;
    readCapability?: ReadCapability;
    messages?: { role: 'user' | 'assistant'; content: string }[];
  }): string;
}

export function recordingTranscriptStore(
  failOn: readonly PiggyStoreMethod[] = [],
): RecordingTranscriptStore {
  const appends: RecordedAppend[] = [];
  const conversations = new Map<string, RecordedConversation>();
  const linked: string[] = [];
  const logged: string[] = [];

  function refuse(method: PiggyStoreMethod): void {
    if (failOn.includes(method)) throw new Error(`the store was told to fail on ${method}`);
  }

  function transcriptOf(conversationId: string): RecordedAppend[] {
    return appends.filter((entry) => entry.conversationId === conversationId);
  }

  const store: PiggyTranscriptStore = {
    async create(
      owner: PiggyConversationOwner,
      input: PiggyConversationCreateInput = {},
    ): Promise<PiggyConversationDetail> {
      refuse('create');
      // The caller's id when it brought one, exactly as the column's primary
      // key does — a fake that minted its own would let a relay that loses the
      // client's id pass, and losing it strands every approval mid-turn.
      const id = input.id ?? randomUUID();
      if (conversations.has(id)) throw new Error(`conversation ${id} already exists`);
      conversations.set(id, {
        id,
        userId: owner.userId,
        title: input.title ?? input.firstMessage ?? 'New conversation',
        readCapability: input.readCapability ?? 'book:read',
      });
      const now = new Date().toISOString();
      return {
        id,
        title: conversations.get(id)?.title ?? '',
        model: input.model ?? null,
        mode: input.mode ?? null,
        context: input.context ?? null,
        createdAt: now,
        updatedAt: now,
        messages: [],
      };
    },

    async readCapabilityFor(
      owner: PiggyConversationOwner,
      id: string,
    ): Promise<ReadCapability | null> {
      refuse('readCapabilityFor');
      const conversation = conversations.get(id);
      // The predicate the real store puts in SQL: another person's thread and
      // an id that was never issued are the same answer.
      return conversation && conversation.userId === owner.userId
        ? conversation.readCapability
        : null;
    },

    async promptHistory(
      principal: Principal,
      id: string,
    ): Promise<{ role: 'user' | 'assistant'; content: string }[]> {
      refuse('promptHistory');
      const conversation = conversations.get(id);
      if (!conversation || conversation.userId !== principal.userId) return [];
      const turns: { role: 'user' | 'assistant'; content: string }[] = [];
      for (const entry of transcriptOf(id)) {
        const { role, content } = entry.message;
        // Tool rows are evidence, not context — the same exclusion the real
        // store makes, and the relay is tested against it.
        if ((role === 'user' || role === 'assistant') && content) turns.push({ role, content });
      }
      return turns;
    },

    async appendMessage(
      owner: PiggyConversationOwner,
      conversationId: string,
      message: PiggyMessageInput,
    ): Promise<PiggyTranscriptMessage | null> {
      refuse('appendMessage');
      const conversation = conversations.get(conversationId);
      // Null means "not yours", exactly as the real store's predicate does, so
      // a relay that starts writing into somebody else's thread fails here too.
      if (!conversation || conversation.userId !== owner.userId) return null;
      const seq = transcriptOf(conversationId).length;
      appends.push({ conversationId, message });
      // The conversation keeps the strongest capability any turn in it needed.
      if (message.readCapability === 'economics:read') {
        conversation.readCapability = 'economics:read';
      }
      return {
        id: randomUUID(),
        seq,
        role: message.role,
        content: message.content ?? '',
        reasoning: message.reasoning ?? null,
        model: message.model ?? null,
        mode: message.mode ?? null,
        inputTokens: message.inputTokens ?? null,
        outputTokens: message.outputTokens ?? null,
        costMicroCents: message.costMicroCents ?? null,
        finishReason: message.finishReason ?? null,
        tool: message.tool
          ? {
              callId: message.tool.callId,
              name: message.tool.name,
              arguments: message.tool.arguments ?? null,
              result: message.tool.result ?? null,
              ok: message.tool.ok ?? null,
            }
          : null,
        approval: message.approval
          ? {
              id: message.approval.change.id,
              change: message.approval.change,
              decision: message.approval.decision ?? null,
              decidedAt: message.approval.decidedAt?.toISOString() ?? null,
            }
          : null,
        error: message.error ?? null,
        createdAt: new Date().toISOString(),
      };
    },

    async linkAgentRuns(_owner: PiggyConversationOwner, conversationId: string): Promise<void> {
      refuse('linkAgentRuns');
      linked.push(conversationId);
    },
  };

  return {
    store,
    appends,
    conversations,
    linked,
    seed(conversation): string {
      const id = randomUUID();
      conversations.set(id, {
        id,
        userId: conversation.userId,
        title: conversation.title ?? 'Seeded thread',
        readCapability: conversation.readCapability ?? 'book:read',
      });
      for (const message of conversation.messages ?? []) {
        appends.push({ conversationId: id, message });
      }
      return id;
    },
  };
}