Artur Mukhamadiev
38a8ca67f7
Pinned pi SDK 0.85.1 bridge, Python supervisor, artifact exporter, scripted backend peer, image checks and boundary checks under agent/**. Review findings F1-F3 are recorded in docs/implementation/PI_AGENT_REVIEW.md.
210 lines
11 KiB
TypeScript
210 lines
11 KiB
TypeScript
import { jsonBytes } from './json.js';
|
|
import fs from 'node:fs';
|
|
import { Agent } from '@earendil-works/pi-agent-core';
|
|
import { localTools } from './local-tools.js';
|
|
import { LIMITS, type Warning, type BridgeBaseMessage, type ModelRequest, type ModelResponse } from './types.js';
|
|
import { NDJsonFrameParser, StreamMessageWriter, sanitizeError } from './framing.js';
|
|
import { envelope, confluenceResult } from './validation.js';
|
|
import { createModelProvider } from './model-provider.js';
|
|
import { createConfluenceTools } from './confluence-tools.js';
|
|
import { scanArtifactsDirectory, exportArtifacts, closeArtifacts } from './artifacts.js';
|
|
import { supervisorChannel } from './supervision.js';
|
|
import { WarningCollector } from './warnings.js';
|
|
|
|
export type BridgeState = 'INIT' | 'RUNNING' | 'COLLECTING' | 'COMPLETE' | 'FAILED';
|
|
export interface BridgeOptions {
|
|
stdin?: NodeJS.ReadableStream; stdout?: NodeJS.WritableStream;
|
|
workDir?: string; artifactsDir?: string;
|
|
onChildReap?: () => Promise<void>; onSetDeadline?: (ms: number) => void;
|
|
}
|
|
type Pending = { type: string; tool?: string; transfer?: string; decisions?: string[]; resolve: (x: any) => void; reject: (e: Error) => void };
|
|
export class Bridge {
|
|
private state: BridgeState = 'INIT';
|
|
private seq = 0;
|
|
private seen = new Set<string>();
|
|
private pending = new Map<string, Pending>();
|
|
private writer: StreamMessageWriter;
|
|
private parser = new NDJsonFrameParser();
|
|
private input: NodeJS.ReadableStream;
|
|
private work: string;
|
|
private artifacts: string;
|
|
private agent?: Agent;
|
|
private warnings = new WarningCollector();
|
|
private completing = false;
|
|
private inputEnded = false;
|
|
private timer?: NodeJS.Timeout;
|
|
private expires = Date.now() + LIMITS.MAX_DEADLINE_MS;
|
|
private modelCalls = 0;
|
|
private toolCalls = 0;
|
|
private settled = false;
|
|
private resolve!: () => void;
|
|
private reject!: (e: Error) => void;
|
|
constructor(private options: BridgeOptions = {}) {
|
|
this.input = options.stdin || process.stdin;
|
|
this.writer = new StreamMessageWriter(options.stdout || process.stdout);
|
|
this.work = options.workDir || '/work';
|
|
this.artifacts = options.artifactsDir || `${this.work}/artifacts`;
|
|
}
|
|
nextId(): string { return `a_${++this.seq}`; }
|
|
getState(): BridgeState { return this.state; }
|
|
addWarning(w: Warning): void { this.warnings.add(w); }
|
|
private arm(ms: number): void {
|
|
this.expires = Math.min(this.expires, Date.now() + ms);
|
|
clearTimeout(this.timer);
|
|
this.timer = setTimeout(() => this.failRun('query_timeout'), Math.max(1, this.expires - Date.now()));
|
|
}
|
|
start(): Promise<void> {
|
|
const done = new Promise<void>((resolve, reject) => { this.resolve = resolve; this.reject = reject; });
|
|
this.arm(LIMITS.MAX_DEADLINE_MS);
|
|
this.parser.on('data', frame => {
|
|
try { this.inbound(frame); } catch { this.failRun('invalid_input'); }
|
|
});
|
|
this.parser.on('error', () => this.failRun('invalid_input'));
|
|
this.input.on('error', () => this.failRun('connectivity_failed'));
|
|
this.input.on('close', () => { if (this.state !== 'COMPLETE' && this.state !== 'FAILED' && !this.inputEnded) this.failRun('connectivity_failed'); });
|
|
this.parser.on('end', () => {
|
|
this.inputEnded = true;
|
|
if (this.state === 'COMPLETE') this.finishEOF();
|
|
else if (!this.completing) this.failRun('connectivity_failed');
|
|
// A peer can close stdin synchronously after observing complete, before
|
|
// the write callback. Wait for that write to succeed before resolving.
|
|
});
|
|
this.input.pipe(this.parser);
|
|
return done;
|
|
}
|
|
private finishEOF(): void {
|
|
if (this.settled) return;
|
|
this.settled = true; this.resolve();
|
|
}
|
|
private inbound(frame: BridgeBaseMessage): void {
|
|
if (this.state === 'FAILED') return;
|
|
envelope(frame);
|
|
if (this.seen.has(frame.id) || this.seen.size >= 512) throw new Error('Duplicate or excess IDs');
|
|
this.seen.add(frame.id);
|
|
if (frame.type === 'start') {
|
|
if (this.state !== 'INIT') throw new Error('Unexpected start');
|
|
this.state = 'RUNNING';
|
|
this.arm(frame.payload.remaining_ms);
|
|
this.options.onSetDeadline?.(frame.payload.remaining_ms);
|
|
this.run(frame.payload).catch(() => this.failRun('execution_failed'));
|
|
return;
|
|
}
|
|
const p = this.pending.get(frame.reply_to!);
|
|
if (!p || p.type !== frame.type) throw new Error('Unknown or wrong response');
|
|
if (frame.type === 'artifact_ack' && (p.transfer !== frame.payload.transfer_id || !p.decisions!.includes(frame.payload.decision))) throw new Error('Wrong artifact acknowledgement');
|
|
if (frame.type === 'tool_response' && frame.payload.error === null) confluenceResult(frame.payload.result, p.tool!);
|
|
this.pending.delete(frame.reply_to!);
|
|
if (frame.payload.error) {
|
|
const error = sanitizeError(frame.payload.error.code);
|
|
// Model errors are terminal, while recoverable Confluence failures are
|
|
// returned to the SDK as ordinary failed tool calls.
|
|
p.reject(new Error(error.message));
|
|
if (frame.type === 'model_response') this.failRun(error.code);
|
|
} else p.resolve(frame.payload);
|
|
}
|
|
private async send(type: string, payload: any): Promise<void> {
|
|
if (this.state === 'FAILED') throw new Error('Failed run');
|
|
await this.writer.write({ v: 1, type, id: this.nextId(), payload });
|
|
}
|
|
private request(type: string, payload: any, response: string, transfer?: string, decisions?: string[]): Promise<any> {
|
|
if (this.state === 'FAILED') return Promise.reject(new Error('Failed run'));
|
|
const id = this.nextId();
|
|
return new Promise((resolve, reject) => {
|
|
this.pending.set(id, { type: response, tool: payload.tool, transfer, decisions, resolve, reject });
|
|
this.writer.write({ v: 1, type, id, payload }).catch(() => this.failRun('connectivity_failed'));
|
|
});
|
|
}
|
|
async sendToolRequest(tool: string, parameters: Record<string, any>): Promise<any> {
|
|
if (this.state !== 'RUNNING' || this.pending.size >= 4 || ++this.toolCalls > 100) throw new Error('Tool call limit');
|
|
const payload = { tool, parameters };
|
|
if (jsonBytes(payload, LIMITS.TOOL_REQUEST_PAYLOAD_MAX_BYTES) > LIMITS.TOOL_REQUEST_PAYLOAD_MAX_BYTES) throw new Error('Tool payload limit');
|
|
return (await this.request('tool_request', payload, 'tool_response')).result;
|
|
}
|
|
async sendModelRequest(request: ModelRequest): Promise<ModelResponse> {
|
|
if (this.state !== 'RUNNING' || this.pending.size >= 4 || ++this.modelCalls > 50) { this.failRun('execution_failed'); throw new Error('Model call limit'); }
|
|
if (jsonBytes(request, LIMITS.MODEL_REQUEST_PAYLOAD_MAX_BYTES) > LIMITS.MODEL_REQUEST_PAYLOAD_MAX_BYTES) { this.failRun('model_context_exceeded'); throw new Error('Model payload limit'); }
|
|
return (await this.request('model_request', request, 'model_response')).result;
|
|
}
|
|
private async run(start: any): Promise<void> {
|
|
fs.mkdirSync(this.artifacts, { recursive: true });
|
|
const { model, streamFn } = createModelProvider(start.model, async request => {
|
|
const response = await this.sendModelRequest(request);
|
|
if (response.usage.input_tokens > start.model.context_window_tokens) { this.failRun('model_context_exceeded'); throw new Error('Context limit'); }
|
|
if (response.stop_reason === 'length' || response.usage.output_tokens > start.model.max_output_tokens) { this.failRun('model_output_limit'); throw new Error('Output limit'); }
|
|
return response;
|
|
});
|
|
this.agent = new Agent({ initialState: {
|
|
systemPrompt: start.system_instruction, model,
|
|
tools: [...createConfluenceTools((name, args) => this.sendToolRequest(name, args)), ...localTools(this.work)],
|
|
}, streamFn, toolExecution: 'sequential' });
|
|
await this.agent.prompt(start.prompt);
|
|
await this.agent.waitForIdle();
|
|
if (this.state !== 'RUNNING') return;
|
|
if (this.pending.size) throw new Error('Unanswered calls');
|
|
const final = this.agent.state.messages.at(-1);
|
|
if (!final || final.role !== 'assistant' || final.stopReason !== 'stop') throw new Error('Invalid final state');
|
|
const markdown = final.content.filter(x => x.type === 'text').map(x => x.text).join('');
|
|
if (!markdown.trim()) throw new Error('Empty final answer');
|
|
if (Buffer.byteLength(markdown) > LIMITS.FINAL_MARKDOWN_MAX_BYTES) throw new Error('Answer too large');
|
|
this.agent.clearAllQueues(); this.agent.abort();
|
|
// No host /proc fallback: collection requires a supervisor or an explicit
|
|
// test hook. Production runs always use the immutable PID-1 supervisor.
|
|
if (!this.options.onChildReap) throw new Error('Missing process supervisor');
|
|
await this.options.onChildReap();
|
|
if (this.state !== 'RUNNING') return;
|
|
const scan = scanArtifactsDirectory(this.artifacts);
|
|
try {
|
|
scan.warnings.forEach(w => this.addWarning(w));
|
|
this.state = 'COLLECTING';
|
|
await this.request('collection_start', { markdown, warnings: this.warnings.snapshot() }, 'collection_ready');
|
|
const count = await exportArtifacts(scan.candidates, {
|
|
sendBegin: async (transfer_id, name, size_bytes) => {
|
|
const ack = await this.request('artifact_begin', { transfer_id, name, size_bytes }, 'artifact_ack', transfer_id, ['accept', 'skip']);
|
|
return ack.decision;
|
|
},
|
|
sendChunk: async (transfer_id, index, data_base64) => this.send('artifact_chunk', { transfer_id, index, data_base64 }),
|
|
sendEnd: async (transfer_id, size_bytes, chunks) => {
|
|
await this.request('artifact_end', { transfer_id, size_bytes, chunks }, 'artifact_ack', transfer_id, ['stored']);
|
|
},
|
|
});
|
|
if (this.getState() === 'FAILED') return;
|
|
await this.writer.write({ v: 1, type: 'complete', id: this.nextId(), payload: { accepted_transfer_count: count } }, () => { this.completing = true; });
|
|
if (this.getState() === 'FAILED') return;
|
|
this.state = 'COMPLETE';
|
|
clearTimeout(this.timer);
|
|
if (this.inputEnded) this.finishEOF();
|
|
} finally { closeArtifacts(scan.candidates); }
|
|
}
|
|
async flushOutput(timeoutMs = 1000): Promise<boolean> {
|
|
let timer: NodeJS.Timeout | undefined;
|
|
try {
|
|
return await Promise.race([
|
|
this.writer.close().then(() => true, () => false),
|
|
new Promise<boolean>(resolve => { timer = setTimeout(() => resolve(false), timeoutMs); }),
|
|
]);
|
|
} finally { clearTimeout(timer); }
|
|
}
|
|
failRun(code: string, _raw?: unknown): void {
|
|
if (this.state === 'FAILED' || this.state === 'COMPLETE' || this.settled) return;
|
|
this.state = 'FAILED'; clearTimeout(this.timer);
|
|
this.agent?.clearAllQueues(); this.agent?.abort();
|
|
const error = sanitizeError(code);
|
|
for (const p of this.pending.values()) p.reject(new Error(error.message));
|
|
this.pending.clear();
|
|
// Failure settlement does not depend on a peer continuing to drain stdout.
|
|
this.writer.writeTerminal({ v: 1, type: 'error', id: this.nextId(), payload: error }).catch(() => {});
|
|
this.settled = true; this.reject(new Error(error.message));
|
|
}
|
|
}
|
|
if (/\bbridge\.(ts|js)$/.test(process.argv[1] || '')) {
|
|
const channel = supervisorChannel();
|
|
const bridge = new Bridge({ workDir: process.env.AGENT_WORK_DIR || '/work',
|
|
onChildReap: channel?.reap, onSetDeadline: channel?.deadline });
|
|
bridge.start().then(() => process.exit(0)).catch(async () => {
|
|
// Flush complete NDJSON lines within a bounded budget. The supervisor's
|
|
// independent deadline still owns the maximum container lifetime.
|
|
await bridge.flushOutput();
|
|
process.exit(1);
|
|
});
|
|
}
|