@@ -7,6 +7,7 @@ import { deleteWorkspaceFixture, insertWorkspaceFixture } from '@sim/db/testing/
77 * duplicate-send path does, then answers the tool resume that follows.
88 */
99import { authMock , authMockFns } from '@sim/testing/mocks/auth.mock'
10+ import { envFlagsMock , resetEnvFlagsMock , setEnvFlags } from '@sim/testing/mocks/env-flags.mock'
1011import { afterAll , beforeAll , beforeEach , describe , expect , it , vi } from 'vitest'
1112
1213const { redisUrl, inheritedEnv, worker } = await vi . hoisted ( async ( ) => {
@@ -47,12 +48,13 @@ const { redisUrl, inheritedEnv, worker } = await vi.hoisted(async () => {
4748 REDIS_URL : process . env . REDIS_URL ,
4849 SIM_AGENT_API_URL : process . env . SIM_AGENT_API_URL ,
4950 }
50- process . env . REDIS_URL = url
51+ if ( url ) process . env . REDIS_URL = url
5152 process . env . SIM_AGENT_API_URL = `http://127.0.0.1:${ port } `
5253 return { redisUrl : url , inheritedEnv, worker : Object . assign ( worker , { server } ) }
5354} )
5455
5556vi . mock ( '@/lib/auth' , ( ) => authMock )
57+ vi . mock ( '@/lib/core/config/env-flags' , ( ) => envFlagsMock )
5658
5759import { db } from '@sim/db'
5860import {
@@ -74,20 +76,32 @@ import {
7476 EXHAUSTED_SIM_TOOL_MESSAGE ,
7577 SIM_TOOL_EXECUTION_MAX_ATTEMPTS ,
7678} from '@/lib/mothership/async-runs/execution-lease'
77- import { REVOKED_RUN_ERROR } from '@/lib/mothership/async-runs/orphaned-runs'
79+ import {
80+ ABANDONED_RUN_CURSOR_KEY ,
81+ findAbandonedRecoverableRuns ,
82+ REVOKED_RUN_ERROR ,
83+ } from '@/lib/mothership/async-runs/orphaned-runs'
7884import {
7985 claimToolExecution ,
8086 claimWorkflowToolExecution ,
87+ prepareWorkbenchAccess ,
88+ recordSimSandboxProcess ,
8189 revokeExhaustedSimToolExecutions ,
8290 upsertAsyncToolCall ,
8391} from '@/lib/mothership/async-runs/repository'
92+ import { createTrustedCopilotPrincipal } from '@/lib/mothership/auth/application-delegation'
8493import { chatPubSub } from '@/lib/mothership/chat-status'
8594import { takeOverAbandonedRuns } from '@/lib/mothership/request/application/abandoned-runs'
95+ import {
96+ RUN_TAKEOVER_AUDIENCE ,
97+ takeOverStreamRun ,
98+ } from '@/lib/mothership/request/application/recover-stream'
8699import {
87100 createProviderToolCallIdentity ,
88101 scopeProviderToolCallId ,
89102} from '@/lib/mothership/request/go/tool-call-identity'
90103import { chatStreamLockKey } from '@/lib/mothership/request/session/controller-lease'
104+ import { chatSandboxSessionKey } from '@/lib/mothership/tools/sandbox-session-key'
91105import { GET as streamGET } from '@/app/api/copilot/chat/stream/route'
92106
93107const ownerId = generateId ( )
@@ -119,12 +133,12 @@ function frameFor(streamId: string, chatId: string) {
119133 * chat, and the tool's execution lease lapsed with no result recorded.
120134 */
121135async function runAbandonedMidTool (
122- options : { attempts ?: number ; leaseLive ?: boolean ; toolName ?: string } = { }
136+ options : { attempts ?: number ; leaseLive ?: boolean ; toolName ?: string ; runId ?: string } = { }
123137) {
124138 const chatId = generateId ( )
125139 chatIds . push ( chatId )
126140 const streamId = generateId ( )
127- const runId = generateId ( )
141+ const runId = options . runId ?? generateId ( )
128142 const request = { message : 'Run the report' , userId, messageId : streamId , chatId, workspaceId }
129143 await db
130144 . insert ( copilotChats )
@@ -304,6 +318,8 @@ describe.runIf(Boolean(redisUrl))('taking over a run whose controller died', ()
304318
305319 beforeEach ( async ( ) => {
306320 worker . requests . length = 0
321+ resetEnvFlagsMock ( )
322+ await redis ( ) . del ( ABANDONED_RUN_CURSOR_KEY )
307323 await db
308324 . insert ( permissions )
309325 . values ( {
@@ -471,4 +487,105 @@ describe.runIf(Boolean(redisUrl))('taking over a run whose controller died', ()
471487 await db . update ( copilotRuns ) . set ( { status : 'cancelled' } ) . where ( eq ( copilotRuns . id , runId ) )
472488 }
473489 } )
490+
491+ it ( 'reaches an abandoned run behind live ones, examining a bounded window per scan' , async ( ) => {
492+ /** The lowest ids, so the scan meets the two live runs before the abandoned one. */
493+ const ids = [ 1 , 2 , 3 ] . map ( ( n ) => `00000000-0000-4000-8000-00000000000${ n } ` )
494+ const live = [
495+ await runAbandonedMidTool ( { runId : ids [ 0 ] } ) ,
496+ await runAbandonedMidTool ( { runId : ids [ 1 ] } ) ,
497+ ]
498+ const abandoned = await runAbandonedMidTool ( { runId : ids [ 2 ] } )
499+ for ( const run of live ) {
500+ await redis ( ) . set ( chatStreamLockKey ( run . chatId ) , `${ run . streamId } \nlive` , 'EX' , 60 )
501+ }
502+ try {
503+ const bound = { pageSize : 2 , maxExamined : 2 }
504+ expect ( await findAbandonedRecoverableRuns ( 1 , bound ) ) . toEqual ( [ ] )
505+ expect ( ( await findAbandonedRecoverableRuns ( 1 , bound ) ) . map ( ( run ) => run . id ) ) . toEqual ( [
506+ abandoned . runId ,
507+ ] )
508+ } finally {
509+ for ( const run of [ ...live , abandoned ] ) {
510+ await redis ( ) . del ( chatStreamLockKey ( run . chatId ) )
511+ await db
512+ . update ( copilotRuns )
513+ . set ( { status : 'cancelled' } )
514+ . where ( eq ( copilotRuns . id , run . runId ) )
515+ }
516+ }
517+ } )
518+
519+ it ( "hands a re-run the earlier attempt's workbench commands to stop first" , async ( ) => {
520+ const { chatId, runId, toolCallId } = await runAbandonedMidTool ( { leaseLive : true } )
521+ const sessionKey = chatSandboxSessionKey ( chatId )
522+ try {
523+ await recordSimSandboxProcess ( {
524+ toolCallId,
525+ runId,
526+ userId,
527+ ownerToken : 'dead-controller' ,
528+ process : { id : 'earlier-command' , sandboxId : 'sandbox-1' , sessionKey } ,
529+ } )
530+ await db
531+ . update ( copilotAsyncToolCalls )
532+ . set ( { executionLeaseExpiresAt : sql `clock_timestamp() - interval '1 second'` } )
533+ . where ( eq ( copilotAsyncToolCalls . toolCallId , toolCallId ) )
534+ const successor = { toolCallId, runId, userId, ownerToken : 'successor' }
535+ expect ( await claimToolExecution ( successor ) ) . toEqual ( { outcome : 'claimed' } )
536+ await recordSimSandboxProcess ( {
537+ ...successor ,
538+ process : { id : 'own-command' , sandboxId : 'sandbox-1' , sessionKey } ,
539+ } )
540+
541+ const state = await prepareWorkbenchAccess ( { ...successor , sessionKey } )
542+
543+ expect ( state . processes . map ( ( process ) => process . id ) ) . toEqual ( [ 'earlier-command' ] )
544+ } finally {
545+ await db . update ( copilotRuns ) . set ( { status : 'cancelled' } ) . where ( eq ( copilotRuns . id , runId ) )
546+ }
547+ } )
548+
549+ it ( 'claims no run once its caller is shutting down' , async ( ) => {
550+ const { chatId, streamId, runId, toolCallId } = await runAbandonedMidTool ( )
551+ const before = await storedRun ( runId )
552+ const shutdown = new AbortController ( )
553+ shutdown . abort ( )
554+ const principal = createTrustedCopilotPrincipal (
555+ { userId, workspaceId, chatId, delegationId : `test:${ runId } ` } ,
556+ { audience : RUN_TAKEOVER_AUDIENCE , ttlMs : 60_000 }
557+ )
558+ try {
559+ await expect (
560+ takeOverStreamRun . execute ( {
561+ principal,
562+ input : { streamId } ,
563+ request : { headers : new Headers ( ) , signal : shutdown . signal } ,
564+ } )
565+ ) . rejects . toThrow ( )
566+
567+ expect ( await storedRun ( runId ) ) . toEqual ( before )
568+ expect ( requestsFor ( streamId , '/api/mothership' ) ) . toHaveLength ( 0 )
569+ expect ( ( await storedTool ( toolCallId ) ) . attempts ) . toBe ( 1 )
570+ expect ( await redis ( ) . get ( chatStreamLockKey ( chatId ) ) ) . toBeNull ( )
571+ } finally {
572+ await db . update ( copilotRuns ) . set ( { status : 'cancelled' } ) . where ( eq ( copilotRuns . id , runId ) )
573+ }
574+ } )
575+
576+ it ( 'leaves a run recoverable when its takeover fails for a reason other than access' , async ( ) => {
577+ /** Hosted recovery refuses a run saved without its billing admission. */
578+ setEnvFlags ( { isHosted : true } )
579+ const { streamId, runId } = await runAbandonedMidTool ( )
580+ const before = await storedRun ( runId )
581+ try {
582+ await takeOverAbandonedRuns ( )
583+
584+ expect ( await storedRun ( runId ) ) . toEqual ( before )
585+ expect ( requestsFor ( streamId , '/api/streams/explicit-abort' ) ) . toHaveLength ( 0 )
586+ expect ( requestsFor ( streamId , '/api/mothership' ) ) . toHaveLength ( 0 )
587+ } finally {
588+ await db . update ( copilotRuns ) . set ( { status : 'cancelled' } ) . where ( eq ( copilotRuns . id , runId ) )
589+ }
590+ } )
474591} )
0 commit comments