From 391d677b9edfaaf1a17f39f0620ccab7168caa4a Mon Sep 17 00:00:00 2001 From: Chris Huber Date: Wed, 21 Oct 2026 09:16:33 -0400 Subject: [PATCH] Studio Code: Load pi extensions from Studio's agent directory Studio Code sessions now load the user's pi extensions from ~/.studio/agent/extensions, the way pi loads ~/.pi/agent/extensions. An extension can add tools, register providers, and extend the system prompt on each turn through pi's own API. The agent directory moves off the sites root, which stays an untrusted project, so project extensions do not load. An extension that fails to load is reported and skipped. --- apps/cli/ai/runtimes/pi/index.ts | 21 ++++++-- apps/cli/ai/tests/pi-runtime.test.ts | 72 ++++++++++++++++++++++++- packages/common/lib/well-known-paths.ts | 6 +++ 3 files changed, 93 insertions(+), 6 deletions(-) diff --git a/apps/cli/ai/runtimes/pi/index.ts b/apps/cli/ai/runtimes/pi/index.ts index 38e7f9f0c4..b3ab157eec 100644 --- a/apps/cli/ai/runtimes/pi/index.ts +++ b/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,10 +358,13 @@ 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', @@ -366,7 +372,12 @@ async function createStudioAgentSession( 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' }, } ); diff --git a/apps/cli/ai/tests/pi-runtime.test.ts b/apps/cli/ai/tests/pi-runtime.test.ts index dfbb6a94ab..7fd704c007 100644 --- a/apps/cli/ai/tests/pi-runtime.test.ts +++ b/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( () => {} ); diff --git a/packages/common/lib/well-known-paths.ts b/packages/common/lib/well-known-paths.ts index 52f9c7e364..97fba1db7a 100644 --- a/packages/common/lib/well-known-paths.ts +++ b/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' ); }