diff --git a/.changeset/harness-pi-extension-session.md b/.changeset/harness-pi-extension-session.md new file mode 100644 index 000000000000..56fd1bc60cff --- /dev/null +++ b/.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. diff --git a/content/providers/02-ai-sdk-harnesses/03-pi.mdx b/content/providers/02-ai-sdk-harnesses/03-pi.mdx index 322f48295214..e6679acbb4a9 100644 --- a/content/providers/02-ai-sdk-harnesses/03-pi.mdx +++ b/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. diff --git a/packages/harness-pi/README.md b/packages/harness-pi/README.md index 4954b48903a3..5411e52eeb7e 100644 --- a/packages/harness-pi/README.md +++ b/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. diff --git a/packages/harness-pi/src/index.ts b/packages/harness-pi/src/index.ts index 7413b4568bb0..810c91432e21 100644 --- a/packages/harness-pi/src/index.ts +++ b/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'; diff --git a/packages/harness-pi/src/pi-harness.test-d.ts b/packages/harness-pi/src/pi-harness.test-d.ts index 98da66176999..dc25b6bc06a3 100644 --- a/packages/harness-pi/src/pi-harness.test-d.ts +++ b/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(); + expectTypeOf(session.instructions()).toEqualTypeOf< + string | undefined + >(); + pi.on('agent_start', () => {}); + }, + ], + }); +}); + test('createPi accepts the max thinking level', () => { createPi({ thinkingLevel: 'max' }); }); diff --git a/packages/harness-pi/src/pi-harness.ts b/packages/harness-pi/src/pi-harness.ts index 50afb757a96c..626cf2abef36 100644 --- a/packages/harness-pi/src/pi-harness.ts +++ b/packages/harness-pi/src/pi-harness.ts @@ -4,10 +4,7 @@ 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'; @@ -15,6 +12,7 @@ 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; + readonly extensionFactories?: ReadonlyArray; /** * Absolute sandbox paths that widen or narrow what Pi's native file tools * can reach. `readableRoots` lets `read`, `ls`, `find`, and `grep` reach diff --git a/packages/harness-pi/src/pi-session.test.ts b/packages/harness-pi/src/pi-session.test.ts index 878263def219..1c97d3f7e6f3 100644 --- a/packages/harness-pi/src/pi-session.test.ts +++ b/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; @@ -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 = []; + 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(); diff --git a/packages/harness-pi/src/pi-session.ts b/packages/harness-pi/src/pi-session.ts index 59d87d7ef24c..a4f617f4e94c 100644 --- a/packages/harness-pi/src/pi-session.ts +++ b/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; } +/** + * 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; + export interface PiSessionSettings { readonly auth?: PiAuthenticationMode; readonly credentials?: PiCredentialStore; @@ -285,7 +314,7 @@ export interface PiSessionSettings { readonly mcpServers?: Record; readonly mcpSettings?: PiMcpSettings; readonly providers?: Readonly>; - readonly extensionFactories?: ReadonlyArray; + readonly extensionFactories?: ReadonlyArray; readonly fileToolPathPolicy?: PiFileToolPathPolicy; } @@ -627,6 +656,12 @@ 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 @@ -634,7 +669,11 @@ export async function createPiSession( * 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(