Browse Source

fix: prune unmanaged session read context

Alvin Unreal 3 months ago
parent
commit
e88ba4f7da

+ 110 - 0
src/hooks/task-session-manager/index.test.ts

@@ -111,6 +111,13 @@ describe('task-session-manager hook', () => {
   test('tracks files read by child sessions in resumable prompt context', async () => {
     const { hook } = createHook();
 
+    await hook.event({
+      event: {
+        type: 'session.created',
+        properties: { info: { id: 'child-1', parentID: 'parent-1' } },
+      },
+    });
+
     await hook['tool.execute.after'](
       {
         tool: 'read',
@@ -171,6 +178,13 @@ describe('task-session-manager hook', () => {
   test('accumulates multiple reads and hides tiny read context', async () => {
     const { hook } = createHook();
 
+    await hook.event({
+      event: {
+        type: 'session.created',
+        properties: { info: { id: 'child-1', parentID: 'parent-1' } },
+      },
+    });
+
     await hook['tool.execute.after'](
       { tool: 'read', sessionID: 'child-1', callID: 'read-1' },
       {
@@ -228,6 +242,102 @@ describe('task-session-manager hook', () => {
     expect(prompt).toContain('src/large.ts (12 lines)');
   });
 
+  test('ignores reads from unmanaged child sessions', async () => {
+    const { hook } = createHook({
+      shouldManageSession: (sessionID) => sessionID === 'parent-1',
+    });
+
+    await hook.event({
+      event: {
+        type: 'session.created',
+        properties: { info: { id: 'child-1', parentID: 'other-parent' } },
+      },
+    });
+    await hook['tool.execute.after'](
+      { tool: 'read', sessionID: 'child-1', callID: 'read-1' },
+      {
+        output: [
+          '<path>/tmp/src/index.ts</path>',
+          '<content>',
+          ...Array.from({ length: 12 }, (_, index) => `${index + 1}: line`),
+          '</content>',
+        ].join('\n'),
+      },
+    );
+
+    await hook['tool.execute.before'](
+      { tool: 'task', sessionID: 'parent-1', callID: 'call-1' },
+      { args: { subagent_type: 'explorer', description: 'unmanaged read' } },
+    );
+    await hook['tool.execute.after'](
+      { tool: 'task', sessionID: 'parent-1', callID: 'call-1' },
+      {
+        output:
+          'task_id: child-1 (for resuming to continue this task if needed)',
+      },
+    );
+
+    const system = { system: ['base'] };
+    await hook['experimental.chat.system.transform'](
+      { sessionID: 'parent-1' },
+      system,
+    );
+
+    const prompt = system.system.join('\n');
+    expect(prompt).toContain('exp-1 unmanaged read');
+    expect(prompt).not.toContain('Context read by exp-1');
+  });
+
+  test('prunes read context when remembered sessions are evicted', async () => {
+    const { hook } = createHook();
+
+    for (const index of [1, 2, 3]) {
+      await hook.event({
+        event: {
+          type: 'session.created',
+          properties: {
+            info: { id: `child-${index}`, parentID: 'parent-1' },
+          },
+        },
+      });
+      await hook['tool.execute.after'](
+        { tool: 'read', sessionID: `child-${index}`, callID: `read-${index}` },
+        {
+          output: [
+            `<path>/tmp/src/file-${index}.ts</path>`,
+            '<content>',
+            ...Array.from({ length: 12 }, (_, line) => `${line + 1}: line`),
+            '</content>',
+          ].join('\n'),
+        },
+      );
+      await hook['tool.execute.before'](
+        { tool: 'task', sessionID: 'parent-1', callID: `call-${index}` },
+        { args: { subagent_type: 'explorer', description: `thread ${index}` } },
+      );
+      await hook['tool.execute.after'](
+        { tool: 'task', sessionID: 'parent-1', callID: `call-${index}` },
+        {
+          output: `task_id: child-${index} (for resuming to continue this task if needed)`,
+        },
+      );
+    }
+
+    const system = { system: ['base'] };
+    await hook['experimental.chat.system.transform'](
+      { sessionID: 'parent-1' },
+      system,
+    );
+
+    const prompt = system.system.join('\n');
+    expect(prompt).not.toContain('exp-1 thread 1');
+    expect(prompt).not.toContain('file-1.ts');
+    expect(prompt).toContain('exp-2 thread 2');
+    expect(prompt).toContain('file-2.ts (12 lines)');
+    expect(prompt).toContain('exp-3 thread 3');
+    expect(prompt).toContain('file-3.ts (12 lines)');
+  });
+
   test('drops stale remembered sessions and falls back to fresh', async () => {
     const { hook } = createHook();
 

+ 38 - 2
src/hooks/task-session-manager/index.ts

@@ -100,6 +100,7 @@ export function createTaskSessionManagerHook(
   const pendingCalls = new Map<string, PendingTaskCall>();
   const pendingCallOrder: string[] = [];
   const contextByTask = new Map<string, Map<string, PendingContextFile>>();
+  const pendingManagedTaskIds = new Set<string>();
 
   function addTaskContext(taskId: string, files: ContextFile[]): void {
     if (files.length === 0) return;
@@ -134,6 +135,21 @@ export function createTaskSessionManagerHook(
     }));
   }
 
+  function canTrackTaskContext(taskId: string): boolean {
+    return (
+      pendingManagedTaskIds.has(taskId) || sessionManager.taskIds().has(taskId)
+    );
+  }
+
+  function pruneContext(): void {
+    const remembered = sessionManager.taskIds();
+    for (const taskId of contextByTask.keys()) {
+      if (!pendingManagedTaskIds.has(taskId) && !remembered.has(taskId)) {
+        contextByTask.delete(taskId);
+      }
+    }
+  }
+
   function isMissingRememberedSessionError(output: string): boolean {
     const firstLine = output.split(/\r?\n/, 1)[0]?.trim().toLowerCase() ?? '';
     return (
@@ -221,6 +237,7 @@ export function createTaskSessionManagerHook(
       }
 
       args.task_id = remembered.taskId;
+      pendingManagedTaskIds.add(remembered.taskId);
       sessionManager.markUsed(
         input.sessionID,
         args.subagent_type,
@@ -242,7 +259,7 @@ export function createTaskSessionManagerHook(
       output: { output: unknown; metadata?: unknown },
     ): Promise<void> => {
       if (input.tool.toLowerCase() === 'read') {
-        if (input.sessionID) {
+        if (input.sessionID && canTrackTaskContext(input.sessionID)) {
           addTaskContext(
             input.sessionID,
             extractReadFiles(_ctx.directory, output),
@@ -285,8 +302,10 @@ export function createTaskSessionManagerHook(
         agentType: pending.agentType,
         label: pending.label,
       });
+      pendingManagedTaskIds.delete(taskId);
       const contextFiles = contextFilesForPrompt(contextByTask.get(taskId));
       sessionManager.addContext(taskId, contextFiles);
+      pruneContext();
     },
 
     'experimental.chat.system.transform': async (
@@ -305,9 +324,24 @@ export function createTaskSessionManagerHook(
     event: async (input: {
       event: {
         type: string;
-        properties?: { info?: { id?: string }; sessionID?: string };
+        properties?: {
+          info?: { id?: string; parentID?: string };
+          sessionID?: string;
+        };
       };
     }): Promise<void> => {
+      if (input.event.type === 'session.created') {
+        const info = input.event.properties?.info;
+        if (
+          info?.id &&
+          info.parentID &&
+          options.shouldManageSession(info.parentID)
+        ) {
+          pendingManagedTaskIds.add(info.id);
+        }
+        return;
+      }
+
       if (input.event.type !== 'session.deleted') return;
       const sessionId =
         input.event.properties?.info?.id ?? input.event.properties?.sessionID;
@@ -316,6 +350,8 @@ export function createTaskSessionManagerHook(
       sessionManager.clearParent(sessionId);
       sessionManager.dropTask(sessionId);
       contextByTask.delete(sessionId);
+      pendingManagedTaskIds.delete(sessionId);
+      pruneContext();
 
       for (const [callId, pending] of pendingCalls.entries()) {
         if (pending.parentSessionId !== sessionId) {

+ 12 - 0
src/utils/session-manager.ts

@@ -156,6 +156,18 @@ export class SessionManager {
     }
   }
 
+  taskIds(): Set<string> {
+    const ids = new Set<string>();
+    for (const groups of this.sessionsByParent.values()) {
+      for (const group of groups.values()) {
+        for (const entry of group) {
+          ids.add(entry.taskId);
+        }
+      }
+    }
+    return ids;
+  }
+
   addContext(taskId: string, files: ContextFile[]): void {
     if (files.length === 0) return;