|
|
@@ -0,0 +1,212 @@
|
|
|
+import { afterEach, describe, expect, it } from 'vitest'
|
|
|
+import { Context } from 'cordis'
|
|
|
+import Loader from '@cordisjs/plugin-loader'
|
|
|
+import { agentEvents, type Agent } from '@deepseek-ai/dsh-agent'
|
|
|
+import LlmService, { CallId, type GenerateOptions, LlmAdapter, type StreamChunk } from '@deepseek-ai/dsh-llm'
|
|
|
+import SessionStore, { SessionId } from '@deepseek-ai/dsh-session'
|
|
|
+import type { SessionEvent, SessionHeader } from '@deepseek-ai/dsh-session'
|
|
|
+import SessionPersistence from '@deepseek-ai/dsh-session-persistence'
|
|
|
+import SystemPrompt from '@deepseek-ai/dsh-system-prompt'
|
|
|
+import ToolRegistry from '@deepseek-ai/dsh-tools'
|
|
|
+import * as checkpointPolicy from '../src/index.ts'
|
|
|
+
|
|
|
+const contexts: Context[] = []
|
|
|
+
|
|
|
+class TestPersistence extends SessionPersistence {
|
|
|
+ locate(_meta: SessionHeader): undefined { return undefined }
|
|
|
+ create(_meta: SessionHeader): Promise<void> { return Promise.resolve() }
|
|
|
+ append(_id: SessionId, _events: readonly SessionEvent[]): Promise<void> { return Promise.resolve() }
|
|
|
+ load(_id: SessionId): Promise<{ meta: SessionHeader; events: SessionEvent[] }> {
|
|
|
+ return Promise.reject(new Error('not used'))
|
|
|
+ }
|
|
|
+ list(): Promise<SessionHeader[]> { return Promise.resolve([]) }
|
|
|
+}
|
|
|
+
|
|
|
+class RecordingAdapter extends LlmAdapter {
|
|
|
+ constructor(private readonly order: string[]) { super() }
|
|
|
+ async * stream(_options: GenerateOptions): AsyncIterable<StreamChunk> {
|
|
|
+ this.order.push('adapter')
|
|
|
+ yield { type: 'finish', reason: { kind: 'stop' } }
|
|
|
+ }
|
|
|
+}
|
|
|
+
|
|
|
+async function setup(): Promise<Context> {
|
|
|
+ const ctx = new Context()
|
|
|
+ contexts.push(ctx)
|
|
|
+ await ctx.plugin(SessionStore)
|
|
|
+ await ctx.plugin(LlmService)
|
|
|
+ await ctx.plugin(SystemPrompt)
|
|
|
+ await ctx.plugin(ToolRegistry)
|
|
|
+ await ctx.plugin(TestPersistence)
|
|
|
+ await ctx.plugin(checkpointPolicy)
|
|
|
+ return ctx
|
|
|
+}
|
|
|
+
|
|
|
+async function drain(stream: AsyncIterable<StreamChunk>): Promise<void> {
|
|
|
+ for await (const _chunk of stream) { /* drain */ }
|
|
|
+}
|
|
|
+
|
|
|
+afterEach(async () => {
|
|
|
+ await Promise.all(contexts.splice(0).map(ctx => ctx.fiber.dispose()))
|
|
|
+})
|
|
|
+
|
|
|
+describe('session-checkpoint-policy request boundary', () => {
|
|
|
+ it('awaits the live session checkpoint before constructing the downstream model stream', async () => {
|
|
|
+ const ctx = await setup()
|
|
|
+ const session = ctx.sessions.create(SessionId('request-checkpoint'))
|
|
|
+ session.append('turn/start', { turn: 1, trigger: { kind: 'message', source: { kind: 'user' } } })
|
|
|
+ const gate = Promise.withResolvers<undefined>()
|
|
|
+ const order: string[] = []
|
|
|
+ ctx.on('session/flush', async () => {
|
|
|
+ order.push('flush:start')
|
|
|
+ await gate.promise
|
|
|
+ order.push('flush:end')
|
|
|
+ })
|
|
|
+ ctx.llm.registerAdapter(['mock'], new RecordingAdapter(order))
|
|
|
+
|
|
|
+ const pending = drain(ctx.llm.stream({
|
|
|
+ provider: 'mock', model: 'mock', messages: [], sessionId: session.id,
|
|
|
+ }))
|
|
|
+ await Promise.resolve()
|
|
|
+ expect(order).toEqual(['flush:start'])
|
|
|
+ gate.resolve(undefined)
|
|
|
+ await pending
|
|
|
+ expect(order).toEqual(['flush:start', 'flush:end', 'adapter'])
|
|
|
+ })
|
|
|
+
|
|
|
+ it('delegates a request without a live session without checkpointing', async () => {
|
|
|
+ const ctx = await setup()
|
|
|
+ const order: string[] = []
|
|
|
+ ctx.on('session/flush', () => { order.push('flush') })
|
|
|
+ ctx.llm.registerAdapter(['mock'], new RecordingAdapter(order))
|
|
|
+ await drain(ctx.llm.stream({ provider: 'mock', model: 'mock', messages: [] }))
|
|
|
+ expect(order).toEqual(['adapter'])
|
|
|
+ })
|
|
|
+
|
|
|
+ it('delegates an already-detached session id without checkpointing', async () => {
|
|
|
+ const ctx = await setup()
|
|
|
+ const order: string[] = []
|
|
|
+ ctx.on('session/flush', () => { order.push('flush') })
|
|
|
+ ctx.llm.registerAdapter(['mock'], new RecordingAdapter(order))
|
|
|
+ await drain(ctx.llm.stream({
|
|
|
+ provider: 'mock', model: 'mock', messages: [], sessionId: SessionId('detached'),
|
|
|
+ }))
|
|
|
+ expect(order).toEqual(['adapter'])
|
|
|
+ })
|
|
|
+
|
|
|
+ it('does not dispatch the adapter when the checkpoint rejects', async () => {
|
|
|
+ const ctx = await setup()
|
|
|
+ const session = ctx.sessions.create(SessionId('request-failure'))
|
|
|
+ const order: string[] = []
|
|
|
+ ctx.on('session/flush', () => Promise.reject(new Error('disk unavailable')))
|
|
|
+ ctx.llm.registerAdapter(['mock'], new RecordingAdapter(order))
|
|
|
+ await expect(drain(ctx.llm.stream({
|
|
|
+ provider: 'mock', model: 'mock', messages: [], sessionId: session.id,
|
|
|
+ }))).rejects.toThrow('disk unavailable')
|
|
|
+ expect(order).toEqual([])
|
|
|
+ })
|
|
|
+})
|
|
|
+
|
|
|
+describe('session-checkpoint-policy tool and step boundaries', () => {
|
|
|
+ it('awaits the checkpoint before a top-level tool body', async () => {
|
|
|
+ const ctx = await setup()
|
|
|
+ const session = ctx.sessions.create(SessionId('tool-checkpoint'))
|
|
|
+ const agent = { session } as Agent
|
|
|
+ const gate = Promise.withResolvers<undefined>()
|
|
|
+ const order: string[] = []
|
|
|
+ ctx.on('session/flush', async () => {
|
|
|
+ order.push('flush:start')
|
|
|
+ await gate.promise
|
|
|
+ order.push('flush:end')
|
|
|
+ })
|
|
|
+ ctx.tools.register({
|
|
|
+ name: 'write', description: 'side effect', parameters: {},
|
|
|
+ execute: async () => { order.push('tool'); return [] },
|
|
|
+ })
|
|
|
+
|
|
|
+ const pending = ctx.tools.execute({
|
|
|
+ callId: CallId('write-1'), name: 'write', arguments: {}, agent,
|
|
|
+ })
|
|
|
+ await Promise.resolve()
|
|
|
+ expect(order).toEqual(['flush:start'])
|
|
|
+ gate.resolve(undefined)
|
|
|
+ await expect(pending).resolves.toMatchObject({ isError: false })
|
|
|
+ expect(order).toEqual(['flush:start', 'flush:end', 'tool'])
|
|
|
+ })
|
|
|
+
|
|
|
+ it('turns a rejected checkpoint into an error result without running the tool body', async () => {
|
|
|
+ const ctx = await setup()
|
|
|
+ const session = ctx.sessions.create(SessionId('tool-failure'))
|
|
|
+ const agent = { session } as Agent
|
|
|
+ let ran = false
|
|
|
+ ctx.on('session/flush', () => Promise.reject(new Error('disk unavailable')))
|
|
|
+ ctx.tools.register({
|
|
|
+ name: 'write', description: 'side effect', parameters: {},
|
|
|
+ execute: async () => { ran = true; return [] },
|
|
|
+ })
|
|
|
+ const result = await ctx.tools.execute({
|
|
|
+ callId: CallId('write-2'), name: 'write', arguments: {}, agent,
|
|
|
+ })
|
|
|
+ expect(result.isError).toBe(true)
|
|
|
+ expect(result.content).toEqual([{ type: 'text', text: 'Error: disk unavailable' }])
|
|
|
+ expect(ran).toBe(false)
|
|
|
+ })
|
|
|
+
|
|
|
+ it('reuses the outer checkpoint for a nested tool dispatch', async () => {
|
|
|
+ const ctx = await setup()
|
|
|
+ const session = ctx.sessions.create(SessionId('nested-tool'))
|
|
|
+ const agent = { session } as Agent
|
|
|
+ let flushes = 0
|
|
|
+ ctx.on('session/flush', () => { flushes += 1 })
|
|
|
+ ctx.tools.register({ name: 'nested', description: 'nested', parameters: {}, execute: async () => [] })
|
|
|
+ await ctx.tools.execute({
|
|
|
+ callId: CallId('nested-1'), name: 'nested', arguments: {}, agent,
|
|
|
+ parent: Symbol('outer') as never,
|
|
|
+ })
|
|
|
+ expect(flushes).toBe(0)
|
|
|
+ })
|
|
|
+
|
|
|
+ it('checkpoints the complete recorded step at agent/post-step', async () => {
|
|
|
+ const ctx = await setup()
|
|
|
+ const session = ctx.sessions.create(SessionId('post-step'))
|
|
|
+ const agent = { session } as Agent
|
|
|
+ const flushed: string[] = []
|
|
|
+ ctx.on('session/flush', (current) => { flushed.push(current.id) })
|
|
|
+ await agentEvents(ctx, agent).serial(
|
|
|
+ 'agent/post-step', 1, 1, new AbortController().signal,
|
|
|
+ )
|
|
|
+ expect(flushed).toEqual([session.id])
|
|
|
+ })
|
|
|
+})
|
|
|
+
|
|
|
+describe('session-checkpoint-policy lifecycle', () => {
|
|
|
+ it('removes its wrappers when the owning fiber is disposed', async () => {
|
|
|
+ const ctx = new Context()
|
|
|
+ contexts.push(ctx)
|
|
|
+ await ctx.plugin(SessionStore)
|
|
|
+ await ctx.plugin(LlmService)
|
|
|
+ await ctx.plugin(SystemPrompt)
|
|
|
+ await ctx.plugin(ToolRegistry)
|
|
|
+ await ctx.plugin(TestPersistence)
|
|
|
+ const session = ctx.sessions.create(SessionId('disposed-policy'))
|
|
|
+ let flushes = 0
|
|
|
+ ctx.on('session/flush', () => { flushes += 1 })
|
|
|
+ ctx.llm.registerAdapter(['mock'], new RecordingAdapter([]))
|
|
|
+ const fiber = await ctx.plugin(checkpointPolicy)
|
|
|
+ await drain(ctx.llm.stream({ provider: 'mock', model: 'mock', messages: [], sessionId: session.id }))
|
|
|
+ expect(flushes).toBe(1)
|
|
|
+ await fiber.dispose()
|
|
|
+ await drain(ctx.llm.stream({ provider: 'mock', model: 'mock', messages: [], sessionId: session.id }))
|
|
|
+ expect(flushes).toBe(1)
|
|
|
+ })
|
|
|
+
|
|
|
+ it('keeps the Loader-safe namespace plugin shape', () => {
|
|
|
+ expect('default' in checkpointPolicy).toBe(false)
|
|
|
+ const loader = Object.create(Loader.prototype) as Loader
|
|
|
+ const unwrapped = loader.unwrapExports(checkpointPolicy) as Record<string, unknown>
|
|
|
+ expect(unwrapped).toBe(checkpointPolicy)
|
|
|
+ expect(unwrapped.name).toBe('session-checkpoint-policy')
|
|
|
+ expect(unwrapped.inject).toEqual(['llm', 'sessionPersistence', 'sessions', 'tools'])
|
|
|
+ expect(typeof unwrapped.apply).toBe('function')
|
|
|
+ })
|
|
|
+})
|