CedarCopilot/cedar-mail

/**
 * Which sources a `[N]` in an answer resolves against.
 *
 * The load-bearing case is the last one: web-source numbering restarts every turn, so a single
 * thread-wide list would relabel every badge already on screen the moment a second search ran.
 */
import { resolveCitationsByMessage } from '@/store/messages/resolveMessageCitations';
import type { Citation, Message } from '@/store/messages/MessageTypes';

const web = (index: number, host: string): Citation => ({
  kind: 'web',
  index,
  url: `https://${host}/a`,
  siteName: host,
});

const msg = (id: string, over: Partial<Message> = {}): Message =>
  ({ id, role: 'assistant', type: 'text', content: '', ...over }) as Message;

describe('resolveCitationsByMessage', () => {
  it('gives a message the sources that existed when it was written', () => {
    const sources = [web(1, 'ashbyhq.com'), web(2, 'indexventures.com')];
    const map = resolveCitationsByMessage([
      msg('u1', { role: 'user' }),
      msg('tr1', { role: 'bot', type: 'tool-result', citations: sources }),
      msg('a1'),
    ]);

    expect(map.get('a1')).toBe(sources);
  });

  it('leaves a message before any search with nothing to resolve', () => {
    const map = resolveCitationsByMessage([
      msg('a0'),
      msg('tr1', { role: 'bot', type: 'tool-result', citations: [web(1, 'ashbyhq.com')] }),
    ]);

    expect(map.get('a0')).toBeUndefined();
  });

  it('does not let a later turn renumber the badges in an earlier one', () => {
    const turnOne = [web(1, 'ashbyhq.com'), web(2, 'indexventures.com')];
    const turnTwo = [web(1, 'crunchbase.com')];
    const map = resolveCitationsByMessage([
      msg('tr1', { role: 'bot', type: 'tool-result', citations: turnOne }),
      msg('a1'),
      msg('u2', { role: 'user' }),
      msg('tr2', { role: 'bot', type: 'tool-result', citations: turnTwo }),
      msg('a2'),
    ]);

    // `[2]` in a1 is still indexventures — turn two's list only reaches turn two's answer.
    expect(map.get('a1')).toBe(turnOne);
    expect(map.get('a2')).toBe(turnTwo);
  });

  it('a later turn that ran no search inherits NOTHING', () => {
    const turnOne = [web(1, 'ashbyhq.com')];
    const map = resolveCitationsByMessage([
      msg('tr1', { role: 'bot', type: 'tool-result', citations: turnOne }),
      msg('a1'),
      msg('u2', { role: 'user' }),
      // No tool result this turn — the model just writes prose, and "option [1]" in it is
      // a list marker, not a source. Carrying turn one's list here made it a live pill
      // pointing at an unrelated page.
      msg('a2'),
    ]);

    expect(map.get('a1')).toBe(turnOne);
    expect(map.get('a2')).toBeUndefined();
  });

  it("never maps the user's own message — `[1]` in something a person typed is not a citation", () => {
    const map = resolveCitationsByMessage([
      msg('tr1', { role: 'bot', type: 'tool-result', citations: [web(1, 'ashbyhq.com')] }),
      msg('u2', { role: 'user' }),
    ]);

    expect(map.has('u2')).toBe(false);
  });

  it('a message carrying its own citations keeps them', () => {
    const own = [web(1, 'crunchbase.com')];
    const map = resolveCitationsByMessage([
      msg('tr1', { role: 'bot', type: 'tool-result', citations: [web(1, 'ashbyhq.com')] }),
      msg('cited', { type: 'cited-text', citations: own }),
    ]);

    expect(map.get('cited')).toBe(own);
  });
});