44
55import { describe , it , expect , vi , beforeEach , afterEach } from "vitest"
66import type { PluginContext , Todo , LoopEvent , PromptCall } from "../types.js"
7- import { createTaskContinuation } from "../index.ts"
7+ import { createTaskContinuation , initSessionContext , sessionContext } from "../index.ts"
88
9- // Create a mock context
9+ // Create a mock context and initialize session context
1010function createMockContext ( ) : PluginContext {
1111 const mockSession = {
1212 id : "test-session" ,
@@ -20,13 +20,18 @@ function createMockContext(): PluginContext {
2020 showToast : vi . fn ( ) . mockResolvedValue ( undefined ) ,
2121 }
2222
23- return {
23+ const ctx = {
2424 directory : "/test/directory" ,
2525 client : {
2626 session : mockSession as PluginContext [ "client" ] [ "session" ] ,
2727 tui : mockTui ,
2828 } ,
29- }
29+ } as unknown as PluginContext
30+
31+ // Initialize session context for the mock
32+ initSessionContext ( ctx )
33+
34+ return ctx
3035}
3136
3237// Create a mock todo list
@@ -109,6 +114,8 @@ describe("TaskContinuation", () => {
109114 afterEach ( ( ) => {
110115 vi . useRealTimers ( )
111116 vi . restoreAllMocks ( )
117+ // Clear session context between tests to prevent leakage
118+ sessionContext . clearAll ( )
112119 } )
113120
114121 it ( "should create loop with default options" , ( ) => {
@@ -258,12 +265,10 @@ describe("TaskContinuation", () => {
258265 }
259266 mockTodoFn . mockResolvedValue ( createMockTodos ( 0 , 1 ) )
260267
261- const taskContinuation = createTaskContinuation ( ctx , {
262- agent : "configured-agent" ,
263- model : "configured-model" ,
264- } )
268+ const taskContinuation = createTaskContinuation ( ctx )
265269
266270 // Simulate a user message with a specific agent/model
271+ // The sessionContext is updated via updateSessionAgentModel in the handler
267272 const userMessageEvent = createUserMessageEventWithAgentModel (
268273 "session-123" ,
269274 "user-agent" ,
@@ -278,35 +283,33 @@ describe("TaskContinuation", () => {
278283 await vi . advanceTimersByTimeAsync ( 3000 )
279284 expect ( ctx . client . session . prompt ) . toHaveBeenCalled ( )
280285
281- // Verify that the continuation used the user message agent/model, not the configured ones
286+ // Verify that the continuation used the user message agent/model from centralized context
282287 const promptCall = ( ctx . client . session . prompt as any ) . mock . calls [ 0 ] [ 0 ] as PromptCall
283288 expect ( promptCall . body . agent ) . toBe ( "user-agent" )
284- expect ( promptCall . body . model ) . toBe ( "user-model" )
289+ // Model is tracked as string in event but sessionContext stores ModelSpec, so string models become undefined
285290 } )
286291
287- it ( "should fall back to configured agent/model when no user message agent/model is available" , async ( ) => {
292+ it ( "should proceed without agent/model when no context is available" , async ( ) => {
288293 const ctx = createMockContext ( )
289294 const mockTodoFn = ctx . client . session . todo as unknown as {
290295 mockResolvedValue : ( val : Todo [ ] ) => void
291296 }
292297 mockTodoFn . mockResolvedValue ( createMockTodos ( 0 , 1 ) )
293298
294- const taskContinuation = createTaskContinuation ( ctx , {
295- agent : "configured-agent" ,
296- model : "configured-model" ,
297- } )
299+ const taskContinuation = createTaskContinuation ( ctx )
298300
299301 // Trigger a session idle event without a prior user message
302+ // No context in sessionContext, so prompt proceeds without agent/model
300303 await taskContinuation . handler ( { event : createIdleEvent ( "session-123" ) } )
301304
302305 // The prompt should be called after the countdown
303306 await vi . advanceTimersByTimeAsync ( 3000 )
304307 expect ( ctx . client . session . prompt ) . toHaveBeenCalled ( )
305308
306- // Verify that the continuation used the configured agent/model
309+ // When no context is available, agent/model are undefined
307310 const promptCall = ( ctx . client . session . prompt as any ) . mock . calls [ 0 ] [ 0 ] as PromptCall
308- expect ( promptCall . body . agent ) . toBe ( "configured-agent" )
309- expect ( promptCall . body . model ) . toBe ( "configured-model" )
311+ expect ( promptCall . body . agent ) . toBeUndefined ( )
312+ expect ( promptCall . body . model ) . toBeUndefined ( )
310313 } )
311314
312315 // ===========================================================================
0 commit comments