feat(harness-pi): pass the harness session to extension factories

.changeset/harness-pi-extension-session.md

@@ -0,0 +1,7 @@
+---
+'@ai-sdk/harness-pi': patch
+---
+
+feat(harness-pi): pass the harness session to extension factories
+
+Each `extensionFactories` entry now receives a second argument, a `PiHarnessExtensionSession` with the restricted `sandboxSession`, the `sessionWorkDir` and an `instructions()` getter for the current turn's instructions. An extension can now run work against the session's sandbox with the session's instructions, such as a tool that delegates to a child Pi session. Existing `(pi) => void` factories keep working unchanged.

content/providers/02-ai-sdk-harnesses/03-pi.mdx

@@ -161,6 +161,31 @@ const harness = createPi({
 });
 ```
 
+Each factory also receives the harness session it runs in as a second argument:
+
+- `sandboxSession`: the session's sandbox, restricted to the operations harness
+  tools may use.
+- `sessionWorkDir`: the sandbox directory the session works in.
+- `instructions()`: the instructions of the current turn. Instructions arrive
+  with each turn, so call it while a turn runs, not when the factory runs.
+
+```ts
+const harness = createPi({
+  extensionFactories: [
+    (pi, session) => {
+      pi.on('agent_start', async () => {
+        const { stdout } = await session.sandboxSession.run({
+          command: `git -C ${session.sessionWorkDir} status --short`,
+        });
+        console.log(session.instructions(), stdout);
+      });
+    },
+  ],
+});
+```
+
+Factories that take only the Pi API keep working unchanged.
+
 Routine resource refreshes between turns reuse the active extension runtime and
 do not reinitialize factories. If the underlying Pi session is rebuilt, the
 factories initialize again for the new Pi runtime.

packages/harness-pi/README.md

@@ -92,4 +92,23 @@ const harness = createPi({
 });
 ```
 
+Each factory also receives the harness session it runs in, as a second argument. It carries the restricted `sandboxSession`, the `sessionWorkDir` and an `instructions()` getter for the current turn's instructions. Instructions arrive with each turn, so read them while a turn runs rather than when the factory runs. This extension uses all three to report the sandbox's git status when the agent starts:
+
+```ts
+import { createPi } from '@ai-sdk/harness-pi';
+
+const harness = createPi({
+  extensionFactories: [
+    (pi, session) => {
+      pi.on('agent_start', async () => {
+        const { stdout } = await session.sandboxSession.run({
+          command: `git -C ${session.sessionWorkDir} status --short`,
+        });
+        console.log(session.instructions(), stdout);
+      });
+    },
+  ],
+});
+```
+
 Routine resource refreshes between turns do not reinitialize extension factories. If the underlying Pi session is rebuilt, factories initialize for the new Pi runtime. Extension factories execute in the host Node.js process, so only pass factories you trust. This option does not enable filesystem extension discovery: user, project, personal, and settings-based Pi extensions remain disabled. Themes and prompt templates also remain disabled.

packages/harness-pi/src/index.ts

@@ -9,4 +9,8 @@ export const pi = createPi();
 export { createPi } from './pi-harness';
 export { VERSION } from './version';
 export type { PiHarnessSettings } from './pi-harness';
+export type {
+  PiHarnessExtensionFactory,
+  PiHarnessExtensionSession,
+} from './pi-session';
 export type { PiAuthenticationMode, PiCredentialStore } from './pi-auth';

packages/harness-pi/src/pi-harness.test-d.ts

@@ -1,7 +1,9 @@
+import type { ExtensionFactory } from '@earendil-works/pi-coding-agent';
 import { expectTypeOf, test } from 'vitest';
 import {
   createPi,
   type PiCredentialStore,
+  type PiHarnessExtensionSession,
   type PiHarnessSettings,
 } from './index';
 
@@ -19,6 +21,24 @@ test('PiHarnessSettings accepts readonly extension factory arrays', () => {
   createPi(settings);
 });
 
