deepseek-harness/packages/core/session/tests/fork.spec.ts

import { describe, expect, it } from 'vitest'
import { Context } from 'cordis'
import { CallId } from '@deepseek-ai/dsh-llm'
import SessionStore, { Session, SessionForkError, SessionId } from '@deepseek-ai/dsh-session'
import type { SessionEvent, TurnEndReason } from '@deepseek-ai/dsh-session'

declare module '@deepseek-ai/dsh-session' {
  interface SessionEventMap {
    'test/log-only': { value: string }
  }
}

async function setup(): Promise<{ ctx: Context; sessions: SessionStore }> {
  const ctx = new Context()
  await ctx.plugin(SessionStore)
  return { ctx, sessions: ctx.sessions }
}

function appendClosedTurn(
  session: Session,
  turn: number,
  text = `hello ${turn}`,
  reason: TurnEndReason = { kind: 'completed' },
): void {
  session.append('turn/start', { turn, trigger: { kind: 'message', source: { kind: 'user' } } })
  session.append('user/message', {
    content: [{ type: 'text', text }],
    source: { kind: 'user' },
  }, { surfaceOp: 'append' })
  session.append('turn/end', { turn, reason })
}

function appendOpenTurn(session: Session, turn: number): void {
  session.append('turn/start', { turn, trigger: { kind: 'message', source: { kind: 'user' } } })
  session.append('user/message', {
    content: [{ type: 'text', text: `open ${turn}` }],
    source: { kind: 'user' },
  }, { surfaceOp: 'append' })
}

function firstUserMessage(events: readonly SessionEvent[]): SessionEvent<'user/message'> {
  const event = events.find((e): e is SessionEvent<'user/message'> => e.type === 'user/message')
  if (event === undefined) throw new Error('missing user/message')
  return event
}

function lastSeq(session: Session): number {
  const event = session.events.at(-1)
  if (event === undefined) throw new Error('missing last event')
  return event.seq
}

