packages/agent-core-v2/test/session/sessionTitle/agentTitlePromptSourceService.test.ts

import { afterEach, beforeEach, describe, expect, it } from 'vitest';

import { DisposableStore } from '#/_base/di/lifecycle';
import { createServices, type TestInstantiationService } from '#/_base/di/test';
import { IAgentContextMemoryService } from '#/agent/contextMemory/contextMemory';
import type { ContextMessage } from '#/agent/contextMemory/types';
import { IAgentLoopService, type PromptHandle } from '#/agent/loop/loop';
import type { ContentPart } from '#human/llm/message';
import { IAgentTitlePromptSource } from '#/session/sessionTitle/agentTitlePromptSource';
import { AgentTitlePromptSourceService } from '#/session/sessionTitle/agentTitlePromptSourceService';

const USER_ORIGIN: ContextMessage['origin'] = { kind: 'user' };

function userMessage(
  id: string,
  text: string,
  origin: ContextMessage['origin'] = USER_ORIGIN,
): ContextMessage {
  return {
    id,
    role: 'user',
    content: [{ type: 'text', text }],
    toolCalls: [],
    origin,
  };
}

function assistantMessage(id: string, parts: ContentPart[]): ContextMessage {
  return { id, role: 'assistant', content: parts, toolCalls: [] };
}

function toolMessage(id: string, text: string): ContextMessage {
  return { id, role: 'tool', content: [{ type: 'text', text }], toolCalls: [] };
}

interface MockQueueState {
  active?: PromptHandle;
  pending: {
    id: string;
    message: ContextMessage;
    tracked: true;
    createdAt: string;
    userMessageId: string;
  }[];
}

describe('AgentTitlePromptSource', () => {
  let disposables: DisposableStore;
  let ix: TestInstantiationService;
  let liveMessages: readonly ContextMessage[];
  let queue: MockQueueState;

  beforeEach(() => {
    liveMessages = [];
    queue = { active: undefined, pending: [] };
    disposables = new DisposableStore();
    ix = createServices(disposables, {
      additionalServices: (reg) => {
        reg.definePartialInstance(IAgentContextMemoryService, { get: () => liveMessages });
        reg.definePartialInstance(IAgentLoopService, {
          snapshot: () => ({
            state: 'idle' as const,
            activeTurnId: undefined,
            activePromptId: queue.active?.id,
            queue: queue.pending.map((item) => ({
              message: { role: 'user' as const, content: [...item.message.content] },
              meta: {
                promptId: item.id,
                tracked: item.tracked,
                createdAt: item.createdAt,
                userMessageId: item.userMessageId,
              },
            })),
            notificationCount: 0,
            paused: false,
            hasPendingRequests: queue.pending.length > 0,
            turn: undefined,
            activeTraceId: undefined,
          }),
          promptHandle: (id: string) => (queue.active?.id === id ? queue.active : undefined),
        });
        reg.define(IAgentTitlePromptSource, AgentTitlePromptSourceService);
      },
    });
  });

  afterEach(() => {
    disposables.dispose();
  });

  it('uses client display text consistently for titles, first turns and digests', async () => {
    liveMessages = [userMessage('annotated', '<browser_ref id="internal">Save</browser_ref>', { kind: 'user', clientMetadata: [{ display_text: 'Save button · Rename to Save changes' }] }), assistantMessage('reply', [{ type: 'text', text: 'Done' }])];
    const source = ix.get(IAgentTitlePromptSource);
    await expect(source.firstUserPrompts(3)).resolves.toEqual(['Save button · Rename to Save changes']);
    await expect(source.firstTurnExcerpt()).resolves.toEqual({ user: 'Save button · Rename to Save changes', assistant: 'Done' });
    await expect(source.digestExcerpt()).resolves.toEqual({ turns: [{ user: 'Save button · Rename to Save changes', assistant: 'Done' }] });
  });

  it('returns the first three prompts from the live context and queue in order', async () => {
    liveMessages = [userMessage('one', '第一条')];
    queue = {
      active: undefined,
      pending: [
        {
          id: 'two',
          userMessageId: 'two',
          createdAt: '2026-01-01T00:00:00.000Z',
          tracked: true,
          message: userMessage('two', '第二条'),
        },
        {
          id: 'three',
          userMessageId: 'three',
          createdAt: '2026-01-01T00:00:01.000Z',
          tracked: true,
          message: userMessage('three', '第三条'),
        },
      ],
    };

    await expect(ix.get(IAgentTitlePromptSource).firstUserPrompts(3)).resolves.toEqual([
      '第一条',
      '第二条',
      '第三条',
    ]);
  });

  it('keeps the head user messages of a compacted window, skipping elision and summary', async () => {
    liveMessages = [
      userMessage('head', '开场提问'),
      userMessage('elision', '... omitted ...', { kind: 'injection', variant: 'compaction_elision' }),
      userMessage('tail', '最近的追问'),
      userMessage('summary', ' compaction summary ', { kind: 'compaction_summary' }),
    ];

    await expect(ix.get(IAgentTitlePromptSource).firstUserPrompts(3)).resolves.toEqual([
      '开场提问',
      '最近的追问',
    ]);
  });

  it('returns no title prompts when history contains only slash activations', async () => {
    liveMessages = [
      userMessage('skill', 'expanded skill instructions', {
        kind: 'skill_activation',
        activationId: 'skill-1',
        skillName: 'compact',
        trigger: 'user-slash',
      }),
      userMessage('plugin', 'expanded plugin instructions', {
        kind: 'plugin_command',
        activationId: 'plugin-1',
        pluginId: 'example-plugin',
        commandName: 'run',
        trigger: 'user-slash',
      }),
    ];

    await expect(ix.get(IAgentTitlePromptSource).firstUserPrompts(3)).resolves.toEqual([]);
  });

  it('counts a queued prompt already appended to the context only once', async () => {
    liveMessages = [userMessage('one', '同一条')];
    queue = {
      active: {
        id: 'one',
        userMessageId: 'one',
        createdAt: '2026-01-01T00:00:00.000Z',
        state: 'running',
        message: userMessage('one', '同一条'),
        launched: Promise.resolve(undefined),
        completion: new Promise(() => {}),
      },
      pending: [],
    };

    await expect(ix.get(IAgentTitlePromptSource).firstUserPrompts(3)).resolves.toEqual(['同一条']);
  });

  it('firstTurnExcerpt pairs the opening prompt with the turn’s final assistant text', async () => {
    liveMessages = [
      userMessage('u1', '帮我写一个快排'),
      assistantMessage('a1-think', [{ type: 'think', think: '让我想想' }]),
      assistantMessage('a1-text', [{ type: 'text', text: '好的，先写一版' }]),
      toolMessage('t1', 'tool output'),
      assistantMessage('a2', [
        { type: 'text', text: '这是最终版实现' },
        { type: 'image_url', imageUrl: { url: 'data:image/png;base64,AAAA' } },
      ]),
      userMessage('u2', '再加个单测'),
      assistantMessage('a3', [{ type: 'text', text: '第二轮的回复' }]),
    ];

    await expect(ix.get(IAgentTitlePromptSource).firstTurnExcerpt()).resolves.toEqual({
      user: '帮我写一个快排',
      assistant: '这是最终版实现',
    });
  });

  it('firstTurnExcerpt reports a missing assistant reply until the turn ends', async () => {
    liveMessages = [userMessage('u1', '刚发的问题')];

    await expect(ix.get(IAgentTitlePromptSource).firstTurnExcerpt()).resolves.toEqual({
      user: '刚发的问题',
      assistant: undefined,
    });
  });

  it('digestExcerpt counts a queued prompt already appended to the context only once', async () => {
    liveMessages = [
      userMessage('one', '最早的问题'),
      assistantMessage('a1', [{ type: 'text', text: '第一轮回答' }]),
      userMessage('two', '进行中的问题'),
    ];
    queue = {
      active: {
        id: 'two',
        userMessageId: 'two',
        createdAt: '2026-01-01T00:00:01.000Z',
        state: 'running',
        message: userMessage('two', '进行中的问题'),
        launched: Promise.resolve(undefined),
        completion: new Promise(() => {}),
      },
      pending: [],
    };

    await expect(ix.get(IAgentTitlePromptSource).digestExcerpt()).resolves.toEqual({
      turns: [
        { user: '最早的问题', assistant: '第一轮回答' },
        { user: '进行中的问题', assistant: undefined },
      ],
    });
  });

  it('digestExcerpt pairs every prompt with its own turn’s final assistant text', async () => {
    liveMessages = [
      userMessage('u1', '最初的目标'),
      assistantMessage('a1', [{ type: 'text', text: '第一轮回答' }]),
      userMessage('u2', '中途追问'),
      assistantMessage('a2', [{ type: 'text', text: '中间回答' }]),
      userMessage('u3', '最近的要求'),
      assistantMessage('a3', [{ type: 'think', think: '思考中' }]),
      assistantMessage('a4', [{ type: 'text', text: '最新正文' }]),
    ];

    await expect(ix.get(IAgentTitlePromptSource).digestExcerpt()).resolves.toEqual({
      turns: [
        { user: '最初的目标', assistant: '第一轮回答' },
        { user: '中途追问', assistant: '中间回答' },
        { user: '最近的要求', assistant: '最新正文' },
      ],
    });
  });

  it('digestExcerpt covers every turn, even with a dangling tool-only span', async () => {
    liveMessages = [
      userMessage('u1', '最初的目标'),
      assistantMessage('a1', [{ type: 'text', text: '第一轮回答' }]),
      userMessage('u2', '第二个话题'),
      assistantMessage('a2', [{ type: 'think', think: '只在思考' }]),
      userMessage('u3', '第三个话题'),
      assistantMessage('a3', [{ type: 'text', text: '第三轮回答' }]),
      userMessage('u4', '最新的话题'),
      assistantMessage('a4', [{ type: 'text', text: '最新回答' }]),
    ];

    await expect(ix.get(IAgentTitlePromptSource).digestExcerpt()).resolves.toEqual({
      turns: [
        { user: '最初的目标', assistant: '第一轮回答' },
        { user: '第二个话题', assistant: undefined },
        { user: '第三个话题', assistant: '第三轮回答' },
        { user: '最新的话题', assistant: '最新回答' },
      ],
    });
  });

  it('digestExcerpt keeps a single-prompt conversation and dangling questions', async () => {
    liveMessages = [
      userMessage('u1', '唯一的问题'),
      assistantMessage('a1', [{ type: 'text', text: '唯一的回答' }]),
      userMessage('u2', '还没得到回复的新问题'),
    ];

    await expect(ix.get(IAgentTitlePromptSource).digestExcerpt()).resolves.toEqual({
      turns: [
        { user: '唯一的问题', assistant: '唯一的回答' },
        { user: '还没得到回复的新问题', assistant: undefined },
      ],
    });

    liveMessages = [userMessage('u1', '唯一的问题')];
    await expect(ix.get(IAgentTitlePromptSource).digestExcerpt()).resolves.toEqual({
      turns: [{ user: '唯一的问题', assistant: undefined }],
    });

    liveMessages = [];
    await expect(ix.get(IAgentTitlePromptSource).digestExcerpt()).resolves.toEqual({ turns: [] });
  });
});