+test('createPi accepts plain Pi extension factories and session-aware ones', () => {
+  const piFactory: ExtensionFactory = pi => {
+    pi.on('agent_start', () => {});
+  };
+  createPi({
+    extensionFactories: [
+      piFactory,
+      (pi, session) => {
+        expectTypeOf(session).toEqualTypeOf<PiHarnessExtensionSession>();
+        expectTypeOf(session.instructions()).toEqualTypeOf<
+          string | undefined
+        >();
+        pi.on('agent_start', () => {});
+      },
+    ],
+  });
+});
+
 test('createPi accepts the max thinking level', () => {
   createPi({ thinkingLevel: 'max' });
 });

packages/harness-pi/src/pi-harness.ts

@@ -4,17 +4,15 @@ import {
   type HarnessV1BuiltinTool,
 } from '@ai-sdk/harness';
 import { tool } from '@ai-sdk/provider-utils';
-import type {
-  ExtensionFactory,
-  ProviderConfig,
-} from '@earendil-works/pi-coding-agent';
+import type { ProviderConfig } from '@earendil-works/pi-coding-agent';
 import { z } from 'zod/v4';
 import type { PiAuthenticationMode, PiCredentialStore } from './pi-auth';
 import { piResumeStateSchema, type PiResumeStateData } from './pi-resume-state';
 import {
   createPiSession,
   type PiCacheRetention,
   type PiFileToolPathPolicy,
+  type PiHarnessExtensionFactory,
   type PiMcpSettings,
   type PiThinkingLevel,
 } from './pi-session';