describe('SessionStore.fork', () => {
  it('forks an empty live session as an empty child with lineage metadata', async () => {
    const { ctx, sessions } = await setup()
    const source = ctx.sessions.create(SessionId('empty-parent'), { meta: { cwd: '/workspace' } })

    const child = sessions.fork(source, undefined, SessionId('empty-child'))

    expect(child.events).toEqual([])
    expect(child.header).toMatchObject({
      id: SessionId('empty-child'),
      cwd: '/workspace',
      parentSession: SessionId('empty-parent'),
      seedLength: 0,
    })
  })

  it('forks the latest completed boundary by default into detached frozen seed events', async () => {
    const { ctx, sessions } = await setup()
    const source = ctx.sessions.create(SessionId('parent'), { meta: { cwd: '/workspace' } })
    appendClosedTurn(source, 1, 'hello')

    const child = sessions.fork(SessionId('parent'), undefined, SessionId('child'))

    expect(child.events).toEqual(source.events)
    expect(child.events).not.toBe(source.events)
    expect(child.events[1]).not.toBe(source.events[1])
    expect(() => {
      firstUserMessage(child.events).data.content[0] = { type: 'text', text: 'child mutation' }
    }).toThrow(TypeError)
    expect(firstUserMessage(source.events).data.content).toEqual([{ type: 'text', text: 'hello' }])
    expect(firstUserMessage(child.events).data.content).toEqual([{ type: 'text', text: 'hello' }])
    expect(child.header).toMatchObject({
      id: SessionId('child'),
      cwd: '/workspace',
      parentSession: SessionId('parent'),
      seedLength: source.events.length,
    })
  })

  it('includes stable log-only events appended after a closed turn', async () => {
    const { ctx, sessions } = await setup()
    const source = ctx.sessions.create(SessionId('log-only-parent'))
    appendClosedTurn(source, 1, 'hello')
    source.append('test/log-only', { value: 'after execution' })

    const child = sessions.fork(source, undefined, SessionId('log-only-child'))

    expect(child.events).toEqual(source.events)
    expect(child.events.at(-1)).toMatchObject({
      type: 'test/log-only',
      data: { value: 'after execution' },
    })
  })

  it('forks from an earlier turn boundary even when the source currently has an open tail', async () => {
    const { ctx, sessions } = await setup()
    const source = ctx.sessions.create(SessionId('parent'), { meta: { cwd: '/workspace' } })
    appendClosedTurn(source, 1, 'first')
    const firstBoundary = lastSeq(source)
    appendClosedTurn(source, 2, 'second')
    appendOpenTurn(source, 3)

    const child = sessions.fork(source, firstBoundary, SessionId('child-from-first'))

    expect(child.events).toEqual(source.events.slice(0, firstBoundary + 1))
    expect(child.header.seedLength).toBe(firstBoundary + 1)
    expect(child.deriveMessages()).toEqual([{ role: 'user', content: [{ type: 'text', text: 'first' }] }])
  })

  it('accepts every turn/end reason as an explicit fork boundary', async () => {
    const { ctx, sessions } = await setup()
    const reasons: TurnEndReason[] = [
      { kind: 'completed' },
      { kind: 'aborted' },
      { kind: 'error', step: 1, message: 'model failed', code: 'MODEL' },
      { kind: 'disposed' },
      { kind: 'max-tokens' },
      { kind: 'interrupted' },
    ]

    for (const reason of reasons) {
      const source = ctx.sessions.create(SessionId(`parent-${reason.kind}`))
      appendClosedTurn(source, 1, reason.kind, reason)

      const child = sessions.fork(source, lastSeq(source), SessionId(`child-${reason.kind}`))

      expect(child.events.at(-1)?.type).toBe('turn/end')
      expect(child.header.seedLength).toBe(source.events.length)
    }
  })

  it('rejects invalid boundaries before creating a child', async () => {
    const { ctx, sessions } = await setup()
    const empty = ctx.sessions.create(SessionId('empty'))
    expect(() => sessions.fork(empty, 0, SessionId('empty-child')))
      .toThrow(new SessionForkError('fork boundary 0 does not exist in session "empty" (last seq: none)', 'INVALID_BOUNDARY'))
    expect(ctx.sessions.get(SessionId('empty-child'))).toBeUndefined()

    const source = ctx.sessions.create(SessionId('parent'))
    appendClosedTurn(source, 1)
    expect(() => sessions.fork(source, -1, SessionId('negative')))
      .toThrow(/non-negative safe integer/)
    expect(() => sessions.fork(source, 0.5, SessionId('fraction')))
      .toThrow(/non-negative safe integer/)
    expect(() => sessions.fork(source, Number.MAX_SAFE_INTEGER + 1, SessionId('unsafe')))
      .toThrow(/non-negative safe integer/)
    expect(() => sessions.fork(source, source.seq, SessionId('past-end')))
      .toThrow(new SessionForkError(`fork boundary ${source.seq} does not exist in session "parent" (last seq: ${source.seq - 1})`, 'INVALID_BOUNDARY'))
  })

  it('rejects a corrupted live source whose array index no longer matches event seq', async () => {
    const { ctx, sessions } = await setup()
    const source = ctx.sessions.create(SessionId('corrupt-parent'))
    appendClosedTurn(source, 1)
    const mutableLog = (source as unknown as { log: SessionEvent[] }).log
    mutableLog[2] = { ...mutableLog[2]!, seq: 99 }

    expect(() => sessions.fork(source, 2, SessionId('corrupt-child')))
      .toThrow(new SessionForkError('fork boundary 2 does not match a contiguous event seq in session "corrupt-parent"', 'INVALID_BOUNDARY'))
    expect(ctx.sessions.get(SessionId('corrupt-child'))).toBeUndefined()
  })

  it('rejects an unknown live session id', async () => {
    const { sessions } = await setup()

    expect(() => sessions.fork(SessionId('missing')))
      .toThrow(new SessionForkError('session "missing" not found', 'SESSION_NOT_FOUND'))
  })

  it('rejects a detached Session object that is not live in ctx.sessions', async () => {
    const { sessions } = await setup()
    const detached = new Session(SessionId('detached'))

    expect(() => sessions.fork(detached))
      .toThrow(new SessionForkError('session "detached" not found', 'SESSION_NOT_FOUND'))
  })

  it('rejects a stale Session object whose id is live on a different instance', async () => {
    const { ctx, sessions } = await setup()
    ctx.sessions.create(SessionId('same-id'))
    const stale = new Session(SessionId('same-id'))

    expect(() => sessions.fork(stale))
      .toThrow(new SessionForkError('session "same-id" is not the live store instance', 'SESSION_NOT_LIVE'))
  })

  it('rejects selected slices whose boundary is inside an open turn', async () => {
    const { ctx, sessions } = await setup()
    const cases: [string, (session: Session) => number][] = [
      ['turn/start', (session) => {
        session.append('turn/start', { turn: 1, trigger: { kind: 'message', source: { kind: 'user' } } })
        return lastSeq(session)
      }],
      ['step/start', (session) => {
        session.append('turn/start', { turn: 1, trigger: { kind: 'message', source: { kind: 'user' } } })
        session.append('step/start', { turn: 1, step: 1 })
        return lastSeq(session)
      }],
      ['user/message', (session) => {
        session.append('turn/start', { turn: 1, trigger: { kind: 'message', source: { kind: 'user' } } })
        session.append('user/message', { content: [{ type: 'text', text: 'open' }], source: { kind: 'user' } }, { surfaceOp: 'append' })
        return lastSeq(session)
      }],
      ['assistant/message', (session) => {
        session.append('turn/start', { turn: 1, trigger: { kind: 'message', source: { kind: 'user' } } })
        session.append('step/start', { turn: 1, step: 1 })
        session.append('assistant/message', { provenance: { provider: 'mock', model: 'mock' }, turn: 1, step: 1, content: [{ type: 'text', text: 'partial' }] }, { surfaceOp: 'append' })
        return lastSeq(session)
      }],
      ['tool/call', (session) => {
        const callId = CallId('call-open')
        session.append('turn/start', { turn: 1, trigger: { kind: 'message', source: { kind: 'user' } } })
        session.append('step/start', { turn: 1, step: 1 })
        session.append('assistant/message', { provenance: { provider: 'mock', model: 'mock' },
          turn: 1,
          step: 1,
          content: [{ type: 'tool-call', id: callId, name: 'bash', arguments: '{}' }],
        }, { surfaceOp: 'append' })
        session.append('tool/call', { turn: 1, step: 1, callId, name: 'bash', arguments: '{}' })
        return lastSeq(session)
      }],
    ]

    for (const [lastType, build] of cases) {
      const source = ctx.sessions.create(SessionId(`open-${lastType}`))
      const boundary = build(source)

      expect(() => sessions.fork(source, boundary))
        .toThrow(new SessionForkError(`fork boundary ${boundary} in session "open-${lastType}" ends inside open turn 1`, 'OPEN_TURN'))
    }
  })

  it('rejects a child session id that is already live with a typed fork error', async () => {
    const { ctx, sessions } = await setup()
    const source = ctx.sessions.create(SessionId('parent'))
    appendClosedTurn(source, 1)
    ctx.sessions.create(SessionId('child'))

    expect(() => sessions.fork(source, undefined, SessionId('child')))
      .toThrow(new SessionForkError('session "child" already exists', 'SESSION_ALREADY_EXISTS'))
  })

  it('rejects a duplicate child session id before validating the boundary', async () => {
    const { ctx, sessions } = await setup()
    const source = ctx.sessions.create(SessionId('open-parent'))
    source.append('turn/start', { turn: 1, trigger: { kind: 'message', source: { kind: 'user' } } })
    ctx.sessions.create(SessionId('child'))

    expect(() => sessions.fork(source, undefined, SessionId('child')))
      .toThrow(new SessionForkError('session "child" already exists', 'SESSION_ALREADY_EXISTS'))
  })
})