Artur Mukhamadiev
1d7b867986
The bridge counted tool and model calls against fixed 100 and 50 caps, so CONFLUENCE_WEB_MAX_CONFLUENCE_CALLS / CONFLUENCE_WEB_MAX_MODEL_CALLS above the defaults were cut short inside the container, and the model cap failed the run outright. The backend now passes both values into the container at start, the supervisor forwards them to the bridge, and LIMITS reads them with the same 1-1000 validation as the backend, falling back to 100/50 when absent or invalid. Add a test that loads LIMITS in fresh processes with good and bad values for both budgets and the deadline.
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 > LIMITS.MAX_TOOL_CALLS) 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 > LIMITS.MAX_MODEL_CALLS) { 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);
|
|
});
|
|
}
|