Spaces:
Running
Running
| import { mkdir } from 'node:fs/promises'; | |
| import { join } from 'node:path'; | |
| const send = event => { if (process.connected) process.send(event); }; | |
| let session, running = false, disposed = false, requestNumber = 0, modelRequestNumber = 0, compaction = false, contextWindow = 0; | |
| const pendingTools = new Map(); | |
| const schema = properties => ({ type: 'object', properties, required: Object.keys(properties).filter(key => !['timeout', 'timeout_seconds'].includes(key)), additionalProperties: false }); | |
| const string = description => ({ type: 'string', description }); | |
| export function toolSpecifications(condition, id) { | |
| const tools = [ | |
| ['bash', 'Run shell commands only inside your assigned isolated container checkout. Use for reading, editing, compilation and tests. Each call starts in the checkout.', { command: string('Shell command'), timeout: { type: 'number', minimum: 1, maximum: 600 } }], | |
| ]; | |
| if (condition !== 'single') tools.push( | |
| ['send_message', 'Send an attributed message to another persistent peer. It resumes an idle recipient.', { agent_id: string('Recipient agent ID'), message: string('Message') }], | |
| ['wait_for_message', 'Wait briefly for an incoming message. This tool returns the message or a timeout.', { timeout_seconds: { type: 'number', minimum: 1, maximum: 120 } }], | |
| ['agent_status', 'List persistent agent IDs and lifecycle state.', {}], | |
| ); | |
| if (condition === 'async' && id === 'agent-0') tools.push( | |
| ['spawn_agent', 'Create a persistent child for this task and return immediately. The child sees your task and shared operating rules, not your original conversation. At most five agents may be active.', { task: string('Complete delegated task instructions') }], | |
| ['resume_agent', 'Give an existing persistent child more work without resetting its history or budget.', { agent_id: string('Agent ID'), message: string('New instructions') }], | |
| ['terminate_agent', 'Abort a child and permanently mark it terminated.', { agent_id: string('Child ID') }], | |
| ); | |
| return tools.map(([name, description, properties]) => ({ name, label: name, description, parameters: schema(properties) })); | |
| } | |
| export async function createSession(meta, sdk, emit = send) { | |
| const { config, id, directory } = meta; | |
| await mkdir(directory, { recursive: true }); | |
| const authStorage = sdk.AuthStorage.inMemory(); | |
| authStorage.setRuntimeApiKey('study-proxy', process.env.PI_STUDY_API_KEY ?? 'study-proxy'); | |
| const registry = sdk.ModelRegistry.inMemory(authStorage); | |
| const modelId = config.model_id ?? 'Beam-501B-A23B'; | |
| registry.registerProvider('study-proxy', { | |
| baseUrl: `${config.proxy_url.replace(/\/$/, '')}/runs/${encodeURIComponent(config.run_id)}/agents/${encodeURIComponent(id)}/v1`, | |
| api: 'openai-completions', | |
| apiKey: '$PI_STUDY_API_KEY', | |
| models: [{ id: modelId, name: modelId, reasoning: true, input: ['text'], contextWindow: config.context_window ?? 262144, maxTokens: config.max_output_tokens ?? 16000, | |
| cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, | |
| compat: { supportsStore: false, supportsDeveloperRole: true, supportsReasoningEffort: true, supportsUsageInStreaming: true, maxTokensField: 'max_tokens', reasoningEffortMap: { minimal: 'low', low: 'low', medium: 'medium', high: 'high', xhigh: 'xhigh' } }, | |
| }], | |
| }); | |
| const settings = sdk.SettingsManager.inMemory({ | |
| compaction: { enabled: true, reserveTokens: 20000, keepRecentTokens: 12000 }, | |
| retry: { enabled: false, provider: { maxRetries: 0, timeoutMs: 600000 } }, | |
| enableAnalytics: false, enableInstallTelemetry: false, packages: [], | |
| }); | |
| const loader = new sdk.DefaultResourceLoader({ cwd: directory, agentDir: directory, settingsManager: settings, noExtensions: true, noSkills: true, noPromptTemplates: true, noThemes: true, noContextFiles: true, systemPromptOverride: () => 'You are a coding agent reconstructing software from its observable behavior. Follow the study task and operating instructions. You have only explicitly listed tools. All shell actions execute in the isolated solver container.' }); | |
| await loader.reload(); | |
| const specifications = toolSpecifications(config.condition, id); | |
| const customTools = specifications.map(spec => ({ ...spec, execute: async (_toolCallId, args, signal) => { | |
| const request_id = `${id}-tool-${++requestNumber}`; | |
| return new Promise(resolve => { | |
| const abort = () => { pendingTools.delete(request_id); resolve({ content: [{ type: 'text', text: 'Tool cancelled' }], isError: true }); }; | |
| if (signal?.aborted) { abort(); return; } | |
| signal?.addEventListener('abort', abort, { once: true }); | |
| pendingTools.set(request_id, result => { signal?.removeEventListener('abort', abort); resolve(result); }); | |
| emit({ type: 'tool_request', request_id, name: spec.name, args }); | |
| }); | |
| } })); | |
| const { session: created } = await sdk.createAgentSession({ | |
| cwd: directory, agentDir: directory, authStorage, modelRegistry: registry, | |
| model: registry.find('study-proxy', modelId), thinkingLevel: config.thinking_level ?? 'medium', | |
| noTools: 'builtin', tools: specifications.map(tool => tool.name), customTools, | |
| resourceLoader: loader, settingsManager: settings, | |
| sessionManager: sdk.SessionManager.create(directory, join(directory, 'sessions')), | |
| }); | |
| const enabled = created.getActiveToolNames(); | |
| if (enabled.some(name => !specifications.some(spec => spec.name === name)) || enabled.length !== specifications.length) throw new Error(`Unexpected active tools: ${enabled.join(',')}`); | |
| const originalStream = created.agent.streamFn; | |
| created.agent.streamFn = async (model, context, options = {}) => { | |
| const request_id = `${id}-model-${++modelRequestNumber}`; | |
| const started = performance.now(); | |
| emit({ type: 'model_request_start', request_id, purpose: compaction ? 'compaction' : 'agent', context }); | |
| const priorPayload = options.onPayload; | |
| const stream = await originalStream(model, context, { ...options, maxTokens: config.max_output_tokens ?? 16000, reasoning: config.thinking_level ?? 'medium', | |
| headers: { ...options.headers, 'X-Study-Request-ID': request_id, 'X-Study-Purpose': compaction ? 'compaction' : 'agent', 'X-Study-Context-Window': `${id}-${contextWindow}` }, | |
| onPayload: async (payload, requestModel) => { | |
| const before = priorPayload ? await priorPayload(payload, requestModel) : payload; | |
| const effective = { ...(before ?? payload), temperature: config.temperature ?? 0.7, top_p: config.top_p ?? 0.9, max_tokens: config.max_output_tokens ?? 16000, reasoning_effort: config.thinking_level ?? 'medium' }; | |
| delete effective.max_completion_tokens; | |
| return effective; | |
| }, | |
| }); | |
| void stream.result().then(message => emit({ type: 'model_request_end', request_id, duration_ms: performance.now() - started, usage: message.usage, stop_reason: message.stopReason, error: message.errorMessage }), error => emit({ type: 'model_request_error', request_id, error: String(error) })); | |
| return stream; | |
| }; | |
| created.subscribe(event => { | |
| if (event.type === 'compaction_start') compaction = true; | |
| if (event.type === 'compaction_end') { compaction = false; if (event.result && !event.aborted) contextWindow++; } | |
| if (event.type === 'message_start' && event.message?.role === 'user') { | |
| const content = typeof event.message.content === 'string' ? event.message.content : event.message.content.filter(block => block.type === 'text').map(block => block.text).join('\n'); | |
| const delivered = content.match(/^Message from (agent-\d+) \(([^;]+); event (\d+)\):/); | |
| if (delivered) emit({ type: 'message_delivered', sender_id: delivered[1], kind: delivered[2], causal_event_id: Number(delivered[3]), delivery: 'context_insertion' }); | |
| } | |
| // Final messages preserve usage/thinking; streaming deltas would duplicate large payloads. | |
| if (event.type !== 'message_update' && event.type !== 'tool_execution_update') emit({ type: 'pi_event', event }); | |
| }); | |
| emit({ type: 'session_created', session_id: created.sessionId, session_file: created.sessionFile, tools: enabled }); | |
| return created; | |
| } | |
| async function prompt(message) { | |
| if (disposed) return; | |
| if (running) { await session.steer(message); return; } | |
| running = true; | |
| let error; | |
| try { await session.prompt(message, { expandPromptTemplates: false }); } | |
| catch (caught) { error = String(caught); } | |
| finally { | |
| running = false; | |
| const last = [...session.messages].reverse().find(message => message.role === 'assistant'); | |
| if (last?.stopReason === 'error') error = last.errorMessage ?? 'Provider error'; | |
| send({ type: 'turn_done', answer: last?.content.filter(block => block.type === 'text').map(block => block.text).join('\n') ?? '', error }); | |
| } | |
| } | |
| if (process.send) { | |
| process.on('message', async message => { | |
| try { | |
| if (message.type === 'init') { | |
| const sdk = await import('@earendil-works/pi-coding-agent'); | |
| session = await createSession(message, sdk); send({ type: 'ready' }); | |
| } else if (message.type === 'prompt') { void prompt(message.message).catch(error => send({ type: 'worker_error', error: String(error) })); } | |
| else if (message.type === 'tool_result') { pendingTools.get(message.request_id)?.(message.result); pendingTools.delete(message.request_id); } | |
| else if (message.type === 'abort') { await session?.abort(); session?.abortCompaction(); } | |
| else if (message.type === 'dispose') { disposed = true; await session?.abort(); session?.abortCompaction(); session?.dispose(); process.exit(0); } | |
| } catch (error) { send({ type: 'worker_error', error: String(error) }); } | |
| }); | |
| process.on('disconnect', () => { void session?.abort(); process.exit(0); }); | |
| } | |