@@ -82,9 +80,14 @@ export type PiHarnessSettings = {
   /**
    * Trusted inline Pi extensions loaded for each harness session.
    *
+   * Each factory receives the Pi extension API and the harness session it
+   * runs in: the restricted sandbox session, the session work directory and
+   * a getter for the current turn's instructions. Plain Pi `ExtensionFactory`
+   * functions that take only the API keep working.
+   *
    * Filesystem-discovered user and project extensions remain disabled.
    */
-  readonly extensionFactories?: ReadonlyArray<ExtensionFactory>;
+  readonly extensionFactories?: ReadonlyArray<PiHarnessExtensionFactory>;
   /**
    * Absolute sandbox paths that widen or narrow what Pi's native file tools
    * can reach. `readableRoots` lets `read`, `ls`, `find`, and `grep` reach

packages/harness-pi/src/pi-session.test.ts

@@ -32,7 +32,11 @@ import { tmpdir } from 'node:os';
 import path from 'node:path';
 import { z } from 'zod/v4';
 import { createPi } from './pi-harness';
-import { createPiSession } from './pi-session';
+import {
+  createPiSession,
+  type PiHarnessExtensionFactory,
+  type PiHarnessExtensionSession,
+} from './pi-session';
 import { truncatePiToolOutputHead } from './pi-tool-result';
 
 type FakePiTool = Pick<ToolDefinition, 'name' | 'execute'>;
@@ -276,6 +280,60 @@ describe('createPiSession', () => {
     }
   });
 
+  it('gives inline extensions the harness session and the current turn instructions', async () => {
+    const sandboxSession = createSandboxSession();
+    const sessions: PiHarnessExtensionSession[] = [];
+    const instructionsSeen: Array<string | undefined> = [];
+    const factory: PiHarnessExtensionFactory = (piApi, extensionSession) => {
+      sessions.push(extensionSession);
+      piApi.on('agent_start', () => {
+        instructionsSeen.push(extensionSession.instructions());
+      });
+    };
+    piMock.session = Object.assign(
+      createFakePiSession({
+        promptImplementation: async () => {
+          for (const handler of piMock.extensionHandlers.get('agent_start') ??
+            []) {
+            await handler();
+          }
+        },
+      }).session,
+      { getActiveToolNames: () => [], setActiveToolsByName: vi.fn() },
+    );
+
+    const session = await createPi({ extensionFactories: [factory] }).doStart({
+      sessionId: 'session-extension-harness-session',
+      sandboxSession,
+      sessionWorkDir: '/sandbox/work',
+    });
+
+    try {
+      for (const instructions of ['Answer in French.', 'Answer in German.']) {
+        const control = await session.doPromptTurn({
+          skills: [],
+          prompt: 'Hello.',
+          instructions,
+          tools: [],
+          emit: vi.fn(),
+        });
+        await control.done;
+      }
+
+      expect(sessions).toHaveLength(1);
+      expect(sessions[0]).toMatchObject({
+        sandboxSession: sandboxSession.restricted(),
+        sessionWorkDir: '/sandbox/work',
+      });
+      expect(instructionsSeen).toEqual([
+        'Answer in French.',
+        'Answer in German.',
+      ]);
+    } finally {
+      await session.doDestroy();
+    }
+  });
+
   it('lets inline extensions read a host-backed session workspace', async () => {
     const sessionWorkDir = mkdtempSync(
       path.join(tmpdir(), 'pi-host-workspace-'),
@@ -502,7 +560,7 @@ describe('createPiSession', () => {
     try {
       const extensionFactoryInput = piMock.extensionFactoryInputs.at(-1);
       expect(callOrder).toEqual(['first', 'second']);
-      expect(extensionFactoryInput?.snapshot).toEqual(extensionFactories);
+      expect(extensionFactoryInput?.snapshot).toHaveLength(2);
       expect(extensionFactoryInput?.reference).not.toBe(extensionFactories);
     } finally {
       await session.doDestroy();
@@ -1469,7 +1527,7 @@ describe('createPiSession', () => {
 
     try {
       expect(piMock.extensionFactoryInputs.at(-1)?.snapshot).toEqual([
-        factory,
+        expect.any(Function),
         mcpAdapterMock.mcpExtensionFactory,
       ]);
       expect(factory).toHaveBeenCalledOnce();

packages/harness-pi/src/pi-session.ts

@@ -7,6 +7,7 @@ import {
   SettingsManager,
   type AgentSession,
   type AgentToolResult,
+  type ExtensionAPI,
   type ExtensionFactory,
   type ProviderConfig,
   type Skill,
@@ -275,6 +276,34 @@ export interface PiFileToolPathPolicy {
   readonly deniedRoots?: ReadonlyArray<string>;
 }
 
+/**
+ * The harness session a Pi extension factory runs in.
+ */
+export interface PiHarnessExtensionSession {
+  /**
+   * The session's sandbox, restricted to the operations harness tools may use.
+   */
+  readonly sandboxSession: SandboxSession;
+  /**
+   * The sandbox directory the session works in.
+   */
+  readonly sessionWorkDir: string;
+  /**
+   * The instructions of the session's current turn. Instructions arrive with
+   * each turn, so read them when a turn runs, not when the factory runs.
+   */
+  readonly instructions: () => string | undefined;
+}
+
+/**
+ * A Pi extension factory that also receives the harness session it runs in.
+ * Plain Pi `ExtensionFactory` functions are accepted unchanged.
+ */
+export type PiHarnessExtensionFactory = (
+  pi: ExtensionAPI,
+  session: PiHarnessExtensionSession,
+) => ReturnType<ExtensionFactory>;
+
 export interface PiSessionSettings {
   readonly auth?: PiAuthenticationMode;
   readonly credentials?: PiCredentialStore;
@@ -285,7 +314,7 @@ export interface PiSessionSettings {
   readonly mcpServers?: Record<string, unknown>;
   readonly mcpSettings?: PiMcpSettings;
   readonly providers?: Readonly<Record<string, ProviderConfig>>;
-  readonly extensionFactories?: ReadonlyArray<ExtensionFactory>;
+  readonly extensionFactories?: ReadonlyArray<PiHarnessExtensionFactory>;
   readonly fileToolPathPolicy?: PiFileToolPathPolicy;
 }
 
@@ -627,14 +656,24 @@ export async function createPiSession(
 
   let sessionInstructions: string | undefined;
 
+  const extensionSession: PiHarnessExtensionSession = {
+    sandboxSession: toolSafeSandboxSession,
+    sessionWorkDir,
+    instructions: () => sessionInstructions,
+  };
+
   /*
    * Configured MCP servers are served by an inline Pi extension, so they share
    * the extension runtime with the caller-supplied factories: both are loaded
    * by the resource loader below and both are subject to the reload handling
    * that keeps the active runtime alive across resource-only reloads.
    */
   const extensionFactories: ExtensionFactory[] = [
-    ...(input.settings.extensionFactories ?? []),
+    ...(input.settings.extensionFactories ?? []).map(
+      (factory): ExtensionFactory =>
+        pi =>
+          factory(pi, extensionSession),
+    ),
   ];
   if (hasMcpServers) {
     const { createMcpAdapter } = (await import(