fix(pi): resolve packaged proxy token lazily
This commit is contained in:
@@ -1,8 +1,11 @@
|
||||
import { createServer, type IncomingMessage, type Server, type ServerResponse } from 'node:http';
|
||||
import { mkdtemp, rm } from 'node:fs/promises';
|
||||
import { execFile } from 'node:child_process';
|
||||
import { mkdtemp, readFile, readdir, rm } from 'node:fs/promises';
|
||||
import { tmpdir } from 'node:os';
|
||||
import path from 'node:path';
|
||||
import { promisify } from 'node:util';
|
||||
|
||||
import type { CodingProductComposition } from '../../api/coding-product-services';
|
||||
import { createCodingConversationStore } from '../../coding-projects/conversation-store';
|
||||
import { createCodingProjectAgent } from '../../coding-projects/project-config';
|
||||
import {
|
||||
@@ -12,6 +15,11 @@ import {
|
||||
type CodingProjectStore,
|
||||
} from '../../coding-projects/project-store';
|
||||
import type { ProviderAccount } from '../../shared/providers/types';
|
||||
import { getProviderService } from '../../services/providers/provider-service';
|
||||
import {
|
||||
clearWorksSquareAIGatewayCredential,
|
||||
seedWorksSquareAIGatewayCredential,
|
||||
} from '../../services/works-square-ai-gateway';
|
||||
import type { PrepareConversationInput } from '../contracts';
|
||||
import { PiManagedExtensionHost } from './extension-host';
|
||||
import { PiManagedInputRevisionCoordinator } from './managed-input-revision';
|
||||
@@ -85,6 +93,21 @@ type PressureRun = {
|
||||
finish(options?: { injectFailureAt?: 'parents.settle' }): Promise<PiReleasePressureSnapshot>;
|
||||
};
|
||||
|
||||
type ProxyCompositionRun = {
|
||||
composition: CodingProductComposition;
|
||||
provider: LocalProofProvider;
|
||||
projectId: string;
|
||||
projectPath: string;
|
||||
hostToken: string;
|
||||
activeStatus?: PiReleaseProxyCompositionStatus;
|
||||
};
|
||||
|
||||
type InspectedPiProcess = {
|
||||
processId: number;
|
||||
role: ProofWorkerRole;
|
||||
commandLineContainsHostToken: boolean;
|
||||
};
|
||||
|
||||
export interface PiReleaseManagedTurnProof {
|
||||
runRef: string;
|
||||
cold: boolean;
|
||||
@@ -130,6 +153,40 @@ export interface PiReleaseExtensionProof {
|
||||
};
|
||||
}
|
||||
|
||||
export interface PiReleaseProxyCompositionStatus {
|
||||
conversationCount: number;
|
||||
providerRequests: { parent: number; child: number };
|
||||
activeProviderRequests: { parent: number; child: number };
|
||||
workers: ReturnType<CodingProductComposition['runtime']['getDiagnostics']>['workers'];
|
||||
processes: {
|
||||
supported: boolean;
|
||||
parent: number[];
|
||||
child: number[];
|
||||
};
|
||||
currentHostTokenAvailable: boolean;
|
||||
tokenSafety: {
|
||||
argv: boolean;
|
||||
modelsJson: boolean;
|
||||
logs: boolean;
|
||||
diagnostics: boolean;
|
||||
};
|
||||
realTurnVerified: false;
|
||||
}
|
||||
|
||||
export interface PiReleaseProxyCompositionProof extends PiReleaseProxyCompositionStatus {
|
||||
providerMode: 'loopback-through-authenticated-host-proxy';
|
||||
firstConversation: {
|
||||
count: number;
|
||||
bindingEstablished: boolean;
|
||||
workerStatus: string;
|
||||
runStatus: string;
|
||||
userInputProjected: boolean;
|
||||
parentResponseProjected: boolean;
|
||||
};
|
||||
currentHostTokenUsedBy: ProofWorkerRole[];
|
||||
released: { workers: number };
|
||||
}
|
||||
|
||||
const PROOF_ACCOUNT_ID = 'release-proof-account';
|
||||
const PROOF_AGENT_ID = 'release-proof-agent';
|
||||
const PROOF_MODEL_ID = 'release-proof-model';
|
||||
@@ -147,6 +204,8 @@ const EXPECTED_TURN_MILESTONES: readonly ProofMilestone[] = [
|
||||
];
|
||||
|
||||
let pressureRun: PressureRun | null = null;
|
||||
let proxyCompositionRun: ProxyCompositionRun | null = null;
|
||||
const execFileAsync = promisify(execFile);
|
||||
|
||||
async function waitFor(predicate: () => boolean, message: string): Promise<void> {
|
||||
const deadline = Date.now() + 30_000;
|
||||
@@ -532,6 +591,122 @@ function providerRequestCounts(provider: LocalProofProvider): { parent: number;
|
||||
};
|
||||
}
|
||||
|
||||
async function fileContainsValue(filePath: string, value: string): Promise<boolean> {
|
||||
try {
|
||||
return (await readFile(filePath)).includes(Buffer.from(value));
|
||||
} catch (error) {
|
||||
if ((error as NodeJS.ErrnoException).code === 'ENOENT') return false;
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
async function directoryContainsValue(directory: string, value: string): Promise<boolean> {
|
||||
let entries;
|
||||
try {
|
||||
entries = await readdir(directory, { withFileTypes: true });
|
||||
} catch (error) {
|
||||
if ((error as NodeJS.ErrnoException).code === 'ENOENT') return false;
|
||||
throw error;
|
||||
}
|
||||
for (const entry of entries) {
|
||||
const entryPath = path.join(directory, entry.name);
|
||||
if (entry.isDirectory()) {
|
||||
if (await directoryContainsValue(entryPath, value)) return true;
|
||||
} else if (entry.isFile() && await fileContainsValue(entryPath, value)) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
async function inspectWindowsPiProcesses(hostToken: string): Promise<{
|
||||
supported: boolean;
|
||||
processes: InspectedPiProcess[];
|
||||
}> {
|
||||
if (process.platform !== 'win32') return { supported: false, processes: [] };
|
||||
const script = `$parentPid = ${process.pid}; `
|
||||
+ '$items = @(Get-CimInstance Win32_Process -Filter "ParentProcessId = $parentPid" '
|
||||
+ '| Where-Object { $_.CommandLine -match "pi-runtime[\\\\/]dist[\\\\/]cli\\.js" } '
|
||||
+ '| Select-Object ProcessId, CommandLine); '
|
||||
+ '$items | ConvertTo-Json -Compress';
|
||||
const { stdout } = await execFileAsync(
|
||||
'powershell.exe',
|
||||
['-NoProfile', '-NonInteractive', '-Command', script],
|
||||
{ windowsHide: true, maxBuffer: 1024 * 1024 },
|
||||
);
|
||||
const text = stdout.trim();
|
||||
const parsed = text ? JSON.parse(text) as unknown : [];
|
||||
const items = Array.isArray(parsed) ? parsed : [parsed];
|
||||
const processes = items.flatMap((item): InspectedPiProcess[] => {
|
||||
if (!item || typeof item !== 'object' || Array.isArray(item)) return [];
|
||||
const record = item as { ProcessId?: unknown; CommandLine?: unknown };
|
||||
const processId = Number(record.ProcessId);
|
||||
const commandLine = typeof record.CommandLine === 'string' ? record.CommandLine : '';
|
||||
if (!Number.isSafeInteger(processId) || !commandLine) return [];
|
||||
const role = commandLine.includes('--no-session') ? 'child' : 'parent';
|
||||
return [{
|
||||
processId,
|
||||
role,
|
||||
commandLineContainsHostToken: commandLine.includes(hostToken),
|
||||
}];
|
||||
});
|
||||
return { supported: true, processes };
|
||||
}
|
||||
|
||||
function proxyProviderAccount(baseUrl: string): ProviderAccount {
|
||||
return {
|
||||
...providerAccount(baseUrl),
|
||||
label: 'Authenticated Host proxy release proof',
|
||||
metadata: {
|
||||
worksSquareCredentialMode: 'works_square_ai_gateway_proxy',
|
||||
customModels: [PROOF_MODEL_ID],
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
async function proxyCompositionStatus(
|
||||
run: ProxyCompositionRun,
|
||||
): Promise<PiReleaseProxyCompositionStatus> {
|
||||
const diagnostics = run.composition.runtime.getDiagnostics();
|
||||
const processInspection = await inspectWindowsPiProcesses(run.hostToken);
|
||||
const modelsFile = path.join(
|
||||
path.dirname(run.projectPath),
|
||||
'coding-runtime',
|
||||
'pi',
|
||||
'config',
|
||||
'models.json',
|
||||
);
|
||||
const logsDirectory = path.join(path.dirname(run.projectPath), 'logs');
|
||||
const [modelsContainToken, logsContainToken] = await Promise.all([
|
||||
fileContainsValue(modelsFile, run.hostToken),
|
||||
directoryContainsValue(logsDirectory, run.hostToken),
|
||||
]);
|
||||
const providerRequests = providerRequestCounts(run.provider);
|
||||
const parent = processInspection.processes
|
||||
.filter(({ role }) => role === 'parent')
|
||||
.map(({ processId }) => processId)
|
||||
.sort((left, right) => left - right);
|
||||
const child = processInspection.processes
|
||||
.filter(({ role }) => role === 'child')
|
||||
.map(({ processId }) => processId)
|
||||
.sort((left, right) => left - right);
|
||||
return {
|
||||
conversationCount: (await run.composition.conversations.listConversations(run.projectId)).length,
|
||||
providerRequests,
|
||||
activeProviderRequests: run.provider.activeCounts(),
|
||||
workers: diagnostics.workers,
|
||||
processes: { supported: processInspection.supported, parent, child },
|
||||
currentHostTokenAvailable: run.hostToken.length > 0,
|
||||
tokenSafety: {
|
||||
argv: processInspection.processes.every(({ commandLineContainsHostToken }) => !commandLineContainsHostToken),
|
||||
modelsJson: !modelsContainToken,
|
||||
logs: !logsContainToken,
|
||||
diagnostics: !JSON.stringify(diagnostics).includes(run.hostToken),
|
||||
},
|
||||
realTurnVerified: false,
|
||||
};
|
||||
}
|
||||
|
||||
function timelineForTurn(
|
||||
composition: RealProofComposition,
|
||||
conversationId: string,
|
||||
@@ -929,6 +1104,178 @@ export async function runFinalAsarExtensionProof(): Promise<PiReleaseExtensionPr
|
||||
};
|
||||
}
|
||||
|
||||
async function cleanupProxyCompositionRun(run: ProxyCompositionRun): Promise<void> {
|
||||
run.provider.releaseAll();
|
||||
await run.composition.projects.removeProject(run.projectId).catch(() => undefined);
|
||||
await run.provider.close().catch(() => undefined);
|
||||
clearWorksSquareAIGatewayCredential();
|
||||
await getProviderService().deleteAccount(PROOF_ACCOUNT_ID).catch(() => undefined);
|
||||
}
|
||||
|
||||
export async function startFinalAsarProxyCompositionProof(input: {
|
||||
composition: CodingProductComposition;
|
||||
projectPath: string;
|
||||
hostProxyBaseUrl: string;
|
||||
hostToken: string;
|
||||
}): Promise<{
|
||||
projectId: string;
|
||||
agentId: string;
|
||||
providerMode: 'works_square_ai_gateway_proxy';
|
||||
currentHostTokenAvailable: true;
|
||||
realTurnVerified: false;
|
||||
}> {
|
||||
if (proxyCompositionRun) throw new Error('PI proxy composition proof is already running');
|
||||
const hostToken = input.hostToken.trim();
|
||||
if (!hostToken) throw new Error('Current Main Host token is unavailable');
|
||||
const provider = await startLocalProofProvider('subagent');
|
||||
let projectId: string | null = null;
|
||||
try {
|
||||
seedWorksSquareAIGatewayCredential({
|
||||
accessToken: 'release-proof-upstream-only',
|
||||
oneApiBaseUrl: provider.baseUrl,
|
||||
});
|
||||
const providerService = getProviderService();
|
||||
await providerService.createAccount(proxyProviderAccount(input.hostProxyBaseUrl));
|
||||
await providerService.setDefaultAccount(PROOF_ACCOUNT_ID);
|
||||
const project = await input.composition.projects.createProject({ projectPath: input.projectPath });
|
||||
projectId = project.project.id;
|
||||
await createCodingProjectAgent(input.projectPath, {
|
||||
id: PROOF_AGENT_ID,
|
||||
avatarId: 'avatar-01',
|
||||
roleName: 'Packaged proxy proof',
|
||||
name: 'Packaged proxy proof agent',
|
||||
model: {
|
||||
accountId: PROOF_ACCOUNT_ID,
|
||||
modelId: PROOF_MODEL_ID,
|
||||
thinkingLevel: 'medium',
|
||||
},
|
||||
modelResolution: 'resolved',
|
||||
responsibility: {
|
||||
mission: 'Exercise the real packaged Main proxy composition',
|
||||
owns: [],
|
||||
boundaries: [],
|
||||
collaborators: [],
|
||||
principles: [],
|
||||
},
|
||||
prompt: 'Follow the controlled loopback Provider and complete its subagent request.',
|
||||
skillIds: [],
|
||||
});
|
||||
proxyCompositionRun = {
|
||||
composition: input.composition,
|
||||
provider,
|
||||
projectId,
|
||||
projectPath: input.projectPath,
|
||||
hostToken,
|
||||
};
|
||||
return {
|
||||
projectId,
|
||||
agentId: PROOF_AGENT_ID,
|
||||
providerMode: 'works_square_ai_gateway_proxy',
|
||||
currentHostTokenAvailable: true,
|
||||
realTurnVerified: false,
|
||||
};
|
||||
} catch (error) {
|
||||
provider.releaseAll();
|
||||
if (projectId) await input.composition.projects.removeProject(projectId).catch(() => undefined);
|
||||
await provider.close().catch(() => undefined);
|
||||
clearWorksSquareAIGatewayCredential();
|
||||
await getProviderService().deleteAccount(PROOF_ACCOUNT_ID).catch(() => undefined);
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
export async function getFinalAsarProxyCompositionStatus(): Promise<PiReleaseProxyCompositionStatus> {
|
||||
const run = proxyCompositionRun;
|
||||
if (!run) throw new Error('PI proxy composition proof is not running');
|
||||
const status = await proxyCompositionStatus(run);
|
||||
if (status.activeProviderRequests.child > 0
|
||||
&& status.processes.child.length > 0
|
||||
&& status.processes.parent.length > 0) {
|
||||
run.activeStatus = status;
|
||||
}
|
||||
return status;
|
||||
}
|
||||
|
||||
export function releaseFinalAsarProxyCompositionChild(): void {
|
||||
const run = proxyCompositionRun;
|
||||
if (!run) throw new Error('PI proxy composition proof is not running');
|
||||
run.provider.releaseChildren();
|
||||
}
|
||||
|
||||
export async function finishFinalAsarProxyCompositionProof(): Promise<PiReleaseProxyCompositionProof> {
|
||||
const run = proxyCompositionRun;
|
||||
if (!run) throw new Error('PI proxy composition proof is not running');
|
||||
try {
|
||||
const conversations = await run.composition.conversations.listConversations(run.projectId);
|
||||
if (conversations.length !== 1) {
|
||||
throw new Error(`Expected one first proxy Conversation, received ${conversations.length}`);
|
||||
}
|
||||
const conversation = conversations[0]!;
|
||||
let snapshot = await run.composition.conversations.getSnapshot(conversation.id);
|
||||
const deadline = Date.now() + 30_000;
|
||||
while (Date.now() < deadline && (
|
||||
snapshot.run.status !== 'idle'
|
||||
|| !JSON.stringify(snapshot.nodes).includes('REAL_PARENT_COMPLETE')
|
||||
)) {
|
||||
await delay(20);
|
||||
snapshot = await run.composition.conversations.getSnapshot(conversation.id);
|
||||
}
|
||||
const serializedNodes = JSON.stringify(snapshot.nodes);
|
||||
if (snapshot.run.status !== 'idle' || !serializedNodes.includes('REAL_PARENT_COMPLETE')) {
|
||||
throw new Error('First proxy Conversation did not settle with the controlled parent response');
|
||||
}
|
||||
const binding = await run.composition.projects.conversationStore(run.projectPath).get(conversation.id);
|
||||
const finalStatus = await proxyCompositionStatus(run);
|
||||
const activeStatus = run.activeStatus ?? finalStatus;
|
||||
const currentHostTokenUsedBy: ProofWorkerRole[] = [
|
||||
...(activeStatus.providerRequests.parent > 0 ? ['parent' as const] : []),
|
||||
...(activeStatus.providerRequests.child > 0 ? ['child' as const] : []),
|
||||
];
|
||||
if (currentHostTokenUsedBy.join(',') !== 'parent,child') {
|
||||
throw new Error('Authenticated Host proxy did not receive both parent and child Provider requests');
|
||||
}
|
||||
if (!binding?.piSessionId || !binding.sessionKey) {
|
||||
throw new Error('First proxy Conversation did not persist its Pi session binding');
|
||||
}
|
||||
const tokenSafety = {
|
||||
argv: activeStatus.tokenSafety.argv && finalStatus.tokenSafety.argv,
|
||||
modelsJson: activeStatus.tokenSafety.modelsJson && finalStatus.tokenSafety.modelsJson,
|
||||
logs: activeStatus.tokenSafety.logs && finalStatus.tokenSafety.logs,
|
||||
diagnostics: activeStatus.tokenSafety.diagnostics && finalStatus.tokenSafety.diagnostics,
|
||||
};
|
||||
if (Object.values(tokenSafety).some((safe) => !safe)) {
|
||||
throw new Error(`Current Main Host token escaped its credential boundary: ${JSON.stringify(tokenSafety)}`);
|
||||
}
|
||||
const proof: Omit<PiReleaseProxyCompositionProof, 'released'> = {
|
||||
...finalStatus,
|
||||
providerRequests: activeStatus.providerRequests,
|
||||
activeProviderRequests: activeStatus.activeProviderRequests,
|
||||
processes: activeStatus.processes,
|
||||
tokenSafety,
|
||||
providerMode: 'loopback-through-authenticated-host-proxy',
|
||||
firstConversation: {
|
||||
count: conversations.length,
|
||||
bindingEstablished: true,
|
||||
workerStatus: snapshot.worker.status,
|
||||
runStatus: snapshot.run.status,
|
||||
userInputProjected: serializedNodes.includes('Exercise the packaged Main proxy Conversation'),
|
||||
parentResponseProjected: serializedNodes.includes('REAL_PARENT_COMPLETE'),
|
||||
},
|
||||
currentHostTokenUsedBy,
|
||||
};
|
||||
await cleanupProxyCompositionRun(run);
|
||||
proxyCompositionRun = null;
|
||||
return {
|
||||
...proof,
|
||||
released: { workers: run.composition.runtime.getDiagnostics().workers.length },
|
||||
};
|
||||
} catch (error) {
|
||||
await cleanupProxyCompositionRun(run);
|
||||
proxyCompositionRun = null;
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
export async function startFinalAsarPressureProof(): Promise<PiReleasePressureSnapshot> {
|
||||
if (pressureRun) throw new Error('PI release pressure proof is already running');
|
||||
pressureRun = await startPressureRun();
|
||||
|
||||
Reference in New Issue
Block a user