BoxBrain/tests/conversation.test.ts

import { afterEach, describe, expect, it } from 'vitest';
import type { Fact } from 'identitydb';
import {
  generateSchedule,
  getAvailabilitySnapshot,
  listConversationEntries,
  replyToConversation,
  setAvailabilityStatus,
  startConversation,
  type ConversationMemoryClassificationResult,
  type ConversationMemoryExtractionResult,
  type ConversationMemorySelectionResult,
  type ConversationTurnPlan,
  type ScheduleGenerationResult,
  type StructuredModelAdapter,
} from '../src';
import { persistFactDrafts } from '../src/memory';
import { closeOpenDbs, createDb } from './helpers';

afterEach(closeOpenDbs);

describe('conversation APIs', () => {
  it('blocks replies while the persona is offline', async () => {
    const db = await createDb();
    await db.upsertSpace({
      name: 'persona:minji',
      metadata: { boxbrain: { domain: 'persona.space', personaId: 'minji', displayName: 'Minji' } },
    });
    await setAvailabilityStatus(db, {
      spaceName: 'persona:minji',
      mode: 'offline',
      reason: 'Sleeping',
      effectiveFrom: '2026-05-12T00:00:00.000Z',
      until: '2026-05-12T08:00:00.000Z',
      sourceType: 'manual',
    });

    let responseCalls = 0;
    const result = await replyToConversation(db, {
      spaceName: 'persona:minji',
      counterpartId: 'user:shinwoo',
      counterpartDisplayName: 'Shinwoo',
      message: '지금뭐해',
      currentTime: '2026-05-12T07:00:00.000Z',
      mandatoryMemoryModel: createSelectionModel([]),
      contextualMemoryModel: createSelectionModel([]),
      responseModel: {
        provider: 'fake-structured',
        model: 'response-model',
        async generateObject<TOutput>(): Promise<TOutput> {
          responseCalls += 1;
          return { mode: 'reply', messages: ['안자'] } as TOutput;
        },
      },
    });

    expect(result.blocked).toBe(true);
    expect(result.messages).toEqual([]);
    expect(responseCalls).toBe(0);
  });

  it('retrieves mandatory/contextual memories, generates DM-style replies, and stores the turn', async () => {
    const db = await createDb();
    await seedPersonaMemory(db);
    const mandatoryPrompts: string[] = [];
    const contextualPrompts: string[] = [];
    const responsePrompts: string[] = [];

    const result = await replyToConversation(db, {
      spaceName: 'persona:minji',
      counterpartId: 'user:shinwoo',
      counterpartDisplayName: 'Shinwoo',
      message: '오늘저녁뭐해?',
      currentTime: '2026-05-12T12:00:00.000Z',
      mandatoryMemoryModel: createSelectionModel(['m1', 'm2'], mandatoryPrompts),
      contextualMemoryModel: createSelectionModel(['m3'], contextualPrompts),
      responseModel: createResponseModel(
        {
          mode: 'reply',
          messages: ['저녁엔가족이랑먹어', '왜궁금해'],
        },
        responsePrompts,
      ),
      rng: () => 0,
    });

    expect(mandatoryPrompts[0]).toContain('yesterday, today, and tomorrow schedules');
    expect(contextualPrompts[0]).toContain('오늘저녁뭐해?');
    expect(responsePrompts[0]).toContain('family dinner');
    expect(result.messages).toHaveLength(2);
    expect(result.messages[0]?.replyDelaySeconds).toBe(1);
    expect(result.messages[0]?.typingDelaySeconds).toBeGreaterThan(0);
    expect(result.messages[1]?.replyDelaySeconds).toBe(0);
    expect(result.usedMemories).toHaveLength(3);

    const history = await listConversationEntries(db, { spaceName: 'persona:minji', counterpartId: 'user:shinwoo' });
    expect(history.map((entry) => `${entry.direction}:${entry.text}`)).toEqual([
      'inbound:오늘저녁뭐해?',
      'outbound:저녁엔가족이랑먹어',
      'outbound:왜궁금해',
    ]);
  });

  it('can proactively start a conversation without inbound user text', async () => {
    const db = await createDb();
    await seedPersonaMemory(db);

    const result = await startConversation(db, {
      spaceName: 'persona:minji',
      counterpartId: 'user:shinwoo',
      counterpartDisplayName: 'Shinwoo',
      currentTime: '2026-05-12T15:00:00.000Z',
      mandatoryMemoryModel: createSelectionModel(['m1']),
      contextualMemoryModel: createSelectionModel([]),
      responseModel: createResponseModel({ mode: 'reply', messages: ['지금뭐해'] }),
      rng: () => 0,
    });

    expect(result.messages.map((message) => message.text)).toEqual(['지금뭐해']);
    const history = await listConversationEntries(db, { spaceName: 'persona:minji', counterpartId: 'user:shinwoo' });
    expect(history[0]?.proactive).toBe(true);
  });

  it('classifies inbound and outbound conversation messages, extracts only approved memories, and stores trace metadata', async () => {
    const db = await createDb();
    await seedPersonaMemory(db);
    const classifierPrompts: string[] = [];
    const extractorPrompts: string[] = [];

    await replyToConversation(db, {
      spaceName: 'persona:minji',
      counterpartId: 'user:shinwoo',
      counterpartDisplayName: 'Shinwoo',
      message: '이번주말에 등산 가고 싶어',
      currentTime: '2026-05-12T12:00:00.000Z',
      mandatoryMemoryModel: createSelectionModel(['m1']),
      contextualMemoryModel: createSelectionModel([]),
      responseModel: createResponseModel({
        mode: 'reply',
        messages: ['좋다 나도 산 좋아해', '저녁엔가족이랑먹어'],
      }),
      rng: () => 0,
      memoryPipeline: {
        classifierModel: createMemoryClassifier([
          { shouldRemember: true, reason: 'stores a durable user preference' },
          { shouldRemember: false, reason: 'small talk only' },
          { shouldRemember: true, reason: 'reveals a stable persona routine' },
        ], classifierPrompts),
        extractorModel: createMemoryExtractor([
          {
            facts: [
              {
                domain: 'persona.relationship',
                statement: 'Shinwoo wants to go hiking on weekends.',
                topics: [{ name: 'Shinwoo' }, { name: 'hiking' }],
              },
            ],
          },
          {
            facts: [
              {
                domain: 'persona.biography',
                statement: 'Minji often has family dinner in the evening.',
                topics: [{ name: 'Minji' }, { name: 'family dinner' }],
              },
            ],
          },
        ], extractorPrompts),
      },
    });

    expect(classifierPrompts).toHaveLength(3);
    expect(classifierPrompts[0]).toContain('Direction: inbound');
    expect(classifierPrompts[0]).toContain('이번주말에 등산 가고 싶어');
    expect(classifierPrompts[1]).toContain('Direction: outbound');
    expect(extractorPrompts).toHaveLength(2);
    expect(extractorPrompts[0]).toContain('stores a durable user preference');
    expect(extractorPrompts[1]).toContain('reveals a stable persona routine');

    const facts = await listFactsForSpace(db, 'persona:minji');
    const hikingFact = facts.find((fact) => fact.statement === 'Shinwoo wants to go hiking on weekends.');
    const dinnerFact = facts.find((fact) => fact.statement === 'Minji often has family dinner in the evening.');

    expect(hikingFact?.metadata).toMatchObject({
      boxbrain: {
        domain: 'persona.relationship',
      },
      conversationMemory: {
        turnId: expect.any(String),
        direction: 'inbound',
        counterpartId: 'user:shinwoo',
        counterpartDisplayName: 'Shinwoo',
        occurredAt: '2026-05-12T12:00:00.000Z',
        proactive: false,
        sourceMessage: '이번주말에 등산 가고 싶어',
        classifierReason: 'stores a durable user preference',
      },
    });
    expect(dinnerFact?.metadata).toMatchObject({
      boxbrain: {
        domain: 'persona.biography',
      },
      conversationMemory: {
        turnId: expect.any(String),
        direction: 'outbound',
        counterpartId: 'user:shinwoo',
        counterpartDisplayName: 'Shinwoo',
        occurredAt: '2026-05-12T12:00:00.000Z',
        proactive: false,
        sourceMessage: '저녁엔가족이랑먹어',
        classifierReason: 'reveals a stable persona routine',
      },
    });
  });

  it('deduplicates repeated extracted conversation memories by statement and domain', async () => {
    const db = await createDb();
    await seedPersonaMemory(db);

    await replyToConversation(db, {
      spaceName: 'persona:minji',
      counterpartId: 'user:shinwoo',
      counterpartDisplayName: 'Shinwoo',
      message: '나는 민트초코 좋아해',
      currentTime: '2026-05-12T12:00:00.000Z',
      mandatoryMemoryModel: createSelectionModel(['m1']),
      contextualMemoryModel: createSelectionModel([]),
      responseModel: createResponseModel({ mode: 'reply', messages: ['오 진짜?'] }),
      memoryPipeline: {
        classifierModel: createMemoryClassifier([
          { shouldRemember: true, reason: 'stable preference' },
          { shouldRemember: false, reason: 'reply is not worth storing' },
        ]),
        extractorModel: createMemoryExtractor([
          {
            facts: [
              {
                domain: 'persona.relationship',
                statement: 'Shinwoo likes mint chocolate.',
                topics: [{ name: 'Shinwoo' }, { name: 'mint chocolate' }],
              },
            ],
          },
        ]),
      },
    });

    await replyToConversation(db, {
      spaceName: 'persona:minji',
      counterpartId: 'user:shinwoo',
      counterpartDisplayName: 'Shinwoo',
      message: '나 아직도 민트초코 좋아해',
      currentTime: '2026-05-12T13:00:00.000Z',
      mandatoryMemoryModel: createSelectionModel(['m1']),
      contextualMemoryModel: createSelectionModel([]),
      responseModel: createResponseModel({ mode: 'reply', messages: ['기억하고있어'] }),
      memoryPipeline: {
        classifierModel: createMemoryClassifier([
          { shouldRemember: true, reason: 'same stable preference' },
          { shouldRemember: false, reason: 'reply is not worth storing' },
        ]),
        extractorModel: createMemoryExtractor([
          {
            facts: [
              {
                domain: 'persona.relationship',
                statement: 'Shinwoo likes mint chocolate.',
                topics: [{ name: 'Shinwoo' }, { name: 'mint chocolate' }],
              },
            ],
          },
        ]),
      },
    });

    const facts = await listFactsForSpace(db, 'persona:minji');
    expect(facts.filter((fact) => fact.statement === 'Shinwoo likes mint chocolate.')).toHaveLength(1);
  });

  it('executes availability tool calls after farewell-style refusal messages', async () => {
    const db = await createDb();
    await seedPersonaMemory(db);
    await generateSchedule(db, {
      spaceName: 'persona:minji',
      displayName: 'Minji',
      currentDate: '2026-05-12',
      scope: 'day',
      structuredModel: {
        provider: 'fake-structured',
        model: 'schedule-model',
        async generateObject<TOutput>(): Promise<TOutput> {
          return {
            events: [
              {
                title: 'Exam',
                startAt: '2026-05-12T10:00:00.000Z',
                endAt: '2026-05-12T12:00:00.000Z',
                availabilityMode: 'offline',
                availabilityReason: 'Exam starting',
                kind: 'special',
              },
            ],
          } satisfies ScheduleGenerationResult as TOutput;
        },
      },
    });

    const prompts: string[] = [];
    const result = await startConversation(db, {
      spaceName: 'persona:minji',
      counterpartId: 'user:shinwoo',
      counterpartDisplayName: 'Shinwoo',
      currentTime: '2026-05-12T09:55:00.000Z',
      mandatoryMemoryModel: createSelectionModel(['m1']),
      contextualMemoryModel: createSelectionModel([]),
      responseModel: createResponseModel(
        {
          mode: 'refuse',
          messages: ['나이제가봐야해', '이따연락해'],
          toolCalls: [
            {
              name: 'setAvailabilityStatus',
              arguments: {
                mode: 'offline',
                reason: 'Exam starting',
                effectiveFrom: '2026-05-12T09:58:00.000Z',
                until: '2026-05-12T12:00:00.000Z',
              },
            },
          ],
        },
        prompts,
      ),
    });

    expect(prompts[0]).toContain('next availability transition');
    expect(result.messages.map((message) => message.text)).toEqual(['나이제가봐야해', '이따연락해']);

    const snapshot = await getAvailabilitySnapshot(db, {
      spaceName: 'persona:minji',
      at: '2026-05-12T10:00:00.000Z',
    });
    expect(snapshot.current.mode).toBe('offline');
    expect(snapshot.current.reason).toBe('Exam starting');
  });
});

function createSelectionModel(memoryIds: string[], prompts: string[] = []): StructuredModelAdapter {
  return {
    provider: 'fake-structured',
    model: 'selection-model',
    async generateObject<TOutput>(request: { prompt: string }): Promise<TOutput> {
      prompts.push(request.prompt);
      return { memoryIds } satisfies ConversationMemorySelectionResult as TOutput;
    },
  };
}

function createResponseModel(plan: ConversationTurnPlan, prompts: string[] = []): StructuredModelAdapter {
  return {
    provider: 'fake-structured',
    model: 'response-model',
    async generateObject<TOutput>(request: { prompt: string }): Promise<TOutput> {
      prompts.push(request.prompt);
      return plan as TOutput;
    },
  };
}

function createMemoryClassifier(
  results: ConversationMemoryClassificationResult[],
  prompts: string[] = [],
): StructuredModelAdapter {
  const queue = [...results];
  return {
    provider: 'fake-structured',
    model: 'memory-classifier',
    async generateObject<TOutput>(request: { prompt: string }): Promise<TOutput> {
      prompts.push(request.prompt);
      const result = queue.shift();
      if (!result) {
        throw new Error('No queued conversation memory classification result.');
      }
      return result as TOutput;
    },
  };
}

function createMemoryExtractor(
  results: ConversationMemoryExtractionResult[],
  prompts: string[] = [],
): StructuredModelAdapter {
  const queue = [...results];
  return {
    provider: 'fake-structured',
    model: 'memory-extractor',
    async generateObject<TOutput>(request: { prompt: string }): Promise<TOutput> {
      prompts.push(request.prompt);
      const result = queue.shift();
      if (!result) {
        throw new Error('No queued conversation memory extraction result.');
      }
      return result as TOutput;
    },
  };
}

async function listFactsForSpace(db: Awaited<ReturnType<typeof createDb>>, spaceName: string): Promise<Fact[]> {
  const topics = await db.listTopics({ includeFacts: true, spaceName });
  const byId = new Map<string, Fact>();
  for (const topic of topics) {
    for (const fact of topic.facts) {
      byId.set(fact.id, fact);
    }
  }
  return Array.from(byId.values());
}

async function seedPersonaMemory(db: Awaited<ReturnType<typeof createDb>>) {
  await db.upsertSpace({
    name: 'persona:minji',
    metadata: { boxbrain: { domain: 'persona.space', personaId: 'minji', displayName: 'Minji' } },
  });

  await persistFactDrafts(db, {
    spaceName: 'persona:minji',
    domain: 'persona.biography',
    source: 'boxbrain:test',
    facts: [
      {
        statement: 'Minji prefers quiet evenings and close family dinners.',
        topics: [{ name: 'Minji' }, { name: 'family dinner' }],
      },
      {
        statement: 'Minji knows Shinwoo as a close contact she messages casually.',
        topics: [{ name: 'Minji' }, { name: 'Shinwoo' }],
      },
    ],
  });

  await generateSchedule(db, {
    spaceName: 'persona:minji',
    displayName: 'Minji',
    currentDate: '2026-05-12',
    scope: 'day',
    structuredModel: {
      provider: 'fake-structured',
      model: 'schedule-model',
      async generateObject<TOutput>(): Promise<TOutput> {
        return {
          events: [
            {
              title: 'Family dinner',
              description: 'Dinner at home with family.',
              startAt: '2026-05-12T18:00:00.000Z',
              endAt: '2026-05-12T20:00:00.000Z',
              availabilityMode: 'offline',
              availabilityReason: 'family dinner',
              kind: 'special',
            },
          ],
        } satisfies ScheduleGenerationResult as TOutput;
      },
    },
  });
}