Studio Code: Load pi extensions from Studio's agent directory

apps/cli/ai/runtimes/pi/index.ts

@@ -39,6 +39,7 @@ import {
 } from '@studio/common/ai/models';
 import { findLastAssistant, isEmptyStop } from '@studio/common/ai/session-events';
 import { getAiTracksIdentity } from '@studio/common/ai/tracks-identity';
+import { getAgentDirectory } from '@studio/common/lib/well-known-paths';
 import { type TSchema } from 'typebox';
 import { withDesignSystemPreview } from 'cli/ai/chat-artifacts';
 import { isImageGenerationAvailable } from 'cli/ai/image-generation';
@@ -73,7 +74,6 @@ type StudioModel = StudioWpcomModel | Model< 'anthropic-messages' >;
 type ProviderConfigInput = Parameters< ModelRuntime[ 'registerProvider' ] >[ 1 ];
 
 const STUDIO_WPCOM_PROVIDER = 'studio-wpcom';
-const STUDIO_AGENT_DIR = STUDIO_SITES_ROOT;
 // A reply with nothing in it ends the turn without the user seeing why (see
 // `isEmptyStop`). Resume it with a message the user doesn't see, up to this
 // many times per turn. pi leaves a reply that has only reasoning out of the
@@ -342,11 +342,14 @@ async function createStudioAgentSession(
 	);
 	const modelRuntime = await createModelRuntime( model, family, creds );
 	const settingsManager = createSettingsManager( config.env );
+	// pi extensions from Studio's agent directory load into the session: they can
+	// add tools, register providers and extend the prompt on each turn. The sites
+	// root is not a trusted project, so only the user's own extensions load.
+	const agentDir = getAgentDirectory();
 	const resourceLoader = new DefaultResourceLoader( {
 		cwd: STUDIO_SITES_ROOT,
-		agentDir: STUDIO_AGENT_DIR,
+		agentDir,
 		settingsManager,
-		noExtensions: true,
 		noSkills: true,
 		noPromptTemplates: true,
 		noThemes: true,
@@ -355,18 +358,26 @@ async function createStudioAgentSession(
 		appendSystemPrompt: [],
 	} );
 	await resourceLoader.reload();
+	for ( const { path: extensionPath, error } of resourceLoader.getExtensions().errors ) {
+		console.warn( `Skipped the pi extension ${ extensionPath }: ${ error }` );
+	}
 
 	const result = await createAgentSession( {
 		cwd: STUDIO_SITES_ROOT,
-		agentDir: STUDIO_AGENT_DIR,
+		agentDir,
 		modelRuntime,
 		model,
 		thinkingLevel: 'medium',
 		sessionManager: config.session,
 		settingsManager,
 		resourceLoader,
 		customTools: toolDefinitions,
-		tools: toolDefinitions.map( ( tool ) => tool.name ),
+		tools: [
+			...toolDefinitions.map( ( tool ) => tool.name ),
+			...resourceLoader
+				.getExtensions()
+				.extensions.flatMap( ( extension ) => [ ...extension.tools.keys() ] ),
+		],
 		sessionStartEvent: { type: 'session_start', reason: 'startup' },
 	} );
 

apps/cli/ai/tests/pi-runtime.test.ts

@@ -1,7 +1,9 @@
+import fs from 'fs';
+import path from 'path';
 import { ANTHROPIC_MODELS } from '@earendil-works/pi-ai/providers/anthropic.models';
 import { ModelRegistry, SessionManager } from '@earendil-works/pi-coding-agent';
 import { AI_MODELS } from '@studio/common/ai/models';
-import { beforeEach, describe, expect, it, vi } from 'vitest';
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
 import { runStudioAgentTurn, type StudioAgentTurnConfig } from 'cli/ai/runtimes/pi';
 import type { AgentSessionEvent, CreateAgentSessionOptions } from '@earendil-works/pi-coding-agent';
 import type { AiModelId } from '@studio/common/ai/models';
@@ -63,6 +65,7 @@ vi.mock( '@studio/common/lib/well-known-paths', async ( importOriginal ) => {
 		...actual,
 		getConfigDirectory: () => mocks.configRoot,
 		getAiPayloadsPath: () => `${ mocks.configRoot }/tmp/ai-payloads`,
+		getAgentDirectory: () => `${ mocks.configRoot }/agent`,
 	};
 } );
 
@@ -718,6 +721,73 @@ describe( 'pi runtime', () => {
 	} );
 
 	// Silent header-drop would surface as an opaque 401 from the wpcom proxy.
+	describe( 'pi extensions', () => {
+		const extensionsDir = path.join( mocks.configRoot, 'agent', 'extensions' );
+
+		beforeEach( () => {
+			fs.rmSync( extensionsDir, { recursive: true, force: true } );
+			fs.mkdirSync( extensionsDir, { recursive: true } );
+		} );
+
+		afterEach( () => {
+			fs.rmSync( extensionsDir, { recursive: true, force: true } );
+		} );
+
+		it( "loads the user's extensions from Studio's agent directory and enables their tools", async () => {
+			fs.writeFileSync(
+				path.join( extensionsDir, 'greeter.js' ),
+				`export default function ( pi ) {
+					pi.registerTool( {
+						name: 'greet',
+						label: 'Greet',
+						description: 'Greets.',
+						parameters: { type: 'object', properties: {} },
+						execute: async () => ( { content: [ { type: 'text', text: 'hi' } ], details: undefined } ),
+					} );
+					pi.on( 'before_agent_start', ( event ) => ( { systemPrompt: event.systemPrompt + ' Greet first.' } ) );
+				}`
+			);
+
+			await runRuntime( {
+				prompt: 'hello',
+				env: WPCOM_ENV,
+				model: 'balanced',
+				session: newSession(),
+			} );
+
+			const options = mocks.createAgentSession.mock.calls[ 0 ][ 0 ] as CreateAgentSessionOptions;
+			const extensions = options.resourceLoader!.getExtensions().extensions;
+			expect( extensions.map( ( extension ) => path.basename( extension.path ) ) ).toEqual( [
+				'greeter.js',
+			] );
+			expect( extensions[ 0 ].handlers.has( 'before_agent_start' ) ).toBe( true );
+			expect( options.tools ).toContain( 'greet' );
+			expect( options.tools ).toContain( 'site_list' );
+		} );
+
+		it( 'warns about an extension that fails to load and runs the turn without it', async () => {
+			fs.writeFileSync( path.join( extensionsDir, 'broken.js' ), 'export default (' );
+			const warnSpy = vi.spyOn( console, 'warn' ).mockImplementation( () => {} );
+
+			try {
+				await runRuntime( {
+					prompt: 'hello',
+					env: WPCOM_ENV,
+					model: 'balanced',
+					session: newSession(),
+				} );
+
+				expect( warnSpy.mock.calls.map( ( call ) => call[ 0 ] ).join( '\n' ) ).toMatch(
+					/Skipped the pi extension .*broken\.js/
+				);
+				const options = mocks.createAgentSession.mock.calls[ 0 ][ 0 ] as CreateAgentSessionOptions;
+				expect( options.resourceLoader!.getExtensions().extensions ).toEqual( [] );
+			} finally {
+				warnSpy.mockRestore();
+			}
+		} );
+	} );
+
 	it( 'warns and continues when STUDIO_WPCOM_DEFAULT_HEADERS is malformed', async () => {
 		const warnSpy = vi.spyOn( console, 'warn' ).mockImplementation( () => {} );
 

packages/common/lib/well-known-paths.ts

@@ -38,6 +38,12 @@ export function getSessionsDirectory(): string {
 	return path.join( getConfigDirectory(), 'sessions' );
 }
 
+// Studio Code's pi agent directory. Extensions in its `extensions/` folder load
+// into every Studio Code session, as they would into pi's own `~/.pi/agent`.
+export function getAgentDirectory(): string {
+	return path.join( getConfigDirectory(), 'agent' );
+}
+
 export function getKnowledgeDirectory(): string {
 	return path.join( getConfigDirectory(), 'knowledge' );
 }