Kaynağa Gözat

fix handoff worker lifecycle

Alvin Unreal 2 ay önce
ebeveyn
işleme
8ad421d73d

+ 22 - 3
src/index.ts

@@ -43,6 +43,7 @@ import {
   createCouncilTool,
   createHandoffCommandManager,
   createHandoffSessionTool,
+  createHandoffState,
   createPresetManager,
   createReadSessionTool,
   createWebfetchTool,
@@ -146,6 +147,7 @@ const OhMyOpenCodeLite: Plugin = async (ctx) => {
     typeof createDisplayNameMentionRewriter
   >;
   let handoffCommandManager: ReturnType<typeof createHandoffCommandManager>;
+  let handoffState: ReturnType<typeof createHandoffState>;
 
   // Counters for post-init health check (set inside try, checked outside)
   let toolCount = 0;
@@ -316,7 +318,8 @@ const OhMyOpenCodeLite: Plugin = async (ctx) => {
     presetManager = createPresetManager(ctx, config);
     divoomManager = createDivoomManager(config.divoom);
 
-    handoffCommandManager = createHandoffCommandManager(ctx);
+    handoffState = createHandoffState();
+    handoffCommandManager = createHandoffCommandManager(ctx, handoffState);
 
     toolCount =
       Object.keys(councilTools).length +
@@ -393,8 +396,12 @@ const OhMyOpenCodeLite: Plugin = async (ctx) => {
       ...todoContinuationHook.tool,
       ast_grep_search,
       ast_grep_replace,
-      handoff_session: createHandoffSessionTool(ctx),
-      read_session: createReadSessionTool(ctx.client),
+      handoff_session: createHandoffSessionTool(
+        ctx,
+        handoffState,
+        depthTracker,
+      ),
+      read_session: createReadSessionTool(ctx.client, handoffState),
     },
 
     mcp: mcps,
@@ -808,6 +815,18 @@ const OhMyOpenCodeLite: Plugin = async (ctx) => {
         },
       );
 
+      handoffCommandManager.handleEvent(
+        input as {
+          event: {
+            type: string;
+            properties?: {
+              info?: { id?: string; parentID?: string };
+              sessionID?: string;
+            };
+          };
+        },
+      );
+
       if (
         event.type === 'permission.asked' ||
         event.type === 'question.asked'

+ 49 - 1
src/tools/handoff/command.test.ts

@@ -1,5 +1,6 @@
 import { describe, expect, test } from 'bun:test';
 import { createHandoffCommandManager } from './command';
+import { createHandoffState } from './state';
 
 function createContext() {
   return {
@@ -10,7 +11,10 @@ function createContext() {
 
 describe('createHandoffCommandManager', () => {
   test('registers the /handoff command', () => {
-    const manager = createHandoffCommandManager(createContext());
+    const manager = createHandoffCommandManager(
+      createContext(),
+      createHandoffState(),
+    );
     const config: Record<string, unknown> = {};
 
     manager.registerCommand(config);
@@ -20,4 +24,48 @@ describe('createHandoffCommandManager', () => {
     expect(commands.handoff.template).toContain('handoff_session');
     expect(commands.handoff.template).toContain('$ARGUMENTS');
   });
+
+  test('marks child sessions of handoff workers with the same source', () => {
+    const state = createHandoffState();
+    state.markSession('ses_worker', 'ses_source');
+    const manager = createHandoffCommandManager(createContext(), state);
+
+    manager.handleEvent({
+      event: {
+        type: 'session.created',
+        properties: { info: { id: 'ses_child', parentID: 'ses_worker' } },
+      },
+    });
+
+    expect(state.sourceFor('ses_child')).toBe('ses_source');
+  });
+
+  test('does not mark unrelated child sessions', () => {
+    const state = createHandoffState();
+    const manager = createHandoffCommandManager(createContext(), state);
+
+    manager.handleEvent({
+      event: {
+        type: 'session.created',
+        properties: { info: { id: 'ses_child', parentID: 'ses_parent' } },
+      },
+    });
+
+    expect(state.isHandoffSession('ses_child')).toBe(false);
+  });
+
+  test('unmarks deleted handoff sessions', () => {
+    const state = createHandoffState();
+    state.markSession('ses_worker', 'ses_source');
+    const manager = createHandoffCommandManager(createContext(), state);
+
+    manager.handleEvent({
+      event: {
+        type: 'session.deleted',
+        properties: { info: { id: 'ses_worker' } },
+      },
+    });
+
+    expect(state.isHandoffSession('ses_worker')).toBe(false);
+  });
 });

+ 25 - 0
src/tools/handoff/command.ts

@@ -6,6 +6,7 @@
  */
 
 import type { PluginInput } from '@opencode-ai/plugin';
+import type { HandoffState } from './state';
 
 const COMMAND_NAME = 'handoff';
 
@@ -30,6 +31,7 @@ Call handoff_session with the worker prompt and any clearly relevant files:
  */
 export function createHandoffCommandManager(
   _ctx: PluginInput,
+  state: HandoffState,
   _processedSessions?: Set<string>,
 ) {
   /**
@@ -52,6 +54,29 @@ export function createHandoffCommandManager(
 
   return {
     registerCommand,
+    handleEvent(input: {
+      event: {
+        type: string;
+        properties?: {
+          info?: { id?: string; parentID?: string };
+          sessionID?: string;
+        };
+      };
+    }): void {
+      if (input.event.type === 'session.created') {
+        const info = input.event.properties?.info;
+        if (!info?.id || !info.parentID) return;
+
+        const source = state.sourceFor(info.parentID);
+        if (source) state.markSession(info.id, source);
+        return;
+      }
+
+      if (input.event.type !== 'session.deleted') return;
+      const sessionID =
+        input.event.properties?.info?.id ?? input.event.properties?.sessionID;
+      if (sessionID) state.unmarkSession(sessionID);
+    },
   };
 }
 

+ 1 - 0
src/tools/handoff/index.ts

@@ -14,6 +14,7 @@ export {
   FILE_REGEX,
   parseFileReferences,
 } from './files';
+export { createHandoffState, type HandoffState } from './state';
 export {
   createHandoffSessionTool,
   createReadSessionTool,

+ 25 - 0
src/tools/handoff/state.ts

@@ -0,0 +1,25 @@
+export interface HandoffState {
+  markSession(sessionID: string, sourceSessionID: string): void;
+  unmarkSession(sessionID: string): void;
+  isHandoffSession(sessionID: string): boolean;
+  sourceFor(sessionID: string): string | undefined;
+}
+
+export function createHandoffState(): HandoffState {
+  const sourceBySession = new Map<string, string>();
+
+  return {
+    markSession(sessionID: string, sourceSessionID: string): void {
+      sourceBySession.set(sessionID, sourceSessionID);
+    },
+    unmarkSession(sessionID: string): void {
+      sourceBySession.delete(sessionID);
+    },
+    isHandoffSession(sessionID: string): boolean {
+      return sourceBySession.has(sessionID);
+    },
+    sourceFor(sessionID: string): string | undefined {
+      return sourceBySession.get(sessionID);
+    },
+  };
+}

+ 65 - 35
src/tools/handoff/tools.test.ts

@@ -2,6 +2,8 @@ import { describe, expect, mock, test } from 'bun:test';
 import * as fs from 'node:fs';
 import * as os from 'node:os';
 import * as path from 'node:path';
+import { SubagentDepthTracker } from '../../utils/subagent-depth';
+import { createHandoffState } from './state';
 import { createHandoffSessionTool, createReadSessionTool } from './tools';
 
 function makeTempDir() {
@@ -26,17 +28,22 @@ describe('handoff_session tool', () => {
         ],
       }));
       const sessionAbort = mock(async () => ({}));
-      const tool = createHandoffSessionTool({
-        directory,
-        client: {
-          session: {
-            abort: sessionAbort,
-            create: sessionCreate,
-            messages: sessionMessages,
-            prompt: sessionPrompt,
+      const state = createHandoffState();
+      const tool = createHandoffSessionTool(
+        {
+          directory,
+          client: {
+            session: {
+              abort: sessionAbort,
+              create: sessionCreate,
+              messages: sessionMessages,
+              prompt: sessionPrompt,
+            },
           },
-        },
-      } as any);
+        } as any,
+        state,
+        new SubagentDepthTracker(),
+      );
 
       const result = await tool.execute(
         { prompt: 'Continue implementation', files: ['src/index.ts'] },
@@ -93,30 +100,35 @@ describe('handoff_session tool', () => {
     const directory = makeTempDir();
     try {
       let nestedResult = '';
-      const tool = createHandoffSessionTool({
-        directory,
-        client: {
-          session: {
-            abort: mock(async () => ({})),
-            create: mock(async () => ({ data: { id: 'ses_handoff' } })),
-            messages: mock(async () => ({
-              data: [
-                {
-                  info: { role: 'assistant' },
-                  parts: [{ type: 'text', text: 'done' }],
-                },
-              ],
-            })),
-            prompt: mock(async () => {
-              nestedResult = String(
-                await tool.execute({ prompt: 'nested handoff' }, {
-                  sessionID: 'ses_handoff',
-                } as any),
-              );
-            }),
+      const state = createHandoffState();
+      const tool = createHandoffSessionTool(
+        {
+          directory,
+          client: {
+            session: {
+              abort: mock(async () => ({})),
+              create: mock(async () => ({ data: { id: 'ses_handoff' } })),
+              messages: mock(async () => ({
+                data: [
+                  {
+                    info: { role: 'assistant' },
+                    parts: [{ type: 'text', text: 'done' }],
+                  },
+                ],
+              })),
+              prompt: mock(async () => {
+                nestedResult = String(
+                  await tool.execute({ prompt: 'nested handoff' }, {
+                    sessionID: 'ses_handoff',
+                  } as any),
+                );
+              }),
+            },
           },
-        },
-      } as any);
+        } as any,
+        state,
+        new SubagentDepthTracker(),
+      );
 
       await tool.execute({ prompt: 'outer handoff' }, {
         sessionID: 'ses_old',
@@ -147,13 +159,31 @@ describe('read_session tool', () => {
         },
       ],
     }));
-    const tool = createReadSessionTool({ session: { messages } } as any);
+    const state = createHandoffState();
+    state.markSession('ses_worker', 'ses_old');
 
-    const result = await tool.execute({ sessionID: 'ses_old' }, {} as any);
+    const result = await createReadSessionTool(
+      { session: { messages } } as any,
+      state,
+    ).execute({ sessionID: 'ses_old' }, { sessionID: 'ses_worker' } as any);
 
     expect(result).toContain('## User');
     expect(result).toContain('Hi');
     expect(result).toContain('## Assistant');
     expect(result).toContain('[Tool: read] Read file');
   });
+
+  test('blocks reads outside the source session', async () => {
+    const state = createHandoffState();
+    state.markSession('ses_worker', 'ses_old');
+    const messages = mock(async () => ({ data: [] }));
+
+    const result = await createReadSessionTool(
+      { session: { messages } } as any,
+      state,
+    ).execute({ sessionID: 'ses_other' }, { sessionID: 'ses_worker' } as any);
+
+    expect(result).toContain('can only read the source session');
+    expect(messages).not.toHaveBeenCalled();
+  });
 });

+ 51 - 9
src/tools/handoff/tools.ts

@@ -9,7 +9,9 @@
 import type { PluginInput, ToolDefinition } from '@opencode-ai/plugin';
 import { tool } from '@opencode-ai/plugin';
 import { extractSessionResult, promptWithTimeout } from '../../utils/session';
+import type { SubagentDepthTracker } from '../../utils/subagent-depth';
 import { buildSyntheticFileParts, parseFileReferences } from './files';
+import type { HandoffState } from './state';
 
 export type OpencodeClient = PluginInput['client'];
 const HANDOFF_TIMEOUT_MS = 5 * 60 * 1000;
@@ -19,9 +21,12 @@ const HANDOFF_TIMEOUT_MS = 5 * 60 * 1000;
  *
  * Takes the OpenCode client as a dependency for TUI and session operations.
  */
-export function createHandoffSessionTool(ctx: PluginInput): ToolDefinition {
+export function createHandoffSessionTool(
+  ctx: PluginInput,
+  state: HandoffState,
+  depthTracker?: SubagentDepthTracker,
+): ToolDefinition {
   const client = ctx.client;
-  const activeHandoffSessions = new Set<string>();
 
   return tool({
     description:
@@ -45,9 +50,16 @@ export function createHandoffSessionTool(ctx: PluginInput): ToolDefinition {
         context && typeof context === 'object' && 'sessionID' in context
           ? (context as { sessionID: string }).sessionID
           : 'unknown';
-      if (activeHandoffSessions.has(sessionID)) {
+      if (state.isHandoffSession(sessionID)) {
         return 'Nested handoff is disabled: this session is already a handoff worker. Finish this worker and return its summary to the parent session instead.';
       }
+      if (
+        sessionID !== 'unknown' &&
+        depthTracker &&
+        depthTracker.getDepth(sessionID) + 1 > depthTracker.maxDepth
+      ) {
+        return `Handoff worker blocked: max subagent depth ${depthTracker.maxDepth} would be exceeded.`;
+      }
 
       const sessionReference = `Work on behalf of parent session ${sessionID}. When you lack specific information you can use read_session to get it.`;
       const files = new Set([
@@ -78,7 +90,18 @@ export function createHandoffSessionTool(ctx: PluginInput): ToolDefinition {
         if (!childSessionID) {
           throw new Error('Handoff worker session did not return an id');
         }
-        activeHandoffSessions.add(childSessionID);
+        if (sessionID !== 'unknown' && depthTracker) {
+          const registered = depthTracker.registerChild(
+            sessionID,
+            childSessionID,
+          );
+          if (!registered) {
+            throw new Error(
+              'Handoff worker blocked: max subagent depth exceeded',
+            );
+          }
+        }
+        state.markSession(childSessionID, sessionID);
 
         await promptWithTimeout(
           client,
@@ -118,10 +141,16 @@ export function createHandoffSessionTool(ctx: PluginInput): ToolDefinition {
         ].join('\n');
       } finally {
         if (childSessionID) {
-          activeHandoffSessions.delete(childSessionID);
-          client.session
-            .abort({ path: { id: childSessionID }, query: { directory } })
-            .catch(() => {});
+          try {
+            await client.session.abort({
+              path: { id: childSessionID },
+              query: { directory },
+            });
+            state.unmarkSession(childSessionID);
+          } catch {
+            // Keep the handoff marker if abort fails; session.deleted cleanup
+            // will remove it when OpenCode eventually deletes the session.
+          }
         }
       }
     },
@@ -204,7 +233,10 @@ function formatTranscript(
  *
  * Takes the OpenCode client as a dependency for session.messages() calls.
  */
-export function createReadSessionTool(client: OpencodeClient): ToolDefinition {
+export function createReadSessionTool(
+  client: OpencodeClient,
+  state: HandoffState,
+): ToolDefinition {
   return tool({
     description:
       "Read the conversation transcript from a previous session. Use this when you need specific information from the source session that wasn't included in the handoff summary.",
@@ -228,6 +260,16 @@ export function createReadSessionTool(client: OpencodeClient): ToolDefinition {
         typeof (context as { directory?: unknown }).directory === 'string'
           ? (context as { directory: string }).directory
           : undefined;
+      const callerSessionID =
+        context && typeof context === 'object' && 'sessionID' in context
+          ? (context as { sessionID?: string }).sessionID
+          : undefined;
+      if (!callerSessionID || !state.isHandoffSession(callerSessionID)) {
+        return 'read_session is only available from handoff worker sessions.';
+      }
+      if (state.sourceFor(callerSessionID) !== args.sessionID) {
+        return 'read_session can only read the source session for this handoff worker.';
+      }
 
       try {
         const response = (await client.session.messages({

+ 1 - 0
src/tools/index.ts

@@ -5,6 +5,7 @@ export type { HandoffCommandManager } from './handoff';
 export {
   createHandoffCommandManager,
   createHandoffSessionTool,
+  createHandoffState,
   createReadSessionTool,
 } from './handoff';
 export type { PresetManager } from './preset-manager';