dsh-superagent/tests/role-models.spec.ts

import { afterEach, describe, expect, it, vi } 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 '@deepseek-ai/dsh-agent'
import { Session, SessionId } from '@deepseek-ai/dsh-session'
import { ReasoningEffortId } from '@deepseek-ai/dsh-llm'
import { alignRoleModel, initializeRoleModel } from '../src/app/role-models.ts'
const contexts: Context[] = []
afterEach(async () => { await Promise.all(contexts.splice(0).map(ctx => ctx.fiber.dispose())); vi.restoreAllMocks() })
async function fixture() {
  const ctx = new Context(); contexts.push(ctx); await ctx.plugin(SystemPrompt)
  const resolve = vi.fn(async value => value)
  const flush = vi.fn(async () => {})
  ctx.reflect.provide('llm', { resolveCallConfig: resolve })
  ctx.reflect.provide('sessionPersistence', { flush })
  const session = Session.create(SessionId('role-model'))
  const selectModel = vi.fn(async (request: { provider: string; model: string; reasoningEffort?: unknown }) => ({
    selected: { provider: request.provider, model: request.model, ...(request.reasoningEffort === void 0 ? {} : { reasoningEffort: request.reasoningEffort }) } }))
  const sessionController: { agents?: { selectForNextRequest(agent: Agent, selection: ModelSelection): void }; selectModel: typeof selectModel } = {
    agents: { selectForNextRequest: (target, selection) => { target.session.append('model/selection', selection) } }, selectModel }
  ctx.reflect.provide('sessionController', sessionController)
  const agent = { ctx, session, id: 'role-model' } as unknown as Agent
  const host: ModelSelectionRef = { get current() { return session.requestHeader()?.config ?? { provider: 'host', model: 'default', reasoningEffort: ReasoningEffortId('high') } }, assembled: undefined }
  installModelSelection(ctx, host)
  const request = () => agentEvents(ctx, agent).waterfall('agent/request', { turn: 1, step: 1, signal: new AbortController().signal }, async () => ({ provider: 'seed', model: 'seed' }))
  const record = (selection: ModelSelection) => { const event = session.append('request/header', { header: { config: selection }, reason: 'initial' }); ctx.emit('session/event', session, event) }
  return { ctx, agent, host, resolve, flush, request, record, sessionController }
}
it('keeps prompt variables and initial routing on the role model without saving the host default', async () => {
  const { ctx, agent, request, record, flush } = await fixture()
  await initializeRoleModel(ctx, agent, { provider: 'worker-provider', model: 'worker-model' })
  expect((await ctx.systemPrompt.assemble()).variables).toMatchObject({ provider: 'worker-provider', model: 'worker-model' })
  const selected = await request()
  expect(selected).toEqual({ provider: 'worker-provider', model: 'worker-model' })
  expect(agent.session.snapshotEvents().some(event => event.type === 'model/selection')).toBe(true)
  expect(flush).toHaveBeenCalledOnce()
  record(selected)
  // A retry of the same assembled step keeps the same model even after the header is persisted.
  expect(await request()).toEqual(selected)
  await ctx.systemPrompt.assemble()
  expect(await request()).toEqual(selected)
})
it('lets a later explicit selection win without changing the in-flight assembled step', async () => {
  const { ctx, agent, host, request } = await fixture()
  await initializeRoleModel(ctx, agent, { provider: 'worker', model: 'one' })
  await ctx.systemPrompt.assemble()
  const manual = { provider: 'manual', model: 'two', reasoningEffort: ReasoningEffortId('low') }
  Object.defineProperty(host, 'current', { value: manual, configurable: true })
  const event = agent.session.append('model/selection', manual); ctx.emit('session/event', agent.session, event)
  expect(await request()).toEqual({ provider: 'worker', model: 'one' })
  await ctx.systemPrompt.assemble()
  expect(await request()).toEqual(manual)
})
it.each(['header', 'selection'])('preserves a resumed Session with existing %s', async kind => {
  const { ctx, agent, resolve } = await fixture()
  const existing = { provider: 'existing', model: 'keep' }
  if (kind === 'header') agent.session.append('request/header', { header: { config: existing }, reason: 'initial' })
  else agent.session.append('model/selection', existing)
  await initializeRoleModel(ctx, agent, { provider: 'worker', model: 'new' })
  expect(resolve).not.toHaveBeenCalled()
})
it('inherits the host when unset and fails unavailable routes before recording a selection', async () => {
  const { ctx, agent, resolve, request } = await fixture()
  await initializeRoleModel(ctx, agent, null)
  await ctx.systemPrompt.assemble()
  expect(await request()).toMatchObject({ provider: 'host', model: 'default' })
  resolve.mockRejectedValueOnce(new Error('unknown provider'))
  await expect(initializeRoleModel(ctx, agent, { provider: 'unknown', model: 'x' })).rejects.toThrow('unknown provider')
  expect(agent.session.snapshotEvents().some(event => event.type === 'model/selection')).toBe(false)
})
describe('alignRoleModel', () => {
  const selections = (agent: Agent) => agent.session.snapshotEvents().filter(event => event.type === 'model/selection')
  it('re-points a reused Session whose durable header drifted from the configured model', async () => {
    const { ctx, agent, flush, record } = await fixture()
    record({ provider: 'deepseek-official', model: 'deepseek-flash', reasoningEffort: ReasoningEffortId('low') })
    await alignRoleModel(ctx, agent, { provider: 'wp', model: 'wm' })
    const events = selections(agent)
    expect(events).toHaveLength(1)
    expect(events[0]?.data).toEqual({ provider: 'wp', model: 'wm' })
    expect(flush).toHaveBeenCalledOnce()
  })
  it('is an idempotent no-op when the durable selection already matches the configured model', async () => {
    const { ctx, agent, resolve, flush } = await fixture()
    agent.session.append('model/selection', { provider: 'wp', model: 'wm' })
    flush.mockClear()
    await alignRoleModel(ctx, agent, { provider: 'wp', model: 'wm' })
    expect(resolve).toHaveBeenCalledOnce()
    expect(selections(agent)).toHaveLength(1)
    expect(flush).not.toHaveBeenCalled()
  })
  it('returns without resolving or recording anything when no model is configured', async () => {
    const { ctx, agent, resolve, flush } = await fixture()
    await alignRoleModel(ctx, agent, null)
    expect(resolve).not.toHaveBeenCalled()
    expect(selections(agent)).toHaveLength(0)
    expect(flush).not.toHaveBeenCalled()
  })
  it('treats an equal configured reasoning effort as the same route', async () => {
    const { ctx, agent, flush } = await fixture()
    agent.session.append('model/selection', { provider: 'wp', model: 'wm', reasoningEffort: ReasoningEffortId('low') })
    await alignRoleModel(ctx, agent, { provider: 'wp', model: 'wm', reasoningEffort: 'low' })
    expect(selections(agent)).toHaveLength(1)
    expect(flush).not.toHaveBeenCalled()
  })
  it('re-points when only the reasoning effort differs and records the configured effort', async () => {
    const { ctx, agent, flush } = await fixture()
    agent.session.append('model/selection', { provider: 'wp', model: 'wm', reasoningEffort: ReasoningEffortId('low') })
    await alignRoleModel(ctx, agent, { provider: 'wp', model: 'wm', reasoningEffort: 'high' })
    const events = selections(agent)
    expect(events).toHaveLength(2)
    expect(events[1]?.data).toEqual({ provider: 'wp', model: 'wm', reasoningEffort: ReasoningEffortId('high') })
    expect(flush).toHaveBeenCalledOnce()
  })
  it('follows a selection that supersedes an older header instead of the header', async () => {
    const { ctx, agent, flush, record } = await fixture()
    record({ provider: 'old', model: 'header' })
    agent.session.append('model/selection', { provider: 'wp', model: 'wm' })
    flush.mockClear()
    await alignRoleModel(ctx, agent, { provider: 'wp', model: 'wm' })
    expect(selections(agent)).toHaveLength(1)
    expect(flush).not.toHaveBeenCalled()
  })
  it('treats a Session with no durable route as drift and seeds the configured model', async () => {
    const { ctx, agent, flush } = await fixture()
    await alignRoleModel(ctx, agent, { provider: 'wp', model: 'wm' })
    const events = selections(agent)
    expect(events).toHaveLength(1)
    expect(events[0]?.data).toEqual({ provider: 'wp', model: 'wm' })
    expect(flush).toHaveBeenCalledOnce()
  })
  it('prefers the internal Controller primitive over the default-saving selectModel', async () => {
    const { ctx, agent, sessionController } = await fixture()
    const primitive = vi.fn()
    sessionController.agents = { selectForNextRequest: primitive }
    await alignRoleModel(ctx, agent, { provider: 'wp', model: 'wm' })
    expect(primitive).toHaveBeenCalledWith(agent, { provider: 'wp', model: 'wm' })
    expect(sessionController.selectModel).not.toHaveBeenCalled()
  })
  it('falls back to selectModel when the internal Controller surface is absent', async () => {
    const { ctx, agent, sessionController } = await fixture()
    delete sessionController.agents
    await alignRoleModel(ctx, agent, { provider: 'wp', model: 'wm' })
    expect(sessionController.selectModel).toHaveBeenCalledWith(expect.objectContaining({ sessionId: 'role-model', provider: 'wp', model: 'wm' }))
  })
})