upstream-affine/packages/backend/server/src/__tests__/copilot/conversation-host.spec.ts

import '../../plugins/copilot/runtime/capability-runtime';

import ava from 'ava';

import {
  CopilotMessageNotFound,
  CopilotSessionNotFound,
  type Mutex,
} from '../../base';
import type { CompatSubmissionStore } from '../../plugins/copilot/compat/submission-store';
import type { ConversationPolicy } from '../../plugins/copilot/conversation/policy';
import type { Turn } from '../../plugins/copilot/core';
import type { AdmittedAttachmentSource } from '../../plugins/copilot/runtime/hosts/attachment-admission';
import { ConversationHost } from '../../plugins/copilot/runtime/hosts/conversation-host';
import {
  ChatSession,
  type ChatSessionService,
} from '../../plugins/copilot/session';

const test = ava;

function fixture(
  options: {
    failFirstAcceptedWrite?: boolean;
    failFirstAppend?: boolean;
    mode?: 'canonical' | 'personal';
    admittedAttachments?: AdmittedAttachmentSource[];
    revokeDuringAdmission?: boolean;
    promptAction?: string;
  } = {}
) {
  const sessionId = 'session-1';
  const token = 'submission-1';
  const durable = new Map<string, Turn>();
  const accepted = new Map<
    string,
    { userId: string; sessionId: string; turnId: string }
  >();
  const submissions = new Map([
    [
      token,
      {
        id: token,
        userId: 'user-1',
        workspaceId: 'workspace-1',
        sessionId,
        content: 'hello',
        attachments: options.admittedAttachments?.length
          ? ['data:text/plain;base64,aGVsbG8=']
          : [],
        params: {
          tone: 'brief',
          scopeSelectors: [{ kind: 'document', id: 'doc-2' }],
        },
        createdAt: new Date('2026-01-01T00:00:00.000Z'),
      },
    ],
  ]);
  let appendCount = 0;
  let quota = true;
  let accessAllowed = true;
  let acceptedWriteCount = 0;
  let artifactWrites = 0;
  const titleCalls: unknown[] = [];
  const chatSession = new ChatSession(
    {
      sessionId,
      userId: 'user-1',
      workspaceId: 'workspace-1',
      docId: 'doc-1',
      focus: { selectors: [] },
      prompt: {
        name: 'Chat With AFFiNE AI',
        action: options.promptAction,
        config: {},
        paramKeys: [],
        params: {},
      },
      turns: [],
    },
    () => []
  );
  const sessions = {
    getOwnedScope: async (id: string, userId: string) =>
      id === sessionId && userId === 'user-1'
        ? { workspaceId: 'workspace-1', docId: 'doc-1' }
        : undefined,
    getInScope: async (input: { sessionId: string; userId: string }) =>
      input.sessionId === sessionId && input.userId === 'user-1'
        ? chatSession
        : undefined,
    findTurnByCompatSubmissionId: async (
      _sessionId: string,
      _userId: string,
      _workspaceId: string,
      id: string
    ) => durable.get(id),
    appendTurn: async (input: { compatSubmissionId: string; turn: Turn }) => {
      appendCount += 1;
      if (options.failFirstAppend && appendCount === 1) {
        throw new Error('durable append failed');
      }
      const stored = { ...input.turn, id: `turn-${appendCount}` };
      durable.set(input.compatSubmissionId, stored);
      return stored;
    },
    getMessage: async (
      _sessionId: string,
      _userId: string,
      _workspaceId: string,
      turnId: string
    ) => [...durable.values()].find(turn => turn.id === turnId),
    revertLatestMessage: async () => {},
    generateSessionTitle: async (input: unknown) => {
      titleCalls.push(input);
    },
  } as unknown as ChatSessionService;
  const submissionStore = {
    get: async (id: string, userId: string) => {
      const value = submissions.get(id);
      return value?.userId === userId ? value : undefined;
    },
    getAccepted: async (id: string, userId: string) => {
      const value = accepted.get(id);
      return value?.userId === userId
        ? { ...value, acceptedAt: new Date('2026-01-01T00:00:00.000Z') }
        : undefined;
    },
    markAccepted: async (
      id: string,
      userId: string,
      value: { sessionId: string; turnId: string }
    ) => {
      acceptedWriteCount += 1;
      if (options.failFirstAcceptedWrite && acceptedWriteCount === 1) {
        throw new Error('accepted cache write failed');
      }
      accepted.set(id, { ...value, userId });
      submissions.delete(id);
    },
  } as unknown as CompatSubmissionStore;
  const mutex = {
    acquire: async () => ({ async [Symbol.asyncDispose]() {} }),
  } as unknown as Mutex;
  const policy = {
    hasQuota: async () => quota,
    shouldScheduleTitle: (prompt: { action?: string }) => !prompt.action,
  } as unknown as ConversationPolicy;
  const runtime = {
    putWorkspaceArtifact: async () => {
      artifactWrites++;
      return {
        id: 'artifact-1',
        canonicalMediaType: 'text/plain',
      };
    },
    compileTurnScope: async (input: {
      selectors: unknown[];
      preferredSourceIds?: string[];
    }) => ({
      version: 1,
      resolvedAt: '2026-01-01T00:00:00.000Z',
      selectors: input.selectors,
      requiredDocIds: [],
      requiredArtifactIds: [],
      preferredSourceIds: input.preferredSourceIds ?? [],
      retrieval: {
        mode: input.selectors.length ? 'required' : 'workspace',
        requiredDocIds: [],
        requiredArtifactIds: [],
        preferredSourceIds: input.preferredSourceIds ?? [],
      },
    }),
  };
  const attachmentAdmission = {
    admitPromptAttachments: async () => {
      if (options.revokeDuringAdmission) accessAllowed = false;
      return options.admittedAttachments ?? [];
    },
  };

  return {
    host: new ConversationHost(
      sessions,
      submissionStore,
      mutex,
      policy,
      runtime as never,
      attachmentAdmission as never,
      {
        sessionResource: async () => {
          if (!accessAllowed) throw new Error('permission denied');
          return options.mode ?? 'canonical';
        },
      } as never
    ),
    sessionId,
    token,
    durable,
    accepted,
    submissions,
    session: chatSession,
    appendCount: () => appendCount,
    artifactWrites: () => artifactWrites,
    titleCalls,
    setQuota: (value: boolean) => {
      quota = value;
    },
    revokeAccess: () => {
      accessAllowed = false;
    },
  };
}

test('compat submission becomes one durable user turn and replays idempotently', async t => {
  const state = fixture();

  const first = await state.host.prepareTurn('user-1', state.sessionId, {
    messageId: state.token,
  });
  t.is(first.latestTurn?.content, 'hello');
  t.deepEqual(first.latestTurn?.metadata, { tone: 'brief' });
  t.deepEqual(first.latestTurn?.scopeSnapshot?.selectors, [
    { kind: 'document', id: 'doc-2', source: 'draft' },
  ]);
  t.is(state.appendCount(), 1);
  t.false(state.submissions.has(state.token));
  t.truthy(state.accepted.get(state.token));

  state.setQuota(false);
  const replay = await state.host.prepareTurn('user-1', state.sessionId, {
    messageId: state.token,
  });
  t.is(replay.latestTurn?.id, first.latestTurn?.id);
  t.true(replay.quotaBackedRoutesAllowed);
  t.is(state.appendCount(), 1);
});

test('durable compat turn recovers after accepted-cache write failure', async t => {
  const state = fixture({ failFirstAcceptedWrite: true });

  await t.throwsAsync(
    state.host.prepareTurn('user-1', state.sessionId, {
      messageId: state.token,
    }),
    { message: 'accepted cache write failed' }
  );
  t.is(state.appendCount(), 1);
  t.truthy(state.durable.get(state.token));

  const recovered = await state.host.prepareTurn('user-1', state.sessionId, {
    messageId: state.token,
  });
  t.is(recovered.latestTurn?.id, state.durable.get(state.token)?.id);
  t.is(state.appendCount(), 1);
  t.truthy(state.accepted.get(state.token));
});

test('compat submission remains retryable when durable append fails', async t => {
  const state = fixture({ failFirstAppend: true });

  await t.throwsAsync(
    state.host.prepareTurn('user-1', state.sessionId, {
      messageId: state.token,
    }),
    { message: 'durable append failed' }
  );
  t.true(state.submissions.has(state.token));
  t.false(state.accepted.has(state.token));

  const recovered = await state.host.prepareTurn('user-1', state.sessionId, {
    messageId: state.token,
  });
  t.is(recovered.latestTurn?.content, 'hello');
  t.is(state.durable.size, 1);
});

test('compat submission cannot be consumed by another session', async t => {
  const state = fixture();
  const other = fixture();
  other.submissions.set(state.token, {
    id: state.token,
    userId: 'user-1',
    workspaceId: 'workspace-1',
    sessionId: 'session-other',
    content: 'secret',
    attachments: [],
    params: { tone: 'brief', scopeSelectors: [] },
    createdAt: new Date(),
  });

  await t.throwsAsync(
    other.host.prepareTurn('user-1', other.sessionId, {
      messageId: state.token,
    }),
    { instanceOf: CopilotMessageNotFound }
  );
  await t.throwsAsync(
    state.host.prepareTurn('user-2', state.sessionId, {
      messageId: state.token,
    }),
    { instanceOf: CopilotSessionNotFound }
  );
  t.is(other.appendCount(), 0);
});

test('direct conversation rejects revoked canonical access before durable append', async t => {
  const state = fixture();
  state.revokeAccess();

  await t.throwsAsync(
    state.host.prepareTurn('user-1', state.sessionId, {
      messageId: state.token,
    }),
    { message: 'permission denied' }
  );
  t.is(state.appendCount(), 0);
  t.true(state.submissions.has(state.token));
});

test('workspace context fails closed for personal scope and revoked canonical access', async t => {
  const personal = fixture({ mode: 'personal' });
  const personalError = await t.throwsAsync(
    personal.host.prepareTurn('user-1', personal.sessionId, {
      messageId: personal.token,
    })
  );

  const canonical = fixture({
    admittedAttachments: [
      {
        id: 'attachment-1',
        kind: 'bytes',
        mimeType: 'text/plain',
        size: 5,
        fileName: 'note.txt',
        hash: 'hash-1',
        data: 'aGVsbG8=',
        encoding: 'base64',
      },
    ],
    revokeDuringAdmission: true,
  });
  const canonicalError = await t.throwsAsync(
    canonical.host.prepareTurn('user-1', canonical.sessionId, {
      messageId: canonical.token,
    })
  );

  t.snapshot({
    personal: {
      error: personalError?.message,
      artifactWrites: personal.artifactWrites(),
      durableAppends: personal.appendCount(),
    },
    canonical: {
      error: canonicalError?.message,
      artifactWrites: canonical.artifactWrites(),
      durableAppends: canonical.appendCount(),
    },
  });
});

test('assistant persistence rechecks access and schedules eligible titles', async t => {
  const state = fixture();
  state.revokeAccess();

  await t.throwsAsync(
    state.host.persistAssistantTurn(
      state.session,
      {
        conversationId: state.session.config.sessionId,
        role: 'assistant',
        content: 'must not persist',
        attachments: [],
        metadata: {},
        renderTrace: [],
        toolEvents: [],
        createdAt: new Date(),
      },
      false
    ),
    { message: 'permission denied' }
  );
  t.is(state.appendCount(), 0);

  const eligible = fixture();
  const assistantTurn = {
    conversationId: eligible.sessionId,
    role: 'assistant' as const,
    content: 'answer',
    attachments: [],
    metadata: {},
    renderTrace: [],
    toolEvents: [],
    createdAt: new Date(),
  };
  await eligible.host.persistAssistantTurn(
    eligible.session,
    assistantTurn,
    false
  );
  await eligible.host.persistAssistantTurn(
    eligible.session,
    { ...assistantTurn, content: 'aborted' },
    true
  );

  const action = fixture({ promptAction: 'edit' });
  await action.host.persistAssistantTurn(
    action.session,
    { ...assistantTurn, conversationId: action.sessionId },
    false
  );

  t.deepEqual(eligible.titleCalls, [
    {
      sessionId: 'session-1',
      userId: 'user-1',
      workspaceId: 'workspace-1',
    },
  ]);
  t.deepEqual(action.titleCalls, []);
});