jb/paseo
1
0
Fork 0
Code Issues 1 Pull Requests Actions Packages Projects Releases Wiki Activity

Load Pi extensions for model discovery

Browse Source
This commit is contained in:
Mohamed Boudra
2026-04-22 22:12:49 +07:00
parent de57f301f9
commit 4d5889ff38
2 changed files with 171 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/server/src/server/agent/providers/pi-direct-agent.test.ts
+125
View File
@@ -1,3 +1,6 @@
import { mkdir, mkdtemp, rm, writeFile } from "node:fs/promises";
import { join } from "node:path";
import { tmpdir } from "node:os";
import { describe, expect, test, vi } from "vitest";
import type { Api, Model } from "@mariozechner/pi-ai";
import pino from "pino";
@@ -93,4 +96,126 @@ describe("PiDirectAgentClient", () => {
expect(registry.getAll).not.toHaveBeenCalled();
expect(models.map((model) => model.id)).toEqual(["anthropic/claude-opus-4-5"]);
});
test("loads project extensions before listing available models", async () => {
const testRoot = await mkdtemp(join(tmpdir(), "paseo-pi-extension-"));
const previousAgentDir = process.env.PI_CODING_AGENT_DIR;
try {
const agentDir = join(testRoot, "agent");
const cwd = join(testRoot, "project");
const extensionDir = join(cwd, ".pi", "extensions");
process.env.PI_CODING_AGENT_DIR = agentDir;
await mkdir(extensionDir, { recursive: true });
await writeFile(
join(extensionDir, "dummy-provider.ts"),
`
export default function(pi) {
pi.registerProvider("paseo-dummy", {
baseUrl: "https://example.invalid/v1",
apiKey: "paseo-test-key",
api: "openai-responses",
models: [
{
id: "extension-model",
name: "Extension Model",
reasoning: true,
input: ["text"],
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
contextWindow: 128000,
maxTokens: 4096
}
]
});
}
`,
"utf-8",
);
const client = new PiDirectAgentClient({
logger: pino({ level: "silent" }),
});
const models = await client.listModels({ cwd, force: false });
expect(models.map((model) => model.id)).toContain("paseo-dummy/extension-model");
} finally {
if (previousAgentDir === undefined) {
delete process.env.PI_CODING_AGENT_DIR;
} else {
process.env.PI_CODING_AGENT_DIR = previousAgentDir;
}
await rm(testRoot, { recursive: true, force: true });
}
});
test("creates sessions with project extension models and exposes extension commands", async () => {
const testRoot = await mkdtemp(join(tmpdir(), "paseo-pi-extension-session-"));
const previousAgentDir = process.env.PI_CODING_AGENT_DIR;
try {
const agentDir = join(testRoot, "agent");
const cwd = join(testRoot, "project");
const extensionDir = join(cwd, ".pi", "extensions");
process.env.PI_CODING_AGENT_DIR = agentDir;
await mkdir(extensionDir, { recursive: true });
await writeFile(
join(extensionDir, "dummy-command.ts"),
`
export default function(pi) {
pi.registerProvider("paseo-dummy", {
baseUrl: "https://example.invalid/v1",
apiKey: "paseo-test-key",
api: "openai-responses",
models: [
{
id: "extension-model",
name: "Extension Model",
reasoning: true,
input: ["text"],
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
contextWindow: 128000,
maxTokens: 4096
}
]
});
pi.registerCommand("dummy-command", {
description: "Dummy extension command",
handler: async () => {}
});
}
`,
"utf-8",
);
const client = new PiDirectAgentClient({
logger: pino({ level: "silent" }),
});
const session = await client.createSession({
provider: "pi",
cwd,
model: "paseo-dummy/extension-model",
});
try {
await expect(session.listCommands()).resolves.toContainEqual({
name: "dummy-command",
description: "Dummy extension command",
argumentHint: "",
});
} finally {
await session.close();
}
} finally {
if (previousAgentDir === undefined) {
delete process.env.PI_CODING_AGENT_DIR;
} else {
process.env.PI_CODING_AGENT_DIR = previousAgentDir;
}
await rm(testRoot, { recursive: true, force: true });
}
});
});
packages/server/src/server/agent/providers/pi-direct-agent.ts
+46 -43
View File
@@ -4,12 +4,13 @@ import { join } from "node:path";
import { homedir } from "node:os";
import type { Logger } from "pino";
import {
AuthStorage,
ModelRegistry,
SessionManager,
createAgentSession,
createAgentSessionFromServices,
createAgentSessionServices,
type AgentSession as PiAgentSession,
type AgentSessionEvent,
type AgentSessionServices,
type BashToolInput,
type EditToolInput,
type FindToolInput,
@@ -101,6 +102,11 @@ interface StartTurnResult {
turnId: string;
}
interface PiDirectSessionResult {
session: PiAgentSession;
modelRegistry: ModelRegistry;
}
export type PiDirectSessionAdapter = Pick<
PiAgentSession,
| "abort"
@@ -1359,48 +1365,48 @@ export class PiDirectAgentClient implements AgentClient {
this.runtimeSettings = options.runtimeSettings;
}
private getModelRegistry(): ModelRegistry {
if (!this.modelRegistry) {
this.modelRegistry = ModelRegistry.create(AuthStorage.create());
}
return this.modelRegistry;
private async getSessionServices(cwd: string): Promise<AgentSessionServices> {
return createAgentSessionServices({
cwd,
...(this.modelRegistry ? { modelRegistry: this.modelRegistry } : {}),
});
}
private resolveConfiguredModel(modelId: string | null | undefined): Model<Api> | undefined {
private resolveConfiguredModel(
registry: PiDirectModelRegistry,
modelId: string | null | undefined,
): Model<Api> | undefined {
const parsedReference = parseModelReference(modelId ?? null);
if (!parsedReference) {
return undefined;
}
const registry = this.getModelRegistry();
return findModelInRegistry(registry, parsedReference);
}
private async createSdkSession(config: AgentSessionConfig): Promise<PiAgentSession> {
private async createSdkSession(config: AgentSessionConfig): Promise<PiDirectSessionResult> {
const thinkingLevel =
normalizePiThinkingOption(config.thinkingOptionId) ?? DEFAULT_PI_THINKING_LEVEL;
const modelRegistry = this.getModelRegistry();
const model = this.resolveConfiguredModel(config.model);
const services = await this.getSessionServices(config.cwd);
const model = this.resolveConfiguredModel(services.modelRegistry, config.model);
const sessionOptions = {
cwd: config.cwd,
modelRegistry,
const { session } = await createAgentSessionFromServices({
services,
sessionManager: SessionManager.create(config.cwd),
thinkingLevel,
...(model ? { model } : {}),
};
const { session } = await createAgentSession(sessionOptions);
});
await session.bindExtensions({});
applySystemPrompt(session, config.systemPrompt);
return session;
return { session, modelRegistry: services.modelRegistry };
}
async createSession(
config: AgentSessionConfig,
_launchContext?: AgentLaunchContext,
): Promise<AgentSession> {
const session = await this.createSdkSession(config);
return new PiDirectAgentSession(session, this.getModelRegistry(), config);
const { session, modelRegistry } = await this.createSdkSession(config);
return new PiDirectAgentSession(session, modelRegistry, config);
}
async resumeSession(
@@ -1434,37 +1440,34 @@ export class PiDirectAgentClient implements AgentClient {
modeId: overrides?.modeId,
};
const model = this.resolveConfiguredModel(mergedConfig.model);
const services = await this.getSessionServices(mergedConfig.cwd);
const model = this.resolveConfiguredModel(services.modelRegistry, mergedConfig.model);
const thinkingLevel = normalizePiThinkingOption(mergedConfig.thinkingOptionId);
const sessionOptions = {
cwd: mergedConfig.cwd,
modelRegistry: this.getModelRegistry(),
const { session } = await createAgentSessionFromServices({
services,
sessionManager: resumedManager,
...(model ? { model } : {}),
...(thinkingLevel ? { thinkingLevel } : {}),
};
const { session } = await createAgentSession(sessionOptions);
});
await session.bindExtensions({});
applySystemPrompt(session, mergedConfig.systemPrompt);
return new PiDirectAgentSession(session, this.getModelRegistry(), mergedConfig);
return new PiDirectAgentSession(session, services.modelRegistry, mergedConfig);
}
async listModels(_options: ListModelsOptions): Promise<AgentModelDefinition[]> {
// Pi Direct uses an in-process global registry; cwd/force are intentionally irrelevant.
const models = this.getModelRegistry()
.getAvailable()
.map((model) => ({
provider: PI_PROVIDER,
id: `${model.provider}/${model.id}`,
label: `${model.provider}/${model.name}`,
description: `${model.provider}/${model.id}`,
metadata: {
provider: model.provider,
modelId: model.id,
} satisfies AgentMetadata,
thinkingOptions: model.reasoning ? PI_THINKING_OPTIONS.map(mapThinkingOption) : undefined,
defaultThinkingOptionId: model.reasoning ? DEFAULT_PI_THINKING_LEVEL : undefined,
}));
async listModels(options: ListModelsOptions): Promise<AgentModelDefinition[]> {
const services = await this.getSessionServices(options.cwd);
const models = services.modelRegistry.getAvailable().map((model) => ({
provider: PI_PROVIDER,
id: `${model.provider}/${model.id}`,
label: `${model.provider}/${model.name}`,
description: `${model.provider}/${model.id}`,
metadata: {
provider: model.provider,
modelId: model.id,
} satisfies AgentMetadata,
thinkingOptions: model.reasoning ? PI_THINKING_OPTIONS.map(mapThinkingOption) : undefined,
defaultThinkingOptionId: model.reasoning ? DEFAULT_PI_THINKING_LEVEL : undefined,
}));
return transformPiModels(models);
}