deepseek-harness/packages/api/session-controller/tests/commands-queue-attachment.host.spec.ts

import { Context } from '@deepseek-ai/cordis'
import AgentRegistry, { agentEvents, Inbox } from '@deepseek-ai/dsh-agent'
import type { Agent, ModelSelectionRef } from '@deepseek-ai/dsh-agent'
import { AttachmentError, AttachmentId } from '@deepseek-ai/dsh-attachment'
import type { ImageAttachmentRef } from '@deepseek-ai/dsh-attachment'
import { createUserMessage, MessageId } from '@deepseek-ai/dsh-llm'
import SessionStore, { SessionId } from '@deepseek-ai/dsh-session'
import type { SessionEvent, SessionHeader } from '@deepseek-ai/dsh-session'
import SessionProjectionRegistry from '@deepseek-ai/dsh-session-projection'
import { describe, expect, it, vi } from 'vitest'
import { ApiSessionAgentController } from '../src/agent.ts'
import { SessionCommandController } from '../src/commands.ts'

async function commandHarness(): Promise<{
  ctx: Context
  controller: SessionCommandController
  agent: Agent
  inbox: Inbox
  steer: ReturnType<typeof vi.fn>
  cancel: ReturnType<typeof vi.fn>
}> {
  const ctx = new Context()
  await ctx.plugin(SessionStore)
  await ctx.plugin(SessionProjectionRegistry)
  await ctx.plugin(AgentRegistry)
  const session = ctx.sessions.create(SessionId('commands-session'), { meta: { cwd: '/workspace' } })
  const steer = vi.fn()
  const cancel = vi.fn()
  const agent = {
    id: session.id,
    session,
    inbox: undefined as never,
    status: 'running',
    ctx,
    steer,
    followup: vi.fn(),
    cancel,
  } as unknown as Agent
  Object.assign(agent, { inbox: new Inbox(ctx, agent.session, agentEvents(ctx, agent)) })
  ctx.agents.register(agent)
  ctx.provide('workspaceRegistry', { get: () => undefined, list: () => [] } as never)
  ctx.provide('agentDefaultModel', {
    currentSelection: () => ({ provider: 'fixture', model: 'fixture-model' }),
    saveSelection: () => Promise.resolve(),
  } as never)
  const selection: ModelSelectionRef = {
    current: { provider: 'fixture', model: 'fixture-model' },
    assembled: undefined,
  }
  const agents = {
    resolveAgent: () => Promise.resolve({ agent }),
    selectionFor: () => selection,
    serializeImageAdmission: <Value>(_agent: Agent, operation: () => Promise<Value>) => operation(),
    composeAgent: () => Promise.resolve({ setup: () => {} }),
  } as unknown as ApiSessionAgentController
  return { ctx, controller: new SessionCommandController(ctx, agents, '/workspace'), agent, inbox: agent.inbox, steer, cancel }
}

async function expectFailure(operation: Promise<unknown>, code: string): Promise<void> {
  await expect(operation).rejects.toMatchObject({ failure: { code } })
}

describe('Session queue commands', () => {
  it('edits, removes, steers, and rejects stale queue occurrences', async () => {
    const { ctx, controller, agent, inbox, steer, cancel } = await commandHarness()
    const queued = createUserMessage({ content: [{ type: 'text', text: 'queued' }], source: { kind: 'user' } })
    const nextStep = createUserMessage({ content: [{ type: 'text', text: 'step' }], source: { kind: 'user' } })
    inbox.append('next-turn', queued)
    inbox.append('next-step', nextStep)

    await expectFailure(Promise.resolve().then(() => controller.updateQueue({
      sessionId: agent.id,
      itemId: queued.id,
      action: {
        kind: 'edit',
        content: [{
          type: 'image',
          attachment: {
            attachmentId: AttachmentId('att-edit'), mediaType: 'image/png', bytes: 1, width: 1, height: 1,
          },
        }],
      },
    })), 'attachment-error')
    await expectFailure(Promise.resolve().then(() => controller.updateQueue({
      sessionId: SessionId('missing'), itemId: queued.id, action: { kind: 'remove' },
    })), 'queue-item-not-found')
    await expectFailure(Promise.resolve().then(() => controller.updateQueue({
      sessionId: agent.id, itemId: MessageId('missing'), action: { kind: 'remove' },
    })), 'queue-item-not-found')
    await expectFailure(Promise.resolve().then(() => controller.updateQueue({
      sessionId: agent.id, itemId: nextStep.id, action: { kind: 'steer' },
    })), 'steer-unavailable')

    Object.assign(agent, { status: 'idle' })
    await expectFailure(Promise.resolve().then(() => controller.updateQueue({
      sessionId: agent.id, itemId: queued.id, action: { kind: 'steer' },
    })), 'steer-unavailable')
    expect(controller.updateQueue({
      sessionId: agent.id,
      itemId: queued.id,
      action: { kind: 'edit', content: [{ type: 'text', text: 'edited' }] },
    })).toEqual({ accepted: true })
    expect(inbox.nextTurn[0]?.content).toEqual([{ type: 'text', text: 'edited' }])
    expect(controller.updateQueue({
      sessionId: agent.id, itemId: nextStep.id, action: { kind: 'remove' },
    })).toEqual({ accepted: true })

    Object.assign(agent, { status: 'running' })
    const steered = inbox.nextTurn[0]
    if (steered === undefined) throw new Error('missing edited queue item')
    expect(controller.updateQueue({
      sessionId: agent.id, itemId: steered.id, action: { kind: 'steer' },
    })).toEqual({ accepted: true })
    expect(steer).toHaveBeenCalledWith(steered)

    await expectFailure(Promise.resolve().then(() => controller.cancel({
      sessionId: SessionId('missing'),
    })), 'session-not-found')
    expect(controller.cancel({ sessionId: agent.id })).toEqual({ accepted: true })
    expect(cancel).toHaveBeenCalledWith({ kind: 'user' }, { keepInbox: true })
    await ctx.fiber.dispose()
  })
})

function imageRef(id: string): ImageAttachmentRef {
  return {
    attachmentId: AttachmentId(id),
    mediaType: 'image/png',
    bytes: 1,
    width: 1,
    height: 1,
  }
}

function event(type: string, seq: number, data: unknown): SessionEvent {
  return { type, seq, time: seq + 1, data } as SessionEvent
}

async function persistedController(
  events: SessionEvent[],
  readImage: (ref: ImageAttachmentRef) => Promise<{ ref: ImageAttachmentRef; data: Uint8Array }>,
): Promise<{ ctx: Context; controller: SessionCommandController; sessionId: SessionId }> {
  const ctx = new Context()
  await ctx.plugin(SessionStore)
  const sessionId = SessionId('cold-attachment')
  const meta: SessionHeader = { version: 0, id: sessionId, createdAt: 1, cwd: '/workspace' }
  ctx.provide('sessionPersistence', {
    list: () => Promise.resolve([meta]),
    inspect: () => Promise.resolve({ meta, events }),
  } as never)
  ctx.provide('attachments', { readImage } as never)
  const agents = { resolveAgent: vi.fn() } as unknown as ApiSessionAgentController
  return { ctx, controller: new SessionCommandController(ctx, agents, '/workspace'), sessionId }
}

describe('Session attachment authorization', () => {
  it('finds references in direct, message, inserted, nested, and streamed content', async () => {
    const nested = imageRef('nested')
    const message = imageRef('message')
    const inserted = imageRef('inserted')
    const streamed = imageRef('streamed')
    const events = [
      event('fixture/direct', 0, {
        content: [null, [], { type: 'tool-result', content: [{ type: 'text', text: 'none' }] }, {
          type: 'tool-result', content: [{ type: 'image', attachment: nested }],
        }],
      }),
      event('assistant/message', 1, { message: { content: [{ type: 'image', attachment: message }] } }),
      event('agent/inbox/spliced', 2, { inserted: [{ content: [{ type: 'image', attachment: inserted }] }] }),
      event('assistant/chunk', 3, {
        chunk: { type: 'block-end', block: { type: 'image', attachment: streamed } },
      }),
    ]
    const readImage = vi.fn((ref: ImageAttachmentRef) => Promise.resolve({ ref, data: Uint8Array.of(1) }))
    const { ctx, controller, sessionId } = await persistedController(events, readImage)

    for (const ref of [nested, message, inserted, streamed]) {
      await expect(controller.attachment({ sessionId, attachmentId: ref.attachmentId }))
        .resolves.toEqual({ attachment: ref, data: 'AQ==' })
    }
    expect(readImage).toHaveBeenCalledTimes(4)
    await ctx.fiber.dispose()
  })

  it('maps missing persistence identities and attachment backend failures', async () => {
    const noPersistence = new Context()
    await noPersistence.plugin(SessionStore)
    const noPersistenceController = new SessionCommandController(
      noPersistence,
      { resolveAgent: vi.fn() } as unknown as ApiSessionAgentController,
      '/workspace',
    )
    await expectFailure(noPersistenceController.attachment({
      sessionId: SessionId('missing'), attachmentId: AttachmentId('att'),
    }), 'internal')

    const missing = new Context()
    await missing.plugin(SessionStore)
    missing.provide('sessionPersistence', {
      list: () => Promise.resolve([]),
      inspect: vi.fn(),
    } as never)
    const missingController = new SessionCommandController(
      missing,
      { resolveAgent: vi.fn() } as unknown as ApiSessionAgentController,
      '/workspace',
    )
    await expectFailure(missingController.attachment({
      sessionId: SessionId('missing'), attachmentId: 'att' as never,
    }), 'session-not-found')

    for (const thrown of [
      new AttachmentError('stored image is unavailable', 'ATTACHMENT_NOT_FOUND'),
      new Error('backend offline'),
    ]) {
      const ref = imageRef(`failure-${thrown.name}`)
      const fixture = await persistedController(
        [event('fixture/content', 0, { content: [{ type: 'image', attachment: ref }] })],
        () => Promise.reject(thrown),
      )
      await expectFailure(fixture.controller.attachment({
        sessionId: fixture.sessionId,
        attachmentId: ref.attachmentId,
      }), thrown instanceof AttachmentError ? 'attachment-error' : 'internal')
      await fixture.ctx.fiber.dispose()
    }
  })
})