fix(coding-agent): isolate auto-learn capture lifecycle

packages/coding-agent/src/sdk.ts

@@ -1069,6 +1069,8 @@ export interface AutoLearnCaptureRunnerOptions {
 	sourceAgent: Agent;
 	captureTools: AgentTool[];
 	createAgent: (options: AgentOptions) => Agent;
+	onPayload?: SimpleStreamOptions["onPayload"];
+	onResponse?: SimpleStreamOptions["onResponse"];
 	createSessionId?: () => string;
 }
 
@@ -1105,6 +1107,8 @@ export function createAutoLearnCaptureRunner(
 			promptCacheKey: captureSessionId,
 			providerSessionState: captureProviderSessionState,
 			getApiKey: requestModel => options.sourceAgent.getApiKey?.(requestModel),
+			onPayload: options.onPayload,
+			onResponse: options.onResponse,
 		});
 		captureAgent.setMetadataResolver(provider => options.sourceAgent.metadataForProvider(provider));
 		const captureMessage: CustomMessage = {
@@ -3013,6 +3017,8 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
 		const runAutoLearnCapture = createAutoLearnCaptureRunner({
 			sourceAgent: agent,
 			captureTools: autoLearnCaptureTools,
+			onPayload,
+			onResponse,
 			createAgent: captureOptions => {
 				const captureModel = captureOptions.initialState?.model;
 				const captureSessionId = captureOptions.sessionId;

packages/coding-agent/src/session/agent-session.ts

@@ -6259,6 +6259,10 @@ export class AgentSession {
 		}
 	}
 
+	#abortAutolearnCapture(): void {
+		this.#autolearnCaptureAbortController?.abort();
+	}
+
 	async #drainAutolearnCapture(): Promise<void> {
 		const task = this.#autolearnCaptureTask;
 		if (!task) return;
@@ -6290,7 +6294,7 @@ export class AgentSession {
 	 */
 	beginDispose(): void {
 		this.#isDisposed = true;
-		this.#autolearnCaptureAbortController?.abort();
+		this.#abortAutolearnCapture();
 		this.#flushPendingIrcAsides();
 		this.yieldQueue.clear();
 		this.agent.setAsideMessageProvider(undefined);
@@ -9125,6 +9129,7 @@ export class AgentSession {
 		// auto-starting a fresh turn during cleanup.
 		this.#abortInProgress = true;
 		try {
+			this.#abortAutolearnCapture();
 			this.abortRetry();
 			this.#promptGeneration++;
 			this.#scheduledHiddenNextTurnGeneration = undefined;
@@ -9146,6 +9151,7 @@ export class AgentSession {
 			this.agent.abort(options?.reason);
 			await postPromptDrain;
 			await this.agent.waitForIdle();
+			await this.#drainAutolearnCapture();
 			await this.#goalRuntime.onTaskAborted({ reason: options?.goalReason ?? "interrupted" });
 			// Clear prompt-in-flight state: waitForIdle resolves when the agent loop's finally
 			// block runs, but nested prompt setup/finalizers may still be unwinding. Without this,
@@ -15749,6 +15755,8 @@ export class AgentSession {
 		// Flush pending writes before branching
 		await this.sessionManager.flush();
 		this.#cancelOwnAsyncJobs();
+		this.#abortAutolearnCapture();
+		await this.#drainAutolearnCapture();
 
 		if (!selectedEntry.parentId) {
 			await this.sessionManager.newSession({ parentSession: previousSessionFile });
@@ -15839,6 +15847,8 @@ export class AgentSession {
 		}
 		await this.sessionManager.flush();
 		this.#cancelOwnAsyncJobs();
+		this.#abortAutolearnCapture();
+		await this.#drainAutolearnCapture();
 
 		this.sessionManager.createBranchedSession(leafId);
 

packages/coding-agent/test/autolearn-controller.test.ts

@@ -1,5 +1,5 @@
 import { describe, expect, it } from "bun:test";
-import { Agent, type AgentMessage, type AgentTool } from "@oh-my-pi/pi-agent-core";
+import { Agent, type AgentMessage, type AgentOptions, type AgentTool } from "@oh-my-pi/pi-agent-core";
 import type { AssistantMessage, FetchImpl, Model, ProviderSessionState, Usage } from "@oh-my-pi/pi-ai";
 import { streamGoogle } from "@oh-my-pi/pi-ai/providers/google";
 import { createMockModel } from "@oh-my-pi/pi-ai/providers/mock";
@@ -442,6 +442,39 @@ describe("isolated auto-learn capture", () => {
 		expect(primaryEvents).toBe(0);
 	});
 
+	it("forwards provider lifecycle hooks to the detached capture", async () => {
+		const captureMock = createMockModel({ responses: [{ content: ["Captured."] }] });
+		const manageSkillTool = captureTool("manage_skill", "Manage reusable skills");
+		const sourceAgent = new Agent({
+			initialState: { model: captureMock, systemPrompt: ["Test"], tools: [manageSkillTool] },
+		});
+		const onPayload: NonNullable<AgentOptions["onPayload"]> = async payload => payload;
+		const onResponse: NonNullable<AgentOptions["onResponse"]> = async () => {};
+		let captureOnPayload: AgentOptions["onPayload"];
+		let captureOnResponse: AgentOptions["onResponse"];
+		const runCapture = createAutoLearnCaptureRunner({
+			sourceAgent,
+			captureTools: [manageSkillTool],
+			onPayload,
+			onResponse,
+			createAgent: options => {
+				captureOnPayload = options.onPayload;
+				captureOnResponse = options.onResponse;
+				return new Agent({
+					...options,
+					convertToLlm,
+					streamFn: captureMock.stream,
+				});
+			},
+		});
+
+		await runCapture("Capture with provider hooks");
+
+		expect(captureMock.calls).toHaveLength(1);
+		expect(captureOnPayload).toBe(onPayload);
+		expect(captureOnResponse).toBe(onResponse);
+	});
+
 	it("adds learn alongside manage_skill when a memory backend provides it", async () => {
 		const model = googleInteractionsModel();
 		const manageSkillTool = captureTool("manage_skill", "Manage reusable skills");