351 lines
12 KiB
TypeScript
351 lines
12 KiB
TypeScript
import type { ExtensionAPI } from '@earendil-works/pi-coding-agent';
|
|
import { Type } from 'typebox';
|
|
import { createHash, randomUUID } from 'node:crypto';
|
|
import { readFileSync, appendFileSync, statSync } from 'node:fs';
|
|
import net from 'node:net';
|
|
import { fileURLToPath } from 'node:url';
|
|
import { resolve } from 'node:path';
|
|
|
|
const SELF = resolve(fileURLToPath(import.meta.url).split('?')[0]!);
|
|
const LOG = process.env['GATE0_PI_LOG'];
|
|
const CONTEXT_BLOCK =
|
|
process.env['GATE0_PI_CONTEXT_BLOCK'] ??
|
|
[
|
|
'GATE0_PI_ATOMIC_BEGIN',
|
|
'segment-01=alpha-7e31',
|
|
'segment-02=middle-9c42',
|
|
'segment-03=omega-5b83',
|
|
'GATE0_PI_ATOMIC_END',
|
|
].join('\n');
|
|
|
|
let sequence = 0;
|
|
|
|
function sha(value: string | Buffer): string {
|
|
return createHash('sha256').update(value).digest('hex');
|
|
}
|
|
|
|
function procStarttime(): number {
|
|
const text = readFileSync(`/proc/${process.pid}/stat`, 'utf8');
|
|
const close = text.lastIndexOf(')');
|
|
const fields = text.slice(close + 2).trim().split(/\s+/);
|
|
return Number(fields[19]);
|
|
}
|
|
|
|
function log(event: string, details: Record<string, unknown> = {}): void {
|
|
if (!LOG) return;
|
|
sequence += 1;
|
|
appendFileSync(
|
|
LOG,
|
|
`${JSON.stringify({ seq: sequence, event, pid: process.pid, starttime_ticks: procStarttime(), ...details })}\n`,
|
|
);
|
|
}
|
|
|
|
function argvExtensions(): string[] {
|
|
const result: string[] = [];
|
|
for (let i = 0; i < process.argv.length; i += 1) {
|
|
if (process.argv[i] === '--extension' || process.argv[i] === '-e') {
|
|
const candidate = process.argv[i + 1];
|
|
if (candidate) result.push(resolve(candidate));
|
|
}
|
|
}
|
|
return result;
|
|
}
|
|
|
|
interface SourceValidation {
|
|
ok: boolean;
|
|
reason: string;
|
|
fragment?: string;
|
|
}
|
|
|
|
function validateSources(): SourceValidation {
|
|
const manifestPath = process.env['GATE0_SOURCE_MANIFEST'];
|
|
if (!manifestPath) return { ok: true, reason: 'no-manifest-probe-disabled' };
|
|
try {
|
|
const manifest = JSON.parse(readFileSync(manifestPath, 'utf8')) as {
|
|
maxBytes: number;
|
|
fragments: Array<{ path: string; sha256: string }>;
|
|
};
|
|
for (const fragment of manifest.fragments) {
|
|
let fileStat;
|
|
try {
|
|
fileStat = statSync(fragment.path);
|
|
} catch {
|
|
return { ok: false, reason: 'missing', fragment: fragment.path };
|
|
}
|
|
if (!fileStat.isFile()) {
|
|
return { ok: false, reason: 'not-regular-file', fragment: fragment.path };
|
|
}
|
|
if (fileStat.size > manifest.maxBytes) {
|
|
return { ok: false, reason: 'oversize', fragment: fragment.path };
|
|
}
|
|
const bytes = readFileSync(fragment.path);
|
|
if (sha(bytes) !== fragment.sha256) {
|
|
return { ok: false, reason: 'hash-mismatch', fragment: fragment.path };
|
|
}
|
|
}
|
|
return { ok: true, reason: 'all-fragments-valid' };
|
|
} catch (error) {
|
|
return { ok: false, reason: `manifest-error:${error instanceof Error ? error.name : 'unknown'}` };
|
|
}
|
|
}
|
|
|
|
function brokerRequest(payload: Record<string, unknown>): Promise<Record<string, unknown>> {
|
|
const socketPath = process.env['GATE0_GENERATION_SOCKET'];
|
|
if (!socketPath) return Promise.resolve({ skipped: true });
|
|
return new Promise((resolvePromise, reject) => {
|
|
const socket = net.createConnection(socketPath);
|
|
let buffer = '';
|
|
socket.setEncoding('utf8');
|
|
socket.on('connect', () => socket.write(`${JSON.stringify(payload)}\n`));
|
|
socket.on('data', (chunk) => {
|
|
buffer += chunk;
|
|
const newline = buffer.indexOf('\n');
|
|
if (newline < 0) return;
|
|
socket.end();
|
|
resolvePromise(JSON.parse(buffer.slice(0, newline)) as Record<string, unknown>);
|
|
});
|
|
socket.on('error', reject);
|
|
});
|
|
}
|
|
|
|
function markerPaths(value: unknown, path = '$'): string[] {
|
|
const matches: string[] = [];
|
|
if (typeof value === 'string') {
|
|
if (value.includes(CONTEXT_BLOCK)) matches.push(path);
|
|
return matches;
|
|
}
|
|
if (Array.isArray(value)) {
|
|
value.forEach((item, index) => matches.push(...markerPaths(item, `${path}[${index}]`)));
|
|
return matches;
|
|
}
|
|
if (value && typeof value === 'object') {
|
|
for (const [key, item] of Object.entries(value as Record<string, unknown>)) {
|
|
matches.push(...markerPaths(item, `${path}.${key}`));
|
|
}
|
|
}
|
|
return matches;
|
|
}
|
|
|
|
function assistantToolIds(message: unknown): string[] {
|
|
if (!message || typeof message !== 'object') return [];
|
|
const candidate = message as { role?: string; content?: unknown };
|
|
if (candidate.role !== 'assistant' || !Array.isArray(candidate.content)) return [];
|
|
return candidate.content
|
|
.filter(
|
|
(block): block is { type: 'toolCall'; id: string } =>
|
|
Boolean(
|
|
block &&
|
|
typeof block === 'object' &&
|
|
(block as { type?: string }).type === 'toolCall' &&
|
|
typeof (block as { id?: unknown }).id === 'string',
|
|
),
|
|
)
|
|
.map((block) => block.id);
|
|
}
|
|
|
|
function assistantText(message: unknown): string {
|
|
if (!message || typeof message !== 'object') return '';
|
|
const candidate = message as { role?: string; content?: unknown };
|
|
if (candidate.role !== 'assistant' || !Array.isArray(candidate.content)) return '';
|
|
return candidate.content
|
|
.filter(
|
|
(block): block is { type: 'text'; text: string } =>
|
|
Boolean(
|
|
block &&
|
|
typeof block === 'object' &&
|
|
(block as { type?: string }).type === 'text' &&
|
|
typeof (block as { text?: unknown }).text === 'string',
|
|
),
|
|
)
|
|
.map((block) => block.text)
|
|
.join('');
|
|
}
|
|
|
|
export default function register(pi: ExtensionAPI) {
|
|
const localProviderUrl = process.env['GATE0_LOCAL_PROVIDER_URL'];
|
|
if (localProviderUrl) {
|
|
pi.registerProvider('gate0-local', {
|
|
baseUrl: localProviderUrl,
|
|
apiKey: 'gate0-probe-not-a-secret',
|
|
api: 'openai-completions',
|
|
models: [
|
|
{
|
|
id: 'gate0-model',
|
|
name: 'Gate0 deterministic local model',
|
|
reasoning: false,
|
|
input: ['text'],
|
|
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
|
|
contextWindow: 32_000,
|
|
maxTokens: 1_024,
|
|
},
|
|
],
|
|
});
|
|
}
|
|
|
|
const extensions = argvExtensions();
|
|
const lastPosition = extensions.length > 0 && extensions.at(-1) === SELF;
|
|
interface RequestCycle {
|
|
nonce: string;
|
|
verified: boolean;
|
|
sourceReason: string;
|
|
}
|
|
let buildingCycle: RequestCycle | undefined;
|
|
const inFlightCycles: RequestCycle[] = [];
|
|
const toolNonce = new Map<string, { nonce: string; verified: boolean; sourceReason: string }>();
|
|
|
|
pi.on('session_start', async (event) => {
|
|
const broker = await brokerRequest({ action: 'lifecycle', phase: 'start', reason: event.reason });
|
|
log('session_start', {
|
|
reason: event.reason,
|
|
extensions,
|
|
self: SELF,
|
|
lastPosition,
|
|
gateState: lastPosition ? 'UNVERIFIED_READY' : 'CLOSED_NOT_LAST',
|
|
broker,
|
|
});
|
|
});
|
|
|
|
pi.on('session_shutdown', async (event) => {
|
|
const broker = await brokerRequest({ action: 'lifecycle', phase: 'shutdown', reason: event.reason });
|
|
log('session_shutdown', { reason: event.reason, broker });
|
|
});
|
|
|
|
pi.on('context', async (event) => {
|
|
const validation = validateSources();
|
|
buildingCycle = {
|
|
nonce: randomUUID(),
|
|
sourceReason: validation.reason,
|
|
verified: lastPosition && validation.ok,
|
|
};
|
|
const inputJson = JSON.stringify(event.messages);
|
|
const injected = {
|
|
role: 'custom' as const,
|
|
customType: 'gate0-context',
|
|
content: CONTEXT_BLOCK,
|
|
display: false,
|
|
timestamp: Date.now(),
|
|
};
|
|
const outputMessages = buildingCycle.verified
|
|
? [...event.messages, injected]
|
|
: [...event.messages];
|
|
const outputPrefix = outputMessages.slice(0, event.messages.length);
|
|
const sourceBroker = validation.ok
|
|
? { action: 'none', reason: 'source-valid' }
|
|
: await brokerRequest({ action: 'source-invalid', reason: validation.reason });
|
|
log('context_return', {
|
|
requestNonce: buildingCycle.nonce,
|
|
sourceValidation: validation,
|
|
sourceBroker,
|
|
lastPosition,
|
|
promotion: false,
|
|
injectionDecision: buildingCycle.verified ? 'ONE_ATOMIC_AGENT_MESSAGE' : 'REFUSED',
|
|
inputCount: event.messages.length,
|
|
outputCount: outputMessages.length,
|
|
prefixHashBefore: sha(inputJson),
|
|
prefixHashAfter: sha(JSON.stringify(outputPrefix)),
|
|
prefixPreservedByReturn: sha(inputJson) === sha(JSON.stringify(outputPrefix)),
|
|
blockLength: CONTEXT_BLOCK.length,
|
|
blockSha256: sha(CONTEXT_BLOCK),
|
|
});
|
|
return { messages: outputMessages };
|
|
});
|
|
|
|
pi.on('before_provider_request', async (event) => {
|
|
const paths = markerPaths(event.payload);
|
|
const cycle = buildingCycle;
|
|
buildingCycle = undefined;
|
|
if (cycle) inFlightCycles.push(cycle);
|
|
log('before_provider_request', {
|
|
requestNonce: cycle?.nonce,
|
|
inFlightDepth: inFlightCycles.length,
|
|
markerOccurrences: paths.length,
|
|
markerPaths: paths,
|
|
finalPayloadValid: Boolean(cycle?.verified && paths.length === 1),
|
|
});
|
|
});
|
|
|
|
pi.on('after_provider_response', async (event) => {
|
|
const cycle = inFlightCycles[0];
|
|
log('after_provider_response', {
|
|
requestNonce: cycle?.nonce,
|
|
status: event.status,
|
|
assistantContentAvailableAtThisHook: false,
|
|
timing: 'headers/status before stream consumption',
|
|
});
|
|
});
|
|
|
|
pi.on('message_end', async (event) => {
|
|
const role = (event.message as { role?: string }).role;
|
|
const ids = assistantToolIds(event.message);
|
|
const text = assistantText(event.message);
|
|
const cycle = role === 'assistant' ? inFlightCycles.shift() : undefined;
|
|
if (ids.length > 0 && cycle) {
|
|
for (const id of ids) {
|
|
toolNonce.set(id, {
|
|
nonce: cycle.nonce,
|
|
verified: cycle.verified,
|
|
sourceReason: cycle.sourceReason,
|
|
});
|
|
}
|
|
}
|
|
log('message_end', {
|
|
role,
|
|
assistantContentObserved: role === 'assistant',
|
|
requestNonce: cycle?.nonce,
|
|
inFlightDepthAfter: inFlightCycles.length,
|
|
toolCallIds: ids,
|
|
nonceMappings: ids.map((id) => ({ toolCallId: id, requestNonce: cycle?.nonce })),
|
|
exactContextBlockCopied: text.includes(CONTEXT_BLOCK),
|
|
assistantTextSha256: text ? sha(text) : null,
|
|
});
|
|
});
|
|
|
|
pi.on('tool_call', async (event) => {
|
|
const mapping = toolNonce.get(event.toolCallId);
|
|
const allowed = Boolean(lastPosition && mapping?.verified);
|
|
log('tool_call', {
|
|
toolCallId: event.toolCallId,
|
|
toolName: event.toolName,
|
|
mapping: mapping ?? null,
|
|
allowed,
|
|
reason: !lastPosition
|
|
? 'closed-not-last'
|
|
: !mapping
|
|
? 'unknown-tool-call-id'
|
|
: !mapping.verified
|
|
? `unverified-source:${mapping.sourceReason}`
|
|
: 'exact-tool-call-id-mapped-to-verified-request-nonce',
|
|
});
|
|
if (!allowed) return { block: true, reason: 'Gate0 probe refused unverified tool batch' };
|
|
});
|
|
|
|
pi.on('agent_settled', async () => {
|
|
log('agent_settled', { retainedNonceMappingsBeforeClear: toolNonce.size });
|
|
toolNonce.clear();
|
|
});
|
|
|
|
pi.registerTool({
|
|
name: 'gate0_nonce_probe',
|
|
label: 'Gate0 Nonce Probe',
|
|
description: 'Gate0-only harmless tool used to prove toolCallId to request-nonce correlation.',
|
|
parameters: Type.Object({ label: Type.String() }),
|
|
async execute(toolCallId, params) {
|
|
const broker = await brokerRequest({ action: 'promote-probe' });
|
|
log('tool_execute', { toolCallId, label: params.label, broker });
|
|
return {
|
|
content: [{ type: 'text', text: `gate0_nonce_probe executed for ${params.label}` }],
|
|
details: { harmless: true },
|
|
};
|
|
},
|
|
});
|
|
|
|
pi.registerCommand('gate0-reload', {
|
|
description: 'Trigger a real same-PID Pi extension/runtime reload.',
|
|
handler: async (_args, ctx) => {
|
|
log('reload_command_before');
|
|
await ctx.reload();
|
|
return;
|
|
},
|
|
});
|
|
}
|