feat(coding): add model tools and device packages

This commit is contained in:
2026-09-02 18:35:26 +08:00
parent 841273b433
commit 664b8823a0
60 changed files with 3847 additions and 2068 deletions

View File

@@ -40,6 +40,8 @@ interface WorkerRegistrationRecord {
projectWriteLeaseToolNames: string[];
role: 'parent' | 'child';
effectiveSnapshot?: EffectivePluginSnapshot;
devicePackageGeneration?: number;
devicePackageIds?: string[];
contextFile: string;
runId: string | null;
leases: Map<string, PiProjectWriteLease>;
@@ -64,6 +66,8 @@ export interface RegisterPiExtensionWorkerInput {
tools?: readonly CodingPluginToolDefinition[];
/** Exact Main-owned resolver output used for this worker generation. */
effectiveSnapshot?: EffectivePluginSnapshot;
devicePackageGeneration?: number;
devicePackageIds?: readonly string[];
extensionsDir: string;
role?: 'parent' | 'child';
runId?: string;
@@ -252,6 +256,10 @@ export class PiManagedExtensionHost {
projectWriteLeaseToolNames,
role,
...(input.effectiveSnapshot ? { effectiveSnapshot: input.effectiveSnapshot } : {}),
...(input.devicePackageGeneration === undefined
? {}
: { devicePackageGeneration: input.devicePackageGeneration }),
...(input.devicePackageIds ? { devicePackageIds: [...input.devicePackageIds] } : {}),
contextFile,
runId: role === 'child'
? input.runId as string
@@ -424,6 +432,7 @@ export class PiManagedExtensionHost {
}
const productResult = await this.productTools.execute(value.toolName, {
conversationId: record.conversationId,
workerGeneration: record.generation,
runId: value.runId,
resourceId: value.resourceId,
projectId: record.projectId,
@@ -596,6 +605,10 @@ export class PiManagedExtensionHost {
skillIds: record.skillIds,
...(record.catalogRevision === undefined ? {} : { catalogRevision: record.catalogRevision }),
...(record.effectiveSnapshot ? { effectivePluginSnapshot: record.effectiveSnapshot } : {}),
...(record.devicePackageGeneration === undefined
? {}
: { devicePackageGeneration: record.devicePackageGeneration }),
...(record.devicePackageIds ? { devicePackageIds: record.devicePackageIds } : {}),
allowedToolNames: record.allowedToolNames,
tools: record.tools,
projectWriteLeaseToolNames: record.projectWriteLeaseToolNames,

View File

@@ -0,0 +1,255 @@
import type { CodingPluginToolDefinition } from '../../../../shared/coding-plugins';
import type { ModelToolDetailsV1, ModelWebSearchFailureV1 } from '../../../../shared/model-tools';
import type { ProviderAccount } from '../../../shared/providers/types';
import type {
PiProviderDescriptor,
PiProviderSelection,
PiWorkerCredentialProjection,
} from '../provider-config';
import {
ModelWebSearchError,
createModelWebSearchAdapter,
type FrozenSelectedModel,
type ModelWebSearchAdapter,
} from './web-search';
const WEB_SEARCH_TOOL: CodingPluginToolDefinition = Object.freeze({
name: 'web_search',
label: 'Web search',
description: [
'Search the public internet with the currently selected model for current information, facts, and sources.',
'Use agent_browser only for interactive browser debugging, authenticated pages, and UI actions.',
'When this tool fails, report the failure. Do not use agent_browser as a fallback.',
].join(' '),
capabilityId: 'model.web-search',
operation: 'search',
roles: ['parent'],
mutation: 'read',
projectWriteLease: false,
permissions: [],
inputSchema: Object.freeze({
type: 'object',
additionalProperties: false,
required: ['query'],
properties: {
query: { type: 'string', minLength: 1, maxLength: 2_000 },
},
}),
});
export interface ModelToolInvocationContext {
conversationId: string;
workerGeneration: number;
runId: string;
resourceId: string;
}
export interface ModelToolInvocationResult {
content: Array<{ type: 'text'; text: string }>;
details: ModelToolDetailsV1;
}
export interface RegisterModelToolWorkerInput {
conversationId: string;
generation: number;
account: ProviderAccount;
descriptor: PiProviderDescriptor;
selection: PiProviderSelection;
credential: PiWorkerCredentialProjection;
}
export interface ModelToolWorkerRegistration {
tools: readonly CodingPluginToolDefinition[];
dispose(): void;
}
export interface ModelToolRegistryOptions {
adapter?: ModelWebSearchAdapter;
}
export interface ModelToolRegistryPort {
registerWorker(input: RegisterModelToolWorkerInput): ModelToolWorkerRegistration;
invoke(
toolName: string,
context: ModelToolInvocationContext,
input: unknown,
signal?: AbortSignal,
): Promise<ModelToolInvocationResult>;
}
interface WorkerBinding {
selectedModel: FrozenSelectedModel;
}
function resolvedHeaderValue(
value: string,
credential: PiWorkerCredentialProjection,
): string | null {
if (!value.startsWith('$')) return value;
return credential.env[value.slice(1)]?.trim() || null;
}
function freezeSelectedModel(input: RegisterModelToolWorkerInput): FrozenSelectedModel | null {
if (
input.account.id !== input.selection.accountId
|| input.descriptor.accountId !== input.selection.accountId
|| input.descriptor.runtimeProviderId !== input.selection.runtimeProviderId
) {
return null;
}
const capability = input.account.metadata
?.worksSquareModelCapabilities
?.[input.selection.modelId]
?.webSearch;
if (!capability || !input.descriptor.baseUrl) return null;
const expectedApi = capability.adapter === 'bailian-chat-completions'
? 'openai-completions'
: 'openai-responses';
if (input.descriptor.api !== expectedApi) return null;
const headers: Record<string, string> = {};
for (const [name, rawValue] of Object.entries(input.descriptor.headers)) {
const value = resolvedHeaderValue(rawValue, input.credential);
if (!value) return null;
headers[name] = value;
}
const hasAuthorization = Object.keys(headers).some((name) => name.toLowerCase() === 'authorization');
const apiKey = input.descriptor.apiKeyEnv
? input.credential.env[input.descriptor.apiKeyEnv]?.trim()
: undefined;
if (!hasAuthorization && apiKey) headers.Authorization = `Bearer ${apiKey}`;
return Object.freeze({
accountId: input.account.id,
runtimeProviderId: input.selection.runtimeProviderId,
modelId: input.selection.modelId,
generation: input.generation,
baseUrl: input.descriptor.baseUrl,
headers: Object.freeze(headers),
capability: Object.freeze({ ...capability }),
});
}
function failure(
modelId: string,
error: ModelWebSearchError,
): ModelToolInvocationResult {
const details: ModelWebSearchFailureV1 = {
schema: 'makelore-model-tool.v1',
tool: 'web_search',
status: 'failed',
modelId,
error: {
code: error.code,
message: error.message,
httpStatus: error.status,
retryable: error.retryable,
},
};
return {
content: [{ type: 'text', text: `${error.code}: ${error.message}` }],
details,
};
}
function queryFromInput(value: unknown): string | null {
if (!value || typeof value !== 'object' || Array.isArray(value)) return null;
const record = value as Record<string, unknown>;
if (Object.keys(record).some((key) => key !== 'query')) return null;
const query = typeof record.query === 'string' ? record.query.trim() : '';
return query && query.length <= 2_000 ? query : null;
}
function successText(answer: string, sources: readonly { title: string; url: string }[]): string {
if (sources.length === 0) return answer;
return `${answer}\n\nSources:\n${sources.map(({ title, url }) => `- ${title}: ${url}`).join('\n')}`;
}
export class ModelToolRegistry implements ModelToolRegistryPort {
private readonly adapter: ModelWebSearchAdapter;
private readonly bindings = new Map<string, WorkerBinding>();
constructor(options: ModelToolRegistryOptions = {}) {
this.adapter = options.adapter ?? createModelWebSearchAdapter();
}
registerWorker(input: RegisterModelToolWorkerInput): ModelToolWorkerRegistration {
const selectedModel = freezeSelectedModel(input);
if (!selectedModel) return { tools: [], dispose: () => undefined };
const binding = { selectedModel };
this.bindings.set(input.conversationId, binding);
return {
tools: [WEB_SEARCH_TOOL],
dispose: () => {
if (this.bindings.get(input.conversationId) === binding) {
this.bindings.delete(input.conversationId);
}
},
};
}
async invoke(
toolName: string,
context: ModelToolInvocationContext,
input: unknown,
signal = new AbortController().signal,
): Promise<ModelToolInvocationResult> {
const binding = this.bindings.get(context.conversationId);
const modelId = binding?.selectedModel.modelId ?? 'unknown';
if (toolName !== 'web_search') {
return failure(modelId, new ModelWebSearchError(
'model_web_search_unsupported',
400,
false,
'The requested model tool is unavailable',
));
}
if (!binding || binding.selectedModel.generation !== context.workerGeneration) {
return failure(modelId, new ModelWebSearchError(
'model_context_changed',
409,
false,
'The selected model context changed before Web Search completed',
));
}
const query = queryFromInput(input);
if (!query) {
return failure(modelId, new ModelWebSearchError(
'model_web_search_invalid_result',
502,
false,
'Web Search query must contain between 1 and 2000 characters',
));
}
try {
const details = await this.adapter.search({
query,
selectedModel: binding.selectedModel,
parentTurnId: context.runId,
toolCallId: context.resourceId,
}, signal);
if (this.bindings.get(context.conversationId) !== binding) {
return failure(modelId, new ModelWebSearchError(
'model_context_changed',
409,
false,
'The selected model context changed while Web Search was running',
));
}
return {
content: [{ type: 'text', text: successText(details.answer, details.sources) }],
details,
};
} catch (error) {
const normalized = error instanceof ModelWebSearchError
? error
: new ModelWebSearchError(
'model_web_search_unavailable',
503,
true,
'The selected model Web Search transport is unavailable',
);
return failure(modelId, normalized);
}
}
}

View File

@@ -0,0 +1,326 @@
import type { ImportedModelWebSearchCapability } from '../../../../shared/imported-model-profile';
import type {
ModelWebSearchErrorCode,
ModelWebSearchSuccessV1,
} from '../../../../shared/model-tools';
import { proxyAwareFetch, runWithDeadline } from '../../../utils/proxy-fetch';
const MAX_QUERY_CHARS = 2_000;
const MAX_ANSWER_CHARS = 20_000;
const MAX_SOURCE_COUNT = 20;
const MAX_SOURCE_TITLE_CHARS = 240;
const MAX_SOURCE_URL_CHARS = 2_048;
const MAX_RESPONSE_BYTES = 1_048_576;
const DEFAULT_TIMEOUT_MS = 120_000;
type FetchImplementation = (
input: string | URL,
init?: RequestInit,
) => Promise<Response>;
export type FrozenSelectedModel = Readonly<{
accountId: string;
runtimeProviderId: string;
modelId: string;
generation: number;
baseUrl: string;
headers: Readonly<Record<string, string>>;
capability?: ImportedModelWebSearchCapability;
}>;
export class ModelWebSearchError extends Error {
constructor(
public readonly code: ModelWebSearchErrorCode,
public readonly status: 400 | 409 | 429 | 502 | 503,
public readonly retryable: boolean,
message: string,
) {
super(message);
this.name = 'ModelWebSearchError';
}
}
export interface ModelWebSearchAdapter {
search(input: Readonly<{
query: string;
selectedModel: FrozenSelectedModel;
parentTurnId: string;
toolCallId: string;
}>, signal: AbortSignal): Promise<ModelWebSearchSuccessV1>;
}
export interface CreateModelWebSearchAdapterOptions {
fetchImpl?: FetchImplementation;
timeoutMs?: number;
}
function asRecord(value: unknown): Record<string, unknown> | null {
return value && typeof value === 'object' && !Array.isArray(value)
? value as Record<string, unknown>
: null;
}
function endpoint(baseUrl: string, pathname: 'responses' | 'chat/completions'): string {
return `${baseUrl.trim().replace(/\/+$/, '')}/${pathname}`;
}
function requestHeaders(headers: Readonly<Record<string, string>>): Record<string, string> {
const result = { ...headers };
const contentType = Object.keys(result).find((name) => name.toLowerCase() === 'content-type');
if (contentType) delete result[contentType];
result['Content-Type'] = 'application/json';
return result;
}
async function readBoundedJson(response: Response): Promise<unknown> {
const contentLength = Number(response.headers.get('content-length'));
if (Number.isFinite(contentLength) && contentLength > MAX_RESPONSE_BYTES) {
throw new ModelWebSearchError(
'model_web_search_invalid_result',
502,
false,
'The selected model returned an oversized Web Search response',
);
}
if (!response.body) {
throw new ModelWebSearchError(
'model_web_search_invalid_result',
502,
false,
'The selected model returned an empty Web Search response',
);
}
const reader = response.body.getReader();
const chunks: Uint8Array[] = [];
let byteLength = 0;
try {
while (true) {
const { value, done } = await reader.read();
if (done) break;
if (!value) continue;
byteLength += value.byteLength;
if (byteLength > MAX_RESPONSE_BYTES) {
await reader.cancel();
throw new ModelWebSearchError(
'model_web_search_invalid_result',
502,
false,
'The selected model returned an oversized Web Search response',
);
}
chunks.push(value);
}
} finally {
reader.releaseLock();
}
const bytes = new Uint8Array(byteLength);
let offset = 0;
for (const chunk of chunks) {
bytes.set(chunk, offset);
offset += chunk.byteLength;
}
try {
return JSON.parse(new TextDecoder().decode(bytes)) as unknown;
} catch {
throw new ModelWebSearchError(
'model_web_search_invalid_result',
502,
false,
'The selected model returned an invalid Web Search response',
);
}
}
function normalizedSourceUrl(rawUrl: unknown): { identity: string; url: string } | null {
if (typeof rawUrl !== 'string') return null;
const url = rawUrl.trim();
if (!url || url.length > MAX_SOURCE_URL_CHARS) return null;
try {
const parsed = new URL(url);
if ((parsed.protocol !== 'http:' && parsed.protocol !== 'https:')
|| parsed.username || parsed.password) {
return null;
}
const identity = new URL(parsed.toString());
identity.hash = '';
return { identity: identity.toString(), url: parsed.toString() };
} catch {
return null;
}
}
function sourceCandidate(value: unknown): { title: string; url: string; identity: string } | null {
const record = asRecord(value);
const normalized = normalizedSourceUrl(record?.url);
if (!record || !normalized) return null;
const rawTitle = typeof record.title === 'string' ? record.title.trim() : '';
let title = rawTitle || new URL(normalized.url).hostname;
if (!title || title.length > MAX_SOURCE_TITLE_CHARS) return null;
title = title.slice(0, MAX_SOURCE_TITLE_CHARS);
return { title, url: normalized.url, identity: normalized.identity };
}
function collectSources(payload: Record<string, unknown>): Array<{ title: string; url: string }> {
const candidates: unknown[] = [];
const choices = Array.isArray(payload.choices) ? payload.choices : [];
for (const choice of choices) {
const message = asRecord(asRecord(choice)?.message);
if (Array.isArray(message?.sources)) candidates.push(...message.sources);
}
const searchInfo = asRecord(payload.search_info ?? payload.searchInfo);
if (Array.isArray(searchInfo?.search_results)) candidates.push(...searchInfo.search_results);
if (Array.isArray(searchInfo?.searchResults)) candidates.push(...searchInfo.searchResults);
const output = Array.isArray(payload.output) ? payload.output : [];
for (const item of output) {
const outputItem = asRecord(item);
const action = asRecord(outputItem?.action);
if (Array.isArray(action?.sources)) candidates.push(...action.sources);
const content = Array.isArray(outputItem?.content) ? outputItem.content : [];
for (const part of content) {
const annotations = Array.isArray(asRecord(part)?.annotations)
? asRecord(part)?.annotations as unknown[]
: [];
for (const annotation of annotations) {
const annotationRecord = asRecord(annotation);
if (annotationRecord?.type === 'url_citation') candidates.push(annotationRecord);
}
}
}
const seen = new Set<string>();
const sources: Array<{ title: string; url: string }> = [];
for (const candidate of candidates) {
const source = sourceCandidate(candidate);
if (!source || seen.has(source.identity)) continue;
seen.add(source.identity);
sources.push({ title: source.title, url: source.url });
if (sources.length === MAX_SOURCE_COUNT) break;
}
return sources;
}
function answerFromPayload(payload: Record<string, unknown>): string | null {
if (typeof payload.output_text === 'string' && payload.output_text.trim()) {
return payload.output_text.trim();
}
const choices = Array.isArray(payload.choices) ? payload.choices : [];
for (const choice of choices) {
const content = asRecord(asRecord(choice)?.message)?.content;
if (typeof content === 'string' && content.trim()) return content.trim();
}
const parts: string[] = [];
const output = Array.isArray(payload.output) ? payload.output : [];
for (const item of output) {
const content = Array.isArray(asRecord(item)?.content) ? asRecord(item)?.content as unknown[] : [];
for (const part of content) {
const text = asRecord(part)?.text;
if (typeof text === 'string' && text.trim()) parts.push(text.trim());
}
}
return parts.length > 0 ? parts.join('\n') : null;
}
function mapHttpError(response: Response): ModelWebSearchError {
if (response.status === 429) {
return new ModelWebSearchError(
'model_web_search_rate_limited',
429,
false,
'The selected model rate-limited Web Search',
);
}
return new ModelWebSearchError(
'model_web_search_unavailable',
503,
true,
'The selected model Web Search transport is unavailable',
);
}
export function createModelWebSearchAdapter(
options: CreateModelWebSearchAdapterOptions = {},
): ModelWebSearchAdapter {
const fetchImpl = options.fetchImpl ?? proxyAwareFetch;
const timeoutMs = options.timeoutMs ?? DEFAULT_TIMEOUT_MS;
return {
async search(input, signal) {
const query = input.query.trim();
if (!query || query.length > MAX_QUERY_CHARS) {
throw new ModelWebSearchError(
'model_web_search_invalid_result',
502,
false,
'Web Search query must contain between 1 and 2000 characters',
);
}
const capability = input.selectedModel.capability;
if (!capability || capability.supportsForcedSearch !== true) {
throw new ModelWebSearchError(
'model_web_search_unsupported',
400,
false,
'The selected model does not support Web Search',
);
}
const isResponses = capability.adapter === 'openai-responses'
|| capability.adapter === 'bailian-responses';
const url = endpoint(input.selectedModel.baseUrl, isResponses ? 'responses' : 'chat/completions');
const body = isResponses
? {
model: input.selectedModel.modelId,
input: query,
tools: [{ type: 'web_search' }],
tool_choice: 'required',
}
: {
model: input.selectedModel.modelId,
messages: [{ role: 'user', content: query }],
stream: false,
enable_search: true,
search_options: { forced_search: true },
};
let response: Response;
try {
response = await runWithDeadline(
async (deadlineSignal) => await fetchImpl(url, {
method: 'POST',
headers: requestHeaders(input.selectedModel.headers),
body: JSON.stringify(body),
signal: deadlineSignal,
}),
timeoutMs,
signal,
);
} catch (error) {
if (error instanceof ModelWebSearchError) throw error;
throw new ModelWebSearchError(
'model_web_search_unavailable',
503,
true,
'The selected model Web Search transport is unavailable',
);
}
if (!response.ok) throw mapHttpError(response);
const payload = asRecord(await readBoundedJson(response));
const answer = payload ? answerFromPayload(payload) : null;
if (!payload || !answer || answer.length > MAX_ANSWER_CHARS) {
throw new ModelWebSearchError(
'model_web_search_invalid_result',
502,
false,
'The selected model returned an invalid Web Search result',
);
}
return {
schema: 'makelore-model-tool.v1',
tool: 'web_search',
status: 'succeeded',
modelId: input.selectedModel.modelId,
answer,
sources: collectSources(payload),
sourceMode: capability.sourceMode,
};
},
};
}

View File

@@ -25,6 +25,9 @@ import { BUNDLED_CODING_SKILL_IDS } from '../../../shared/coding-skills';
import { PiAgentBrowserTool } from './extensions/agent-browser';
import { reportChangedFiles } from './extensions/changed-file';
import { projectTaskState } from './extensions/task-state';
import type { ModelToolRegistryPort } from './model-tools/model-tool-registry';
import type { DevicePackageTools } from '../../coding-packages/device-package-tools';
import { DEVICE_PACKAGE_TOOL_NAMES } from '../../../shared/device-packages';
export type PiProductToolName =
| 'agent_browser'
@@ -45,6 +48,7 @@ export function isPiProductToolName(value: unknown): value is PiProductToolName
export interface PiProductToolContext {
conversationId: string;
workerGeneration?: number;
runId: string;
resourceId: string;
projectId: string;
@@ -67,6 +71,8 @@ export interface PiProductToolsOptions {
getPluginSkillSources?(): readonly ProductCodingPluginSkillSource[]
| Promise<readonly ProductCodingPluginSkillSource[]>;
capabilityRegistry?: CodingCapabilityRegistry;
modelToolRegistry?: ModelToolRegistryPort;
devicePackageTools?: DevicePackageTools;
}
export class PiProductTools {
@@ -143,6 +149,19 @@ export class PiProductTools {
context: PiProductToolContext,
input: unknown,
): Promise<PiProductToolResult> {
if (DEVICE_PACKAGE_TOOL_NAMES.includes(toolName as typeof DEVICE_PACKAGE_TOOL_NAMES[number])) {
if (!this.options.devicePackageTools) throw new Error('Device package management is unavailable');
return await this.options.devicePackageTools.invoke(toolName, context.runId, input);
}
if (toolName === 'web_search') {
if (!this.options.modelToolRegistry) throw new Error('Model Web Search is unavailable');
return await this.options.modelToolRegistry.invoke('web_search', {
conversationId: context.conversationId,
workerGeneration: context.workerGeneration ?? 0,
runId: context.runId,
resourceId: context.resourceId,
}, input);
}
if (toolName === 'agent_browser') {
return await this.browser.execute(context, input);
}

View File

@@ -51,6 +51,9 @@ export interface MaterializePiAgentResourcesOptions {
skillRoots?: readonly string[];
/** The exact resolver output for this worker generation. */
effectiveSnapshot?: EffectivePluginSnapshot;
/** Device-installed packages resolved once for this parent worker generation. */
devicePackageGeneration?: number;
devicePackageIds?: readonly string[];
revision: PiManagedInputRevision;
}
@@ -65,6 +68,8 @@ export interface PiAgentResourceManifest {
catalogRevision: number;
revision: PiManagedInputRevision;
effectivePluginSnapshot?: EffectivePluginSnapshot;
devicePackageGeneration?: number;
devicePackageIds?: string[];
}
export interface PiAgentResourceSnapshot {
@@ -79,6 +84,8 @@ export interface PiAgentResourceSnapshot {
catalogRevision: number;
revision: PiManagedInputRevision;
effectivePluginSnapshot?: EffectivePluginSnapshot;
devicePackageGeneration?: number;
devicePackageIds?: string[];
summary: {
projectId: string;
agentId: string;
@@ -87,6 +94,8 @@ export interface PiAgentResourceSnapshot {
catalogRevision: number;
revision: PiManagedInputRevision;
effectivePluginSnapshot?: EffectivePluginSnapshot;
devicePackageGeneration?: number;
devicePackageIds?: string[];
};
}
@@ -285,6 +294,10 @@ export async function materializePiAgentResources(
catalogRevision: resolvedCatalogRevision,
revision: { ...options.revision },
...(options.effectiveSnapshot ? { effectivePluginSnapshot: structuredClone(options.effectiveSnapshot) } : {}),
...(options.devicePackageGeneration === undefined
? {}
: { devicePackageGeneration: catalogRevision(options.devicePackageGeneration) }),
...(options.devicePackageIds ? { devicePackageIds: [...options.devicePackageIds] } : {}),
};
await Promise.all([
atomicWriteText(promptPath, options.prompt),
@@ -303,6 +316,10 @@ export async function materializePiAgentResources(
catalogRevision: resolvedCatalogRevision,
revision: { ...options.revision },
...(options.effectiveSnapshot ? { effectivePluginSnapshot: structuredClone(options.effectiveSnapshot) } : {}),
...(options.devicePackageGeneration === undefined
? {}
: { devicePackageGeneration: catalogRevision(options.devicePackageGeneration) }),
...(options.devicePackageIds ? { devicePackageIds: [...options.devicePackageIds] } : {}),
summary: {
projectId,
agentId,
@@ -311,6 +328,10 @@ export async function materializePiAgentResources(
catalogRevision: resolvedCatalogRevision,
revision: { ...options.revision },
...(options.effectiveSnapshot ? { effectivePluginSnapshot: structuredClone(options.effectiveSnapshot) } : {}),
...(options.devicePackageGeneration === undefined
? {}
: { devicePackageGeneration: catalogRevision(options.devicePackageGeneration) }),
...(options.devicePackageIds ? { devicePackageIds: [...options.devicePackageIds] } : {}),
},
};
}

View File

@@ -85,6 +85,9 @@ import {
} from './session-projector';
import { PiManagedExtensionHost } from './extension-host';
import type { PiSubagentScheduler } from './subagent';
import type { ModelToolRegistryPort } from './model-tools/model-tool-registry';
import type { DevicePackageManager } from '../../coding-packages/device-package-manager';
import type { CodingPluginToolDefinition } from '../../../shared/coding-plugins';
import type { CodingCapabilityRegistry, ResolvedWorkerResources } from '../../coding-plugins/registry';
import {
PiInteractionStore,
@@ -152,6 +155,9 @@ export interface PiManagedWorkerOpenerOptions {
onTelemetry?: (event: PiRuntimeTelemetryEvent) => void;
extensionHost: PiManagedExtensionHost;
capabilityRegistry?: CodingCapabilityRegistry;
modelToolRegistry?: ModelToolRegistryPort;
devicePackageManager?: Pick<DevicePackageManager, 'resolveEnabledResources' | 'registerActiveWorker'>;
devicePackageTools?: readonly CodingPluginToolDefinition[];
}
interface PiRpcSessionStateProjection {
@@ -253,12 +259,19 @@ export function createPiManagedWorkerOpener(
role: 'parent',
})
: fallbackWorkerResources(registered.agent.skillIds);
const deviceResources = options.devicePackageManager
? await options.devicePackageManager.resolveEnabledResources()
: undefined;
const combinedSkillEntries = [
...workerResources.skillEntries,
...(deviceResources?.skillEntries ?? []),
];
const resources = await materializePiAgentResources({
userDataDir: options.userDataDir,
projectId: input.conversation.projectId,
agentId: registered.agent.id,
prompt: registered.agent.prompt,
skillEntries: workerResources.skillEntries,
skillEntries: combinedSkillEntries,
catalogRevision: workerResources.catalogRevision,
bundledSkillsDir: options.bundledSkillsDir,
...(workerResources.skillRoots
@@ -267,6 +280,12 @@ export function createPiManagedWorkerOpener(
...(workerResources.effectiveSnapshot
? { effectiveSnapshot: workerResources.effectiveSnapshot }
: {}),
...(deviceResources
? {
devicePackageGeneration: deviceResources.generation,
devicePackageIds: deviceResources.packageIds,
}
: {}),
revision: input.revision,
});
const credential = await buildPiWorkerCredentialProjection({
@@ -277,19 +296,43 @@ export function createPiManagedWorkerOpener(
? { localProxyCredential: await options.getLocalProxyCredential() }
: {}),
});
const extension = await options.extensionHost.registerWorker({
const modelTools = options.modelToolRegistry?.registerWorker({
conversationId: input.conversation.conversationId,
generation: input.generation,
projectId: input.conversation.projectId,
projectPath: registered.projectPath,
skillEntries: workerResources.skillEntries,
catalogRevision: workerResources.catalogRevision,
tools: workerResources.tools,
...(workerResources.effectiveSnapshot
? { effectiveSnapshot: workerResources.effectiveSnapshot }
: {}),
extensionsDir: managedPaths.extensionsDir,
account,
descriptor,
selection,
credential,
});
let extension;
try {
extension = await options.extensionHost.registerWorker({
conversationId: input.conversation.conversationId,
generation: input.generation,
projectId: input.conversation.projectId,
projectPath: registered.projectPath,
skillEntries: combinedSkillEntries,
catalogRevision: workerResources.catalogRevision,
tools: [
...workerResources.tools,
...(modelTools?.tools ?? []),
...(options.devicePackageTools ?? []),
],
...(workerResources.effectiveSnapshot
? { effectiveSnapshot: workerResources.effectiveSnapshot }
: {}),
...(deviceResources
? {
devicePackageGeneration: deviceResources.generation,
devicePackageIds: deviceResources.packageIds,
}
: {}),
extensionsDir: managedPaths.extensionsDir,
});
} catch (error) {
modelTools?.dispose();
throw error;
}
recordManagedMilestone(
options.onTelemetry,
input,
@@ -317,6 +360,9 @@ export function createPiManagedWorkerOpener(
additionalArgs: [
...buildPiManagedInputArgs(selection, resources),
'--extension', extension.extensionPath,
...(deviceResources?.extensionPaths.flatMap((extensionPath) => [
'--extension', extensionPath,
]) ?? []),
...(input.fork ? ['--fork', input.fork.sourceSession.piSessionId] : []),
'--session-id', sessionKey,
],
@@ -328,6 +374,9 @@ export function createPiManagedWorkerOpener(
const releaseActivePluginReleases = workerResources.effectiveSnapshot
? options.registerActivePluginReleases?.(workerResources.effectiveSnapshot.pluginReleaseIds)
: undefined;
const releaseActiveDevicePackages = deviceResources
? options.devicePackageManager?.registerActiveWorker(deviceResources.packageRefs)
: undefined;
let managedResourcesDisposed = false;
const disposeManagedResources = async (): Promise<void> => {
if (managedResourcesDisposed) return;
@@ -335,7 +384,15 @@ export function createPiManagedWorkerOpener(
try {
await extension.dispose();
} finally {
await releaseActivePluginReleases?.();
try {
modelTools?.dispose();
} finally {
try {
await releaseActiveDevicePackages?.();
} finally {
await releaseActivePluginReleases?.();
}
}
}
};
let unsubscribeExtensionInvalidation = process.subscribeInvalidation(() => {
@@ -1214,6 +1271,10 @@ export class PiConversationRuntime implements CodingConversationRuntime {
this.pool.markResourcesStale();
}
async refreshResources(): Promise<void> {
await this.pool.refreshResources();
}
subscribe(listener: (patch: ConversationPatchEnvelope) => void): () => void {
this.listeners.add(listener);
return () => this.listeners.delete(listener);

View File

@@ -356,6 +356,18 @@ export class PiWorkerPool {
return this.revisions.markResourcesStale();
}
async refreshResources(): Promise<void> {
this.revisions.markResourcesStale();
const idleWorkers = [...this.workers.values()].filter((record) => (
(record.state === 'ready' || record.state === 'idle')
&& !this.activeRuns.has(record.conversation.conversationId)
));
await Promise.all(idleWorkers.map(async (record) => {
if (this.workers.get(record.conversation.conversationId) !== record) return;
await this.ensureFresh(record);
}));
}
async fork(
sourceConversationId: string,
conversation: PrepareConversationInput,