Browse Source

fix: bound session read context storage

Alvin Unreal 3 months ago
parent
commit
6e47d060d2

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

@@ -246,6 +246,51 @@ describe('task-session-manager hook', () => {
     expect(prompt).toContain('src/large.ts (12 lines)');
   });
 
+  test('counts overlapping repeated reads once per unique line', async () => {
+    const { hook } = createHook();
+
+    await hook.event({
+      event: {
+        type: 'session.created',
+        properties: { info: { id: 'child-1', parentID: 'parent-1' } },
+      },
+    });
+    for (const call of ['read-1', 'read-2']) {
+      await hook['tool.execute.after'](
+        { tool: 'read', sessionID: 'child-1', callID: call },
+        {
+          output: [
+            '<path>/tmp/src/repeat.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: 'repeat reads' } },
+    );
+    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,
+    );
+
+    expect(system.system.join('\n')).toContain('src/repeat.ts (12 lines)');
+    expect(system.system.join('\n')).not.toContain('src/repeat.ts (24 lines)');
+  });
+
   test('uses configured read context thresholds', async () => {
     const { hook } = createHook({
       readContextMinLines: 5,

+ 10 - 7
src/hooks/task-session-manager/index.ts

@@ -39,7 +39,7 @@ const MAX_PENDING_TASK_CALLS = 100;
 
 interface PendingContextFile {
   path: string;
-  lineCount: number;
+  lines: Set<number>;
   lastReadAt: number;
 }
 
@@ -75,18 +75,19 @@ function extractReadFiles(
   return [
     {
       path: normalizePath(root, file),
-      lineCount: countReadLines(output.output),
+      lineCount: countReadLines(output.output).length,
+      lineNumbers: countReadLines(output.output),
       lastReadAt: Date.now(),
     },
   ];
 }
 
-function countReadLines(output: string): number {
+function countReadLines(output: string): number[] {
   const lines = new Set<number>();
   for (const match of output.matchAll(/^([0-9]+):/gm)) {
     lines.add(Number(match[1]));
   }
-  return lines.size;
+  return [...lines];
 }
 
 export function createTaskSessionManagerHook(
@@ -118,10 +119,12 @@ export function createTaskSessionManagerHook(
     for (const file of files) {
       const pending = context.get(file.path) ?? {
         path: file.path,
-        lineCount: 0,
+        lines: new Set<number>(),
         lastReadAt: file.lastReadAt,
       };
-      pending.lineCount += file.lineCount;
+      for (const line of file.lineNumbers ?? []) {
+        pending.lines.add(line);
+      }
       pending.lastReadAt = Math.max(pending.lastReadAt, file.lastReadAt);
       context.set(file.path, pending);
     }
@@ -135,7 +138,7 @@ export function createTaskSessionManagerHook(
     if (!context) return [];
     return [...context.values()].map((file) => ({
       path: file.path,
-      lineCount: file.lineCount,
+      lineCount: file.lines.size,
       lastReadAt: file.lastReadAt,
     }));
   }

+ 29 - 0
src/utils/session-manager.test.ts

@@ -119,6 +119,35 @@ describe('SessionManager', () => {
     expect(prompt).not.toContain('medium.ts');
     expect(prompt).toContain('(+1 more)');
   });
+
+  test('bounds stored read context files to the render cap plus overflow marker', () => {
+    const manager = new SessionManager(2, {
+      readContextMinLines: 1,
+      readContextMaxFiles: 2,
+    });
+
+    const remembered = manager.remember({
+      parentSessionId: 'parent-1',
+      taskId: 'task-1',
+      agentType: 'explorer',
+      label: 'bounded context',
+    });
+    manager.addContext(
+      'task-1',
+      Array.from({ length: 10 }, (_, index) => ({
+        path: `file-${index}.ts`,
+        lineCount: 10,
+        lastReadAt: index,
+      })),
+    );
+
+    expect(remembered.contextFiles).toHaveLength(3);
+    const prompt = manager.formatForPrompt('parent-1') ?? '';
+    expect(prompt).toContain('file-9.ts (10 lines)');
+    expect(prompt).toContain('file-8.ts (10 lines)');
+    expect(prompt).toContain('(+1 more)');
+    expect(prompt).not.toContain('file-0.ts');
+  });
 });
 
 describe('deriveTaskSessionLabel', () => {

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

@@ -3,6 +3,7 @@ import type { AgentName } from '../config';
 export interface ContextFile {
   path: string;
   lineCount: number;
+  lineNumbers?: number[];
   lastReadAt: number;
 }
 
@@ -205,6 +206,7 @@ export class SessionManager {
           }
           match.contextFiles.push({ ...file });
         }
+        this.trimContextFiles(match);
       }
     }
   }
@@ -321,6 +323,18 @@ export class SessionManager {
     }
   }
 
+  private trimContextFiles(entry: RememberedTaskSession): void {
+    if (this.readContextMaxFiles === 0) {
+      entry.contextFiles = [];
+      return;
+    }
+
+    entry.contextFiles = entry.contextFiles
+      .filter((file) => file.lineCount >= this.readContextMinLines)
+      .sort((a, b) => b.lastReadAt - a.lastReadAt)
+      .slice(0, this.readContextMaxFiles + 1);
+  }
+
   private nextOrder(): number {
     this.orderCounter += 1;
     return this.orderCounter;