sessBackend.ts

import * as fs from 'fs';
import * as os from 'os';
import * as path from 'path';
import { execFile } from 'child_process';
import { promisify } from 'util';
import { randomBytes } from 'crypto';
import { atomicJson } from '../journal';
import { AgentSettings, object, Submission } from '../protocol';
import { BackendCapabilities, BackendEvent, ClientReply, InputReply, InspectionRequest, ResizePlotRequest, SessionBackend } from '../backend';
import { SessBridge } from './sessBridge';
import { SessGraphics } from './sessGraphics';
import { Arf } from './arf';
import { rString } from './rCode';

const run = promisify(execFile);
export interface SessOptions {
    frontend: 'r' | 'arf';
    ownership: 'managed' | 'adopted';
    rPath: string;
    library: string;
    resources: string;
    arfPath?: string;
    arfEndpoint?: string;
    plotBackend: 'auto' | 'jgd' | 'standard';
}

export class SessBackend implements SessionBackend {
    readonly ownership: 'managed' | 'adopted';
    readonly capabilities: BackendCapabilities;
    private listeners = new Set<(event: BackendEvent) => void>();
    private bridge: SessBridge;
    private graphics: SessGraphics;
    private frontend: Arf;
    private completion?: Record<string, unknown>;
    private outputDrained = false;
    private stderrFlushed = false;
    private orderedEvents = new Map<number, { event: Record<string, unknown>; bytes: number }>();
    private orderedOutput: ({ eventId: number } | { text: string })[] = [];
    private bufferedBytes = 0;
    private bufferedEventBytes = 0;
    private outputFailed = false;
    private runtime: string;
    private token = randomBytes(32).toString('hex');
    private current?: string;
    private closed = false;
    private disposing?: Promise<void>;
    private startup?: Promise<void>;
    constructor(private settings: AgentSettings, private options: SessOptions) {
        if (options.frontend !== 'arf') { throw new Error('R Interactive now requires arf. Start an arf session or use an ordinary R terminal.'); }
        this.ownership = options.ownership;
        this.capabilities = { streaming: options.ownership === 'managed', stdin: false, debugger: false, tables: true, html: true,
            inspection: true, interrupt: process.platform !== 'win32', jgd: false, plotResize: false, restart: options.ownership === 'managed' };
        this.runtime = fs.mkdtempSync(path.join(os.tmpdir(), 'r-b-'));
        fs.chmodSync(this.runtime, 0o700);
        this.bridge = new SessBridge(this.runtime, this.token, event => this.emit(event),
            message => this.receive(message), (method, params) => this.notification(method, params, this.current));
        this.graphics = new SessGraphics({ ...options, historyLimit: settings.historyLimit }, event => this.emit(event));
        const output = (event: BackendEvent): void => {
            if (this.closed || this.outputFailed && event.type === 'stream') { return; }
            if (event.type === 'stream' && event.channel === 'stdout' && this.ownership === 'managed') {
                this.orderedOutput.push({ text: event.text }); this.bufferedBytes += Buffer.byteLength(event.text);
                this.drainEvents();
            } else { this.emit(event.type === 'stream' ? { ...event, executionId: this.current } : event); }
        };
        this.frontend = new Arf(options.arfPath ?? 'arf', settings.directory, options.arfEndpoint, this.ownership === 'adopted', output,
            this.token, (message, channel) => this.runtimeEvent(message, channel));
    }
    onEvent(listener: (event: BackendEvent) => void): () => void {
        this.listeners.add(listener); return () => this.listeners.delete(listener);
    }
    private emit(event: BackendEvent): void {
        if (this.closed) { return; }
        if ((event.type === 'metadata' || event.type === 'ready') && event.metadata.rPid) { this.frontend.process.setPid(event.metadata.rPid); }
        if (event.type === 'clientRequest' && !event.executionId) { event = { ...event, executionId: this.current }; }
        if (event.type === 'exit') { this.graphics.flush(); this.graphics.dispose(); }
        for (const listener of this.listeners) { listener(event); }
    }
    start(): Promise<void> { return this.startup ??= this.startRuntime(); }
    private async startRuntime(): Promise<void> {
        const libraries = await run(this.options.rPath, ['--vanilla', '--slave', '-e', 'cat(.libPaths(), sep="\\n")'], { timeout: 15000 });
        if (this.closed) { throw new Error('Backend disposed during startup'); }
        const support = libraries.stdout.trim().split(/\r?\n/).filter(Boolean);
        await this.bridge.start();
        if (this.closed) { throw new Error('Backend disposed during startup'); }
        await this.graphics.start();
        if (this.closed) { this.graphics.dispose(); throw new Error('Backend disposed during startup'); }
        this.capabilities.jgd = this.graphics.enabled; this.capabilities.plotResize = this.graphics.enabled;
        const bootstrap = path.join(this.runtime, 'bootstrap.json');
        atomicJson(bootstrap, { sess: path.join(this.runtime, 'sess.sock'),
            jgd: this.graphics.endpoint, token: this.token, useJgd: this.graphics.enabled });
        const code = `base::source(${rString(path.join(this.options.resources, 'interactive-worker.R'))}, local=base::new.env(parent=base::baseenv()))$value(` +
            `${rString(this.options.library)}, ${rString(bootstrap)}, c(${support.map(rString).join(',')}), managed=${this.ownership === 'managed' ? 'TRUE' : 'FALSE'})`;
        await this.frontend.start(code, { ...process.env, SESS_ENDPOINT: path.join(this.runtime, 'sess.sock'), JGD_SOCKET: this.graphics.endpoint });
    }
    async dispatch(submission: Submission): Promise<void> {
        await this.startup;
        if (this.closed) { throw new Error('Backend is disposed'); }
        this.current = submission.id;
        this.completion = undefined; this.stderrFlushed = false; this.outputDrained = false;
        await this.frontend.dispatch(submission);
        this.outputDrained = true;
        if (this.ownership === 'adopted' && this.completion) { this.finish(this.completion); }
    }
    inspect(request: InspectionRequest): Promise<unknown> { return this.bridge.request(request.method, request.params, request.timeout ?? 5000); }
    replyInput(_reply: InputReply): Promise<void> { return Promise.reject(new Error('Notebook console input is not supported by arf IPC. Use an ordinary R terminal.')); }
    replyClientRequest(id: string, reply: ClientReply): Promise<void> { this.bridge.reply(id, reply); return Promise.resolve(); }
    interrupt(): Promise<void> { this.frontend.process.interrupt(); return Promise.resolve(); }
    async stop(options?: { force?: boolean }): Promise<void> {
        // Stop may arrive after the control endpoint opens but before R is spawned.
        await this.startup;
        await this.frontend.process.stop(options?.force);
    }
    resizePlot(request: ResizePlotRequest): Promise<void> { this.graphics.resize(request); return Promise.resolve(); }
    dispose(): Promise<void> { return this.disposing ??= this.disposeRuntime(); }
    private async disposeRuntime(): Promise<void> {
        this.closed = true; this.listeners.clear();
        this.orderedOutput = []; this.orderedEvents.clear(); this.bufferedBytes = 0; this.bufferedEventBytes = 0;
        try {
            if (this.ownership === 'adopted') {
                // Best effort through the existing bridge, without evaluating user code or interrupting R.
                await this.bridge.request('interactive_stop', {}, 500).catch(() => undefined);
            }
        } finally {
            this.frontend.dispose(); this.graphics.dispose(); this.bridge.dispose();
            fs.rmSync(this.runtime, { recursive: true, force: true });
        }
    }
    private runtimeEvent(message: Record<string, unknown>, channel?: string): void {
        if (this.closed || this.outputFailed) { return; }
        const executionId = typeof message.executionId === 'string' && message.executionId ? message.executionId : undefined;
        switch (message.type) {
            case 'event':
                if (channel === 'stdout' && Number.isSafeInteger(message.eventId)) {
                    this.orderedOutput.push({ eventId: Number(message.eventId) }); this.drainEvents();
                }
                break;
            case 'ready':
                this.frontend.ready = true;
                this.emit({ type: 'ready', capabilities: { ...this.capabilities }, metadata: {
                    rPid: Number(message.pid), rVersion: String(message.version), runtimeSessionId: String(message.sessionId),
                    rPath: typeof message.rPath === 'string' ? message.rPath : this.options.rPath,
                    libraryPaths: Array.isArray(message.libraryPaths) ? message.libraryPaths.map(String) : [] } });
                break;
            case 'started': this.frontend.markStarted(); this.emit({ type: 'started', executionId }); break;
            case 'flush':
                if (channel === 'stderr') {
                    this.stderrFlushed = true;
                    if (this.completion) { this.finish(this.completion); }
                }
                break;
            case 'finished':
                this.completion = message;
                if (this.ownership === 'managed' ? this.stderrFlushed : this.outputDrained) { this.finish(message); }
                break;
            case 'stream': this.emit({ type: 'stream', executionId, text: typeof message.text === 'string' ? message.text : '',
                channel: typeof message.channel === 'string' ? message.channel : 'stdout', external: !executionId && this.options.frontend === 'arf' }); break;
            case 'condition': case 'input': case 'display': this.emit({ type: message.type, data: message, executionId }); break;
            case 'truncated': this.emit({ type: 'truncated', message: String(message.message), executionId }); break;
            case 'notification': {
                // R encodes an empty list as []; task callbacks send notifications
                // such as workspace_updated without any parameters.
                const params = message.params;
                this.notification(String(message.method), object(Array.isArray(params) && !params.length ? {} : params ?? {}), executionId);
                break;
            }
            case 'external': this.emit({ type: 'external', code: typeof message.code === 'string' ? message.code : '# Terminal output', success: Boolean(message.success) }); break;
        }
    }
    private receive(message: Record<string, unknown>): void {
        if (this.closed || this.outputFailed) { return; }
        if (this.ownership === 'adopted') { this.runtimeEvent(message); return; }
        if (Number.isSafeInteger(message.eventId)) {
            const eventId = Number(message.eventId), event = object(message.event);
            const previous = this.orderedEvents.get(eventId);
            const bytes = Buffer.byteLength(JSON.stringify(event));
            this.bufferedEventBytes += bytes - (previous?.bytes ?? 0);
            this.orderedEvents.set(eventId, { event, bytes }); this.drainEvents();
        }
    }
    private drainEvents(): void {
        while (this.orderedOutput.length) {
            const item = this.orderedOutput[0];
            if ('eventId' in item) {
                const event = this.orderedEvents.get(item.eventId);
                if (!event) { break; }
                this.orderedOutput.shift(); this.orderedEvents.delete(item.eventId);
                this.bufferedEventBytes -= event.bytes; this.runtimeEvent(event.event);
            } else {
                this.orderedOutput.shift(); this.bufferedBytes -= Buffer.byteLength(item.text);
                this.emit({ type: 'stream', text: item.text, channel: 'stdout', executionId: this.current });
            }
        }
        if (this.bufferedBytes > 4 * 1024 * 1024 || this.bufferedEventBytes > 16 * 1024 * 1024 ||
            this.orderedOutput.length > 1024 || this.orderedEvents.size > 1024) {
            this.outputFailed = true;
            this.orderedOutput = []; this.orderedEvents.clear(); this.bufferedBytes = 0; this.bufferedEventBytes = 0;
            this.emit({ type: 'unavailable', message: 'R output references missing Interactive events; restart the session before submitting more code.' });
        }
    }
    private finish(message: Record<string, unknown>): void {
        const executionId = String(message.executionId);
        if (this.current !== executionId) { return; }
        this.graphics.flush(executionId);
        this.current = undefined; this.completion = undefined;
        this.emit({ type: 'finished', executionId,
            state: message.state === 'success' || message.state === 'interrupted' ? message.state : 'error' });
    }
    private notification(method: string, params: Record<string, unknown>, executionId?: string): void {
        if (['webview', 'page_viewer', 'browser'].includes(method) && typeof params.url === 'string') {
            this.emit({ type: 'display', executionId, data: /^https?:\/\//.test(params.url)
                ? { kind: 'url', url: params.url } : { kind: 'html', file: params.url } });
        } else if (method === 'workspace_updated') { this.emit({ type: 'workspaceChanged' }); }
        else { this.emit({ type: method === 'dataview' ? 'viewer' : 'notification', method, params, executionId }); }
    }
}