From 4e145a0f3f951da7beac4297ae3ee5b44edaf6d8 Mon Sep 17 00:00:00 2001
From: ezl-keygraph <ezhil@keygraph.io>
Date: Mon, 31 Aug 2026 20:36:57 +0530
Subject: [PATCH] fix(pi): give each task sub-session its own resource loader
 to prevent stale extension ctx

---
 apps/worker/src/ai/pi/pi-executor.ts | 4 ++--
 apps/worker/src/ai/pi/task-tool.ts   | 5 +++--
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/apps/worker/src/ai/pi/pi-executor.ts b/apps/worker/src/ai/pi/pi-executor.ts
index 7c1f149d..01a3f258 100644
--- a/apps/worker/src/ai/pi/pi-executor.ts
+++ b/apps/worker/src/ai/pi/pi-executor.ts
@@ -299,7 +299,7 @@ export async function runPiPrompt(
         childUsage.cacheReadTokens += usage.cacheReadTokens;
         childUsage.cacheWriteTokens += usage.cacheWriteTokens;
       },
-      resourceLoader,
+      createResourceLoader: () => buildResourceLoader(sourceDir, logger, agentName),
       ...(cancellationSignal && { cancellationSignal }),
     }),
     createTodoWriteTool(),
@@ -403,7 +403,6 @@ export async function runPiPrompt(
 
     // 6. Run the agent to completion (resolves at agent_end).
     await session.prompt(fullPrompt);
-    session.dispose();
 
     // 7. Surface any error captured during the run.
     if (pendingError) throw pendingError;
@@ -469,6 +468,7 @@ export async function runPiPrompt(
       retryable,
     };
   } finally {
+    session?.dispose();
     cancellationSignal?.removeEventListener('abort', onCancellation);
   }
 }
diff --git a/apps/worker/src/ai/pi/task-tool.ts b/apps/worker/src/ai/pi/task-tool.ts
index 0257b5e4..c904abb5 100644
--- a/apps/worker/src/ai/pi/task-tool.ts
+++ b/apps/worker/src/ai/pi/task-tool.ts
@@ -26,7 +26,7 @@ export interface TaskToolContext {
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   readonly model: Model<any>;
   readonly modelRuntime: ModelRuntime;
-  readonly resourceLoader: ResourceLoader;
+  readonly createResourceLoader: () => Promise<ResourceLoader>;
   readonly parentAgentName: LoggableAgentName;
   readonly workflowLogPath?: string | undefined;
   readonly onDelegationStart?: ((child: string) => Promise<void>) | undefined;
@@ -129,10 +129,11 @@ export function createTaskTool(config: TaskToolContext): ToolDefinition {
       const onCancellation = (): void => abortChildSession();
 
       try {
+        const resourceLoader = await config.createResourceLoader();
         ({ session: subSession } = await createAgentSession({
           cwd: config.cwd,
           agentDir,
-          resourceLoader: config.resourceLoader,
+          resourceLoader,
           model: config.model,
           tools: CHILD_TOOLS,
           modelRuntime: config.modelRuntime,