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' ); }