feat(pi): materialize effective plugin worker tools

This commit is contained in:
2026-08-27 17:24:16 +08:00
parent c4dd8923a0
commit fd891ff3bb
12 changed files with 617 additions and 212 deletions

View File

@@ -1,6 +1,6 @@
// @vitest-environment node
import { mkdtemp, rm, writeFile } from 'node:fs/promises';
import { mkdtemp, readFile, rm, writeFile } from 'node:fs/promises';
import { tmpdir } from 'node:os';
import path from 'node:path';
import { pathToFileURL } from 'node:url';
@@ -11,6 +11,10 @@ import { PiProcessBudget } from '../../electron/coding-runtime/pi/worker-pool';
import type { AgentBrowserModule } from '../../electron/agent-browser';
import { CodingAttachmentStore } from '../../electron/coding-projects/attachment-store';
import { PiProductTools } from '../../electron/coding-runtime/pi/product-tools';
import {
DATA_SERVICE_PLUGIN_DEFINITION,
type CodingPluginToolDefinition,
} from '../../shared/coding-plugins';
type ExtensionHandler = (...arguments_: unknown[]) => Promise<unknown> | unknown;
type ExtensionTool = {
@@ -42,6 +46,83 @@ afterEach(async () => {
});
describe('Makelore Pi extension bundle', () => {
it('materializes only the frozen plugin declarations and lease metadata', async () => {
const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-dynamic-bundle-'));
roots.push(root);
const host = new PiManagedExtensionHost();
hosts.push(host);
const pluginTool: CodingPluginToolDefinition = {
name: 'data_service_inspect',
label: 'Data Service inspect',
description: 'Inspect the active project Data Service instance.',
capabilityId: 'data-service.control',
operation: 'inspect',
roles: ['parent'],
mutation: 'read',
projectWriteLease: true,
permissions: ['project.data.read'],
inputSchema: { type: 'object', additionalProperties: false, properties: {} },
};
const registration = await host.registerWorker({
conversationId: 'dynamic-conversation',
generation: 1,
projectId: 'dynamic-project',
projectPath: root,
extensionsDir: root,
skillEntries: [{ id: 'data-service', entryPath: 'skills/data-service/SKILL.md' }],
catalogRevision: 19,
tools: [pluginTool],
});
await host.bindRun('dynamic-conversation', 1, 'dynamic-run');
const previous = {
bridge: process.env.MAKELORE_PI_BRIDGE_URL,
token: process.env.MAKELORE_PI_WORKER_TOKEN,
context: process.env.MAKELORE_PI_CONTEXT_FILE,
role: process.env.MAKELORE_PI_WORKER_ROLE,
};
Object.assign(process.env, registration.env);
try {
const module = await import(
/* @vite-ignore */ `${pathToFileURL(registration.extensionPath).href}?dynamic=${Date.now()}`
) as {
default(factory: {
registerTool(tool: ExtensionTool): void;
on(event: string, handler: ExtensionHandler): void;
}): void | Promise<void>;
};
const tools = new Map<string, ExtensionTool>();
await module.default({
registerTool: (tool) => tools.set(tool.name, tool),
on: () => undefined,
});
await new Promise<void>((resolve) => setImmediate(resolve));
expect(tools.has('data_service_inspect')).toBe(true);
expect(tools.has('data_service_put_document')).toBe(false);
expect(tools.get('data_service_inspect')?.parameters).toEqual(pluginTool.inputSchema);
const context = JSON.parse(await readFile(registration.env.MAKELORE_PI_CONTEXT_FILE as string, 'utf8'));
expect(context).toMatchObject({
catalogRevision: 19,
allowedToolNames: ['data_service_inspect'],
projectWriteLeaseToolNames: ['data_service_inspect'],
});
const denied = await post(registration, {
action: 'product.invoke', conversationId: 'dynamic-conversation',
workerGeneration: 1, runId: 'dynamic-run', resourceId: 'denied-tool',
toolName: 'data_service_put_document', input: {},
});
expect(denied.status).toBe(403);
} finally {
for (const [key, value] of Object.entries(previous)) {
const environmentKey = key === 'bridge' ? 'MAKELORE_PI_BRIDGE_URL'
: key === 'token' ? 'MAKELORE_PI_WORKER_TOKEN'
: key === 'context' ? 'MAKELORE_PI_CONTEXT_FILE'
: 'MAKELORE_PI_WORKER_ROLE';
if (value === undefined) delete process.env[environmentKey];
else process.env[environmentKey] = value;
}
}
});
it('executes versioned product tools through the authenticated real bundle', async () => {
const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-product-bundle-'));
roots.push(root);
@@ -76,7 +157,10 @@ describe('Makelore Pi extension bundle', () => {
hosts.push(host);
const worker = await host.registerWorker({
conversationId: 'conversation-tools', generation: 1, projectId: 'project-a',
projectPath: root, skillIds: ['agent-browser'], extensionsDir: root,
projectPath: root,
skillEntries: [{ id: 'agent-browser', entryPath: 'agent-browser/SKILL.md' }],
catalogRevision: 3,
extensionsDir: root,
});
await host.bindRun('conversation-tools', 1, 'run-tools');
const previous = {
@@ -97,7 +181,7 @@ describe('Makelore Pi extension bundle', () => {
};
const tools = new Map<string, ExtensionTool>();
const handlers = new Map<string, ExtensionHandler>();
module.default({
await module.default({
registerTool: (tool) => tools.set(tool.name, tool),
on: (event, handler) => handlers.set(event, handler),
});
@@ -192,6 +276,8 @@ describe('Makelore Pi extension bundle', () => {
hosts.push(host);
const extensionWorker = await host.registerWorker({
conversationId: 'conversation-a1', generation: 1, projectId: 'project-a', extensionsDir: root,
tools: [...DATA_SERVICE_PLUGIN_DEFINITION.tools],
catalogRevision: 4,
});
const waitingWorker = await host.registerWorker({
conversationId: 'conversation-a2', generation: 1, projectId: 'project-a', extensionsDir: root,
@@ -217,7 +303,7 @@ describe('Makelore Pi extension bundle', () => {
};
const handlers = new Map<string, ExtensionHandler>();
const tools = new Map<string, ExtensionTool>();
module.default({
await module.default({
registerTool: (tool) => tools.set(tool.name, tool),
on: (event, handler) => handlers.set(event, handler),
});
@@ -237,21 +323,9 @@ describe('Makelore Pi extension bundle', () => {
];
for (const name of dataServiceTools) {
const tool = tools.get(name);
expect(tool?.parameters).toMatchObject({
type: 'object', additionalProperties: false,
});
expect(Object.keys(tool?.parameters?.properties ?? {})).not.toEqual(
expect.arrayContaining(['owner', 'project', 'path', 'token', 'endpoint', 'url', 'handle']),
expect(tool?.parameters).toEqual(
DATA_SERVICE_PLUGIN_DEFINITION.tools.find(({ name: candidate }) => candidate === name)?.inputSchema,
);
const properties = tool?.parameters?.properties as Record<string, Record<string, unknown>>;
if (name === 'data_service_get_document'
|| name === 'data_service_put_document'
|| name === 'data_service_delete_document') {
expect(properties.document_id?.not).toEqual({ enum: ['.', '..'] });
}
if (name === 'data_service_put_document' || name === 'data_service_delete_document') {
expect(properties.if_revision?.maximum).toBe(Number.MAX_SAFE_INTEGER);
}
}
const updates: unknown[] = [];
@@ -355,7 +429,7 @@ describe('Makelore Pi extension bundle', () => {
};
const tools: string[] = [];
const handlers = new Map<string, ExtensionHandler>();
module.default({
await module.default({
registerTool: (tool) => tools.push(tool.name),
on: (event, handler) => handlers.set(event, handler),
});

View File

@@ -107,6 +107,7 @@ describe('managed Pi extension bridge', () => {
)) as Record<string, unknown>;
expect(context).toEqual({
conversationId: 'conversation-a', workerGeneration: 2, role: 'parent', runId: 'run-a',
skillIds: [], allowedToolNames: [], tools: [], projectWriteLeaseToolNames: [],
});
await first.dispose();
const response = await post(replacement, {

View File

@@ -3,6 +3,7 @@ import { tmpdir } from 'node:os';
import path from 'node:path';
import { afterEach, describe, expect, it } from 'vitest';
import YAML from 'yaml';
import { CodingCapabilityRegistryImpl } from '@electron/coding-plugins/registry';
import type { PiProviderSelection } from '@electron/coding-runtime/pi/provider-config';
import {
buildPiManagedInputArgs,
@@ -11,6 +12,8 @@ import {
resolveBundledCodingSkillsDir,
resolveExplicitCodingSkillPaths,
} from '@electron/coding-runtime/pi/resource-loader';
import type { PluginPolicyClientState } from '@electron/services/plugin-policy-client';
import { DATA_SERVICE_PLUGIN_DEFINITION } from '../../shared/coding-plugins';
const temporaryRoots: string[] = [];
@@ -34,12 +37,18 @@ async function fixtureRoot(): Promise<{
mkdir(path.join(projectDir, '.agents', 'skills', 'untrusted-agent-skill'), { recursive: true }),
mkdir(path.join(skillsDir, 'grilling'), { recursive: true }),
mkdir(path.join(skillsDir, 'agent-browser'), { recursive: true }),
mkdir(path.join(path.dirname(skillsDir), 'coding-plugins', 'data-service', 'skills', 'data-service'), { recursive: true }),
]);
await Promise.all([
writeFile(path.join(projectDir, '.pi', 'skills', 'untrusted-project-skill', 'SKILL.md'), 'untrusted', 'utf8'),
writeFile(path.join(projectDir, '.agents', 'skills', 'untrusted-agent-skill', 'SKILL.md'), 'untrusted', 'utf8'),
writeFile(path.join(skillsDir, 'grilling', 'SKILL.md'), '---\nname: grilling\n---\n', 'utf8'),
writeFile(path.join(skillsDir, 'agent-browser', 'SKILL.md'), '---\nname: agent-browser\n---\n', 'utf8'),
writeFile(
path.join(path.dirname(skillsDir), 'coding-plugins', 'data-service', 'skills', 'data-service', 'SKILL.md'),
'---\nname: data-service\n---\n',
'utf8',
),
]);
return { root, userDataDir, projectDir, skillsDir };
}
@@ -53,13 +62,21 @@ describe('Pi managed resource loader', () => {
projectId: 'project-1',
agentId: 'agent-1',
prompt,
skillIds: ['grilling', 'grilling'],
skillEntries: [
{ id: 'grilling', entryPath: 'grilling/SKILL.md' },
{ id: 'grilling', entryPath: 'grilling/SKILL.md' },
],
catalogRevision: 11,
bundledSkillsDir: fixture.skillsDir,
revision: { provider: 3, resources: 7 },
});
expect(await readFile(resources.promptPath, 'utf8')).toBe(prompt);
expect(resources.skillIds).toEqual(['grilling']);
expect(resources.skillEntries).toEqual([
{ id: 'grilling', entryPath: 'grilling/SKILL.md' },
]);
expect(resources.catalogRevision).toBe(11);
expect(resources.skillPaths).toEqual([path.join(fixture.skillsDir, 'grilling', 'SKILL.md')]);
expect(JSON.stringify(resources.summary)).not.toContain(prompt);
const manifest = JSON.parse(await readFile(resources.manifestPath, 'utf8')) as Record<string, unknown>;
@@ -69,6 +86,8 @@ describe('Pi managed resource loader', () => {
agentId: 'agent-1',
promptFile: 'agent-1.md',
skillIds: ['grilling'],
skillEntries: [{ id: 'grilling', entryPath: 'grilling/SKILL.md' }],
catalogRevision: 11,
revision: { provider: 3, resources: 7 },
});
expect(JSON.stringify(manifest)).not.toContain(prompt);
@@ -84,7 +103,8 @@ describe('Pi managed resource loader', () => {
projectId: 'project-1',
agentId: 'agent-1',
prompt,
skillIds: ['agent-browser'],
skillEntries: [{ id: 'agent-browser', entryPath: 'agent-browser/SKILL.md' }],
catalogRevision: 12,
bundledSkillsDir: fixture.skillsDir,
revision: { provider: 1, resources: 1 },
});
@@ -111,16 +131,112 @@ describe('Pi managed resource loader', () => {
expect(serialized).not.toContain(path.join(fixture.projectDir, '.agents'));
});
it('materializes effective plugin Skill entries from their bundled package roots', async () => {
const fixture = await fixtureRoot();
const resources = await materializePiAgentResources({
userDataDir: fixture.userDataDir,
projectId: 'project-1',
agentId: 'agent-1',
prompt: 'plugin prompt',
skillEntries: [{ id: 'data-service', entryPath: 'skills/data-service/SKILL.md' }],
catalogRevision: 13,
bundledSkillsDir: fixture.skillsDir,
revision: { provider: 1, resources: 1 },
});
expect(resources.skillIds).toEqual(['data-service']);
expect(resources.skillPaths).toEqual([
path.join(path.dirname(fixture.skillsDir), 'coding-plugins', 'data-service', 'skills', 'data-service', 'SKILL.md'),
]);
expect(resources.catalogRevision).toBe(13);
});
it('filters a known disabled plugin Skill for the next worker and restores it after re-enable', async () => {
const fixture = await fixtureRoot();
const assignedSkillIds = ['grilling', 'data-service'];
let enabled = true;
const policy: PluginPolicyClientState = {
status: 'current',
revision: 17,
lastVerifiedAt: 1,
catalog: {
schema_version: 1,
catalog_version: 'catalog-17',
pricing_version: null,
plugins: [{
plugin_id: DATA_SERVICE_PLUGIN_DEFINITION.id,
supported_contract_versions: [DATA_SERVICE_PLUGIN_DEFINITION.contractVersion],
status: 'active',
capabilities: DATA_SERVICE_PLUGIN_DEFINITION.skills[0].grants.map((capabilityId) => ({
capability_id: capabilityId,
operations: DATA_SERVICE_PLUGIN_DEFINITION.tools
.filter((tool) => tool.capabilityId === capabilityId)
.map((tool) => ({
operation: tool.operation,
billing: { mode: 'included' as const, entitlement_scope: null, notice: 'Included' },
})),
})),
}],
},
};
const registry = new CodingCapabilityRegistryImpl({
policyClient: { getState: () => policy },
getEnabledPluginIds: async () => enabled ? [DATA_SERVICE_PLUGIN_DEFINITION.id] : [],
adapters: [{
pluginId: DATA_SERVICE_PLUGIN_DEFINITION.id,
async inspect() { return { status: 'ready' }; },
async invoke() {
return {
success: true, status: 200, code: null, error: null, retryable: false,
payload_schema: 'data-service.v1', data: null,
};
},
}],
});
const resolve = async () => await registry.resolveWorkerResources({
projectPath: fixture.projectDir,
assignedSkillIds,
role: 'parent',
});
const first = await resolve();
expect(first.effectiveSkillIds).toEqual(assignedSkillIds);
expect(first.tools).toHaveLength(10);
enabled = false;
const disabled = await resolve();
expect(assignedSkillIds).toEqual(['grilling', 'data-service']);
expect(disabled.effectiveSkillIds).toEqual(['grilling']);
expect(disabled.tools).toEqual([]);
await expect(materializePiAgentResources({
userDataDir: fixture.userDataDir,
projectId: 'project-1',
agentId: 'agent-1',
prompt: 'disabled plugin prompt',
skillEntries: disabled.skillEntries,
catalogRevision: disabled.catalogRevision,
bundledSkillsDir: fixture.skillsDir,
revision: { provider: 1, resources: 2 },
})).resolves.toMatchObject({ skillIds: ['grilling'], catalogRevision: 17 });
enabled = true;
const restored = await resolve();
expect(restored.effectiveSkillIds).toEqual(assignedSkillIds);
expect(restored.tools).toHaveLength(10);
});
it('rejects unknown skills and unsafe managed path segments', async () => {
const fixture = await fixtureRoot();
await expect(resolveExplicitCodingSkillPaths(fixture.skillsDir, ['not-bundled']))
.rejects.toThrow('Unknown bundled coding skill');
await expect(resolveExplicitCodingSkillPaths(fixture.skillsDir, [
{ id: 'not-bundled', entryPath: 'not-bundled/SKILL.md' },
])).rejects.toThrow('Bundled coding Skill entry is not a file');
await expect(materializePiAgentResources({
userDataDir: fixture.userDataDir,
projectId: '../outside',
agentId: 'agent-1',
prompt: '',
skillIds: [],
skillEntries: [],
catalogRevision: 1,
bundledSkillsDir: fixture.skillsDir,
revision: { provider: 1, resources: 1 },
})).rejects.toThrow('Project id');

View File

@@ -244,7 +244,7 @@ describe('Pi worker process', () => {
'--no-context-files',
'--no-approve',
'--tools',
'read,bash,edit,write,grep,find,ls,ask_user,subagent,agent_browser,game_asset_browser,game_asset_review,task_state,changed_file,runtime_context,data_service_configure,data_service_inspect,data_service_list_projects,data_service_get_document,data_service_list_documents,data_service_put_document,data_service_delete_document,data_service_remove_collection,data_service_reset,data_service_remove_project',
'read,bash,edit,write,grep,find,ls,ask_user,subagent,agent_browser,game_asset_browser,game_asset_review,task_state,changed_file,runtime_context',
'--model', 'model-a',
]);
expect(buildPiRpcArgs('sessions', ['--no-session'], ['read', 'grep', 'find', 'ls']))

View File

@@ -6,8 +6,12 @@ import { mkdir, mkdtemp, readFile, rm, writeFile } from 'node:fs/promises';
import { createRequire } from 'node:module';
import { tmpdir } from 'node:os';
import { join, resolve } from 'node:path';
import { PiWorkerProcess } from '../../electron/coding-runtime/pi/worker-process';
import {
PI_CORE_TOOL_NAMES,
PiWorkerProcess,
} from '../../electron/coding-runtime/pi/worker-process';
import { PiManagedExtensionHost } from '../../electron/coding-runtime/pi/extension-host';
import { DATA_SERVICE_PLUGIN_DEFINITION } from '../../shared/coding-plugins';
import {
PI_RUNTIME_MANIFEST,
PI_RUNTIME_VERSION,
@@ -85,6 +89,8 @@ describe('locked Pi worker process smoke', () => {
generation: 1,
projectId: 'real-project',
extensionsDir: join(root, 'extensions'),
catalogRevision: 7,
tools: [...DATA_SERVICE_PLUGIN_DEFINITION.tools],
});
await extensionHost.bindRun('real-conversation', 1, 'real-run');
const probe = await materializeActiveToolsProbe(root);
@@ -94,6 +100,10 @@ describe('locked Pi worker process smoke', () => {
cwd,
configDir,
sessionDir,
tools: [
...PI_CORE_TOOL_NAMES,
...DATA_SERVICE_PLUGIN_DEFINITION.tools.map(({ name }) => name),
],
additionalArgs: [
'--extension', extension.extensionPath,
'--extension', probe.extensionPath,
@@ -118,6 +128,7 @@ describe('locked Pi worker process smoke', () => {
'task_state',
'changed_file',
'runtime_context',
...DATA_SERVICE_PLUGIN_DEFINITION.tools.map(({ name }) => name),
]));
await expect(worker.stop('test_injection')).resolves.toMatchObject({ mode: 'stdin-close', code: 0 });
} finally {