deepseek-harness/packages/core/agent/tests/model-selection.spec.ts

import { describe, expect, it } from 'vitest'
import { Context } from '@deepseek-ai/cordis'
import SystemPrompt from '@deepseek-ai/dsh-system-prompt'
import {
  agentEvents,
  installModelSelection,
  type Agent,
  type ModelSelection,
  type ModelSelectionRef,
} from '../src/index.ts'
import {
  createUserMessage,
  ReasoningEffortId,
  type LlmCallConfig,
  type UserMessage,
} from '@deepseek-ai/dsh-llm'
import { Session, SessionId } from '@deepseek-ai/dsh-session'

const SIGNAL = new AbortController().signal
const INPUT = createUserMessage({
  content: [{ type: 'text', text: 'continue' }],
  source: { kind: 'user' },
})

function createAgent(): Agent {
  return { session: Session.create(SessionId('model-selection')) } as Agent
}

function expectedNotice(from: string, to: string) {
  return {
    content: [{
      type: 'text',
      text: `[model changed: assistant turns above this point were generated by ${from}; the session continues with ${to}]`,
    }],
    source: { kind: 'plugin', plugin: 'model-selection', form: 'notice', summary: `${from} → ${to}` },
  }
}

async function switchHarness(current: ModelSelection, previous?: ModelSelection) {
  const ctx = new Context()
  await ctx.plugin(SystemPrompt)
  const selection: ModelSelectionRef = { current, assembled: undefined }
  const dispose = installModelSelection(ctx, selection)
  const agent = createAgent()
  if (previous !== undefined) {
    agent.session.append('request/header', { header: { config: previous }, reason: 'initial' })
  }
  await ctx.systemPrompt.assemble()
  return { agent, ctx, dispose, selection }
}

async function preStep(
  ctx: Context,
  agent: Agent,
  {
    messages = [INPUT],
    offered = [INPUT],
    step = 1,
    signal = SIGNAL,
  }: {
    messages?: UserMessage[]
    offered?: UserMessage[]
    step?: number
    signal?: AbortSignal
  } = {},
) {
  return agentEvents(ctx, agent).waterfall(
    'agent/pre-step',
    { turn: 1, step, messages: offered, signal },
    () => Promise.resolve({ kind: 'enter' as const, messages }),
  )
}

describe('installModelSelection()', () => {
  it('snapshots prompt variables and request routing together, then disposes its listeners', async () => {
    const ctx = new Context()
    await ctx.plugin(SystemPrompt)
    const selection: ModelSelectionRef = { current: undefined, assembled: undefined }
    const dispose = installModelSelection(ctx, selection)
    const agent = createAgent()
    const seed: LlmCallConfig = { provider: 'seed', model: 'seed', temperature: 0.2 }
    const signal = new AbortController().signal

    expect((await ctx.systemPrompt.assemble()).variables).toEqual({})
    await expect(agentEvents(ctx, agent).waterfall(
      'agent/request', { turn: 1, step: 0, signal }, () => Promise.resolve(seed),
    )).resolves.toBe(seed)

    selection.current = {
      provider: 'alpha',
      model: 'a1',
      reasoningEffort: ReasoningEffortId('high'),
    }
    expect((await ctx.systemPrompt.assemble()).variables).toMatchObject({ provider: 'alpha', model: 'a1' })
    selection.current = { provider: 'beta', model: 'b1' }
    await expect(agentEvents(ctx, agent).waterfall(
      'agent/request', { turn: 1, step: 0, signal }, () => Promise.resolve(seed),
    )).resolves.toEqual({
      provider: 'alpha',
      model: 'a1',
      reasoningEffort: ReasoningEffortId('high'),
      temperature: 0.2,
    })

    expect((await ctx.systemPrompt.assemble()).variables).toMatchObject({ provider: 'beta', model: 'b1' })
    const inherited: LlmCallConfig = {
      provider: 'alpha',
      model: 'a1',
      reasoningEffort: ReasoningEffortId('max'),
      temperature: 0.2,
    }
    await expect(agentEvents(ctx, agent).waterfall(
      'agent/request', { turn: 1, step: 1, signal }, () => Promise.resolve(inherited),
    )).resolves.toEqual({ provider: 'beta', model: 'b1', temperature: 0.2 })

    dispose()
    expect((await ctx.systemPrompt.assemble()).variables).toEqual({})
    await expect(agentEvents(ctx, agent).waterfall(
      'agent/request', { turn: 2, step: 0, signal }, () => Promise.resolve(seed),
    )).resolves.toBe(seed)
    await ctx.fiber.dispose()
  })

  it('announces same-provider and cross-provider route changes from the assembled selection', async () => {
    const { agent, ctx, dispose, selection } = await switchHarness(
      { provider: 'alpha', model: 'a1' },
      { provider: 'alpha', model: 'a0' },
    )
    await expect(preStep(ctx, agent)).resolves.toMatchObject({
      messages: [INPUT, expectedNotice('a0', 'a1')],
    })

    selection.current = { provider: 'beta', model: 'b1' }
    await ctx.systemPrompt.assemble()
    selection.current = { provider: 'alpha', model: 'a2' }
    await expect(preStep(ctx, agent)).resolves.toMatchObject({
      messages: [INPUT, expectedNotice('alpha/a0', 'beta/b1')],
    })
    await ctx.systemPrompt.assemble()
    await expect(preStep(ctx, agent)).resolves.toMatchObject({
      messages: [INPUT, expectedNotice('a0', 'a2')],
    })

    dispose()
    await ctx.fiber.dispose()
  })

  it('does not announce initial, same-route, effort-only, rejected, aborted, or disposed steps', async () => {
    const { agent, ctx, dispose, selection } = await switchHarness({ provider: 'alpha', model: 'a0' })
    await expect(preStep(ctx, agent)).resolves.toMatchObject({ kind: 'enter', messages: [INPUT] })
    agent.session.append('request/header', {
      header: { config: { provider: 'alpha', model: 'a0' } }, reason: 'initial',
    })
    selection.current = {
      provider: 'alpha',
      model: 'a0',
      reasoningEffort: ReasoningEffortId('high'),
    }
    await ctx.systemPrompt.assemble()
    await expect(preStep(ctx, agent)).resolves.toMatchObject({ kind: 'enter', messages: [INPUT] })

    selection.current = { provider: 'alpha', model: 'a1' }
    await ctx.systemPrompt.assemble()
    const rejected = await agentEvents(ctx, agent).waterfall(
      'agent/pre-step',
      { turn: 1, step: 1, messages: [], signal: SIGNAL },
      () => Promise.resolve({ kind: 'reject' as const }),
    )
    expect(rejected).toEqual({ kind: 'reject' })
    const aborted = new AbortController()
    aborted.abort()
    await expect(preStep(ctx, agent, { signal: aborted.signal })).resolves.toMatchObject({ messages: [INPUT] })

    dispose()
    await expect(preStep(ctx, agent)).resolves.toMatchObject({ kind: 'enter', messages: [INPUT] })
    await ctx.fiber.dispose()
  })

  it('preserves empty no-call decisions and announces an empty tool continuation', async () => {
    const { agent, ctx } = await switchHarness(
      { provider: 'alpha', model: 'a1' },
      { provider: 'alpha', model: 'a0' },
    )
    await expect(preStep(ctx, agent, { messages: [] })).resolves.toEqual({ kind: 'enter', messages: [] })
    await expect(preStep(ctx, agent, { messages: [], step: 2 })).resolves.toEqual({ kind: 'enter', messages: [] })
    await expect(preStep(ctx, agent, { messages: [], offered: [], step: 2 })).resolves.toMatchObject({
      messages: [{ source: { summary: 'a0 → a1' } }],
    })
    await ctx.fiber.dispose()
  })
})