@nordbyte/nordrelay 0.2.1

package/dist/pi-session.js

@@ -0,0 +1,573 @@
+import { spawnSync } from "node:child_process";
+import { readFile } from "node:fs/promises";
+import path from "node:path";
+import { PI_AGENT_CAPABILITIES, PI_THINKING_LEVELS, } from "./agent.js";
+import { createDefaultLaunchProfile } from "./codex-launch.js";
+import { resolvePiCli } from "./pi-cli.js";
+import { PiRpcClient } from "./pi-rpc.js";
+import { getPiSession, listPiSessions, listPiWorkspaces, readPiSessionRecord, resolvePiSessionDir, } from "./pi-state.js";
+export class PiSessionService {
+    config;
+    sessionDir;
+    cliPath;
+    launchProfile;
+    rpc = null;
+    currentWorkspace;
+    currentThreadId = null;
+    currentSessionPath;
+    currentModel;
+    currentThinking;
+    currentSessionName;
+    processing = false;
+    cachedStats = {};
+    constructor(config) {
+        this.config = config;
+        const cli = resolvePiCli(process.env, config.piCliPath);
+        if (!cli.path) {
+            throw new Error("Pi CLI not found. Install Pi from https://pi.dev/ or set PI_CLI_PATH.");
+        }
+        this.cliPath = cli.path;
+        this.sessionDir = resolvePiSessionDir({ sessionDir: config.piSessionDir });
+        this.currentWorkspace = config.workspace;
+        this.currentModel = config.piDefaultModel;
+        this.currentThinking = config.piDefaultThinking;
+        this.launchProfile = createDefaultLaunchProfile("workspace-write", "never");
+    }
+    static async create(config, options) {
+        const service = new PiSessionService(config);
+        service.currentWorkspace = options?.workspace ?? config.workspace;
+        service.currentModel = options?.model ?? config.piDefaultModel;
+        service.currentThinking = options?.reasoningEffort ?? config.piDefaultThinking;
+        if (options?.sessionPath) {
+            await service.switchSession(options.sessionPath);
+            return service;
+        }
+        if (options?.resumeThreadId) {
+            await service.resumeThread(options.resumeThreadId);
+            return service;
+        }
+        if (!options?.deferThreadStart) {
+            await service.newThread(service.currentWorkspace, service.currentModel);
+        }
+        return service;
+    }
+    getInfo() {
+        return {
+            agentId: "pi",
+            agentLabel: "Pi",
+            threadId: this.currentThreadId,
+            workspace: this.currentWorkspace,
+            model: this.currentModel,
+            reasoningEffort: this.currentThinking,
+            launchProfileId: "pi-rpc",
+            launchProfileLabel: "Pi RPC",
+            launchProfileBehavior: "rpc / host",
+            sandboxMode: "host",
+            approvalPolicy: "never",
+            fastMode: false,
+            unsafeLaunch: false,
+            sessionPath: this.currentSessionPath,
+            sessionUsage: this.cachedStats.sessionUsage,
+            contextUsage: this.cachedStats.contextUsage,
+            capabilities: PI_AGENT_CAPABILITIES,
+        };
+    }
+    isProcessing() {
+        return this.processing;
+    }
+    getActiveThreadId() {
+        return this.currentThreadId;
+    }
+    hasActiveThread() {
+        return Boolean(this.currentThreadId || this.currentSessionPath);
+    }
+    getCurrentWorkspace() {
+        return this.currentWorkspace;
+    }
+    async prompt(input, callbacks) {
+        if (this.processing) {
+            throw new Error("A Pi turn is already in progress");
+        }
+        await this.ensureSessionStarted();
+        const rpc = this.getRpc();
+        const promptPayload = await this.buildPromptPayload(input);
+        const lastToolOutput = new Map();
+        let didEnd = false;
+        this.processing = true;
+        const off = rpc.onEvent((event) => {
+            try {
+                switch (event.type) {
+                    case "message_update":
+                        handleMessageUpdate(event, callbacks);
+                        break;
+                    case "tool_execution_start": {
+                        const toolCallId = stringValue(event.toolCallId) ?? "tool";
+                        const toolName = stringValue(event.toolName) ?? "tool";
+                        lastToolOutput.set(toolCallId, "");
+                        callbacks.onToolStart(toolName, toolCallId);
+                        break;
+                    }
+                    case "tool_execution_update": {
+                        const toolCallId = stringValue(event.toolCallId) ?? "tool";
+                        const text = extractContentText(objectValue(event.partialResult));
+                        const previous = lastToolOutput.get(toolCallId) ?? "";
+                        const delta = computeTextDelta(previous, text);
+                        lastToolOutput.set(toolCallId, text);
+                        if (delta) {
+                            callbacks.onToolUpdate(toolCallId, delta);
+                        }
+                        break;
+                    }
+                    case "tool_execution_end": {
+                        const toolCallId = stringValue(event.toolCallId) ?? "tool";
+                        const resultText = extractContentText(objectValue(event.result));
+                        const previous = lastToolOutput.get(toolCallId) ?? "";
+                        const delta = computeTextDelta(previous, resultText);
+                        if (delta) {
+                            callbacks.onToolUpdate(toolCallId, delta);
+                        }
+                        callbacks.onToolEnd(toolCallId, Boolean(event.isError));
+                        break;
+                    }
+                    case "turn_end":
+                        this.refreshFromTurnEnd(event, callbacks);
+                        break;
+                    case "agent_end":
+                        didEnd = true;
+                        callbacks.onAgentEnd();
+                        break;
+                    case "extension_error": {
+                        const toolCallId = stringValue(event.id) ?? "extension-error";
+                        callbacks.onToolStart("extension_error", toolCallId);
+                        callbacks.onToolUpdate(toolCallId, stringValue(event.error) ?? "Extension error");
+                        callbacks.onToolEnd(toolCallId, true);
+                        break;
+                    }
+                    default:
+                        break;
+                }
+            }
+            catch (error) {
+                console.error("Failed to handle Pi RPC event:", error);
+            }
+        });
+        try {
+            await rpc.send({ type: "prompt", ...promptPayload }, 30_000);
+            await this.waitForAgentEnd(() => didEnd);
+            await this.refreshState().catch(() => { });
+            await this.refreshStats().catch(() => { });
+        }
+        finally {
+            off();
+            this.processing = false;
+        }
+    }
+    async abort() {
+        if (!this.rpc) {
+            this.processing = false;
+            return;
+        }
+        await this.rpc.send({ type: "abort" }, 10_000).catch(() => { });
+        this.processing = false;
+    }
+    async newThread(workspace, model) {
+        this.ensureIdle("start a new Pi session");
+        this.currentWorkspace = workspace ?? this.currentWorkspace;
+        if (model) {
+            this.currentModel = model;
+        }
+        this.currentThreadId = null;
+        this.currentSessionPath = undefined;
+        this.cachedStats = {};
+        this.restartRpc();
+        await this.refreshState();
+        return this.getInfo();
+    }
+    async resumeThread(threadId) {
+        return this.switchSession(threadId);
+    }
+    async switchSession(threadId) {
+        this.ensureIdle("switch Pi session");
+        const record = getPiSession(threadId, { sessionDir: this.sessionDir });
+        if (!record) {
+            throw new Error(`Unknown Pi session: ${threadId}`);
+        }
+        this.currentWorkspace = record.cwd;
+        this.currentThreadId = record.id;
+        this.currentSessionPath = record.sessionPath;
+        this.currentModel = record.model ?? this.currentModel;
+        this.currentThinking = record.reasoningEffort ?? this.currentThinking;
+        this.cachedStats = {};
+        if (this.rpc) {
+            const result = await this.rpc.send({ type: "switch_session", sessionPath: record.sessionPath }, 30_000);
+            if (objectValue(result.data)?.cancelled === true) {
+                throw new Error("Pi session switch was cancelled by an extension");
+            }
+        }
+        else {
+            this.restartRpc();
+        }
+        await this.refreshState();
+        await this.refreshStats().catch(() => { });
+        return this.getInfo();
+    }
+    listAllSessions(limit) {
+        return listPiSessions(limit ?? 20, { sessionDir: this.sessionDir });
+    }
+    listWorkspaces() {
+        const workspaces = new Set(listPiWorkspaces({ sessionDir: this.sessionDir }));
+        workspaces.add(this.currentWorkspace);
+        workspaces.add(this.config.workspace);
+        return [...workspaces].sort((left, right) => left.localeCompare(right));
+    }
+    listModels() {
+        const result = spawnSync(this.cliPath, ["--list-models"], {
+            cwd: this.currentWorkspace,
+            env: process.env,
+            encoding: "utf8",
+            timeout: 10_000,
+        });
+        if (result.status !== 0 || !result.stdout) {
+            return this.currentModel ? [{ slug: this.currentModel, displayName: this.currentModel }] : [];
+        }
+        const records = [];
+        for (const line of result.stdout
+            .split(/\r?\n/)
+            .slice(1)
+            .map((line) => line.trim())
+            .filter(Boolean)) {
+            const parts = line.split(/\s+/);
+            const provider = parts[0];
+            const model = parts[1];
+            if (!provider || !model) {
+                continue;
+            }
+            const slug = `${provider}/${model}`;
+            const maxInputTokens = parseCompactTokenCount(parts[2]);
+            records.push({
+                slug,
+                displayName: slug,
+                ...(maxInputTokens !== undefined ? { maxInputTokens } : {}),
+            });
+        }
+        if (this.currentModel && !records.some((record) => record.slug === this.currentModel)) {
+            records.unshift({ slug: this.currentModel, displayName: this.currentModel });
+        }
+        return records;
+    }
+    getSessionRecord(threadId) {
+        return getPiSession(threadId, { sessionDir: this.sessionDir });
+    }
+    setModel(slug) {
+        this.currentModel = slug;
+        this.restartRpcIfIdle();
+        return slug;
+    }
+    async setModelForCurrentSession(slug) {
+        this.ensureIdle("change Pi model");
+        this.currentModel = slug;
+        let appliedToActiveThread = false;
+        if (this.rpc && this.currentThreadId) {
+            const { provider, modelId } = splitPiModelSlug(slug);
+            await this.rpc.send({ type: "set_model", provider, modelId }, 30_000);
+            await this.refreshState().catch(() => { });
+            appliedToActiveThread = true;
+        }
+        else {
+            this.restartRpcIfIdle();
+        }
+        return { value: slug, appliedToActiveThread };
+    }
+    setReasoningEffort(effort) {
+        this.currentThinking = normalizePiThinking(effort);
+        this.restartRpcIfIdle();
+    }
+    async setReasoningEffortForCurrentSession(effort) {
+        this.ensureIdle("change Pi thinking level");
+        const level = normalizePiThinking(effort);
+        this.currentThinking = level;
+        let appliedToActiveThread = false;
+        if (this.rpc && this.currentThreadId) {
+            await this.rpc.send({ type: "set_thinking_level", level }, 30_000);
+            await this.refreshState().catch(() => { });
+            appliedToActiveThread = true;
+        }
+        else {
+            this.restartRpcIfIdle();
+        }
+        return { value: level, appliedToActiveThread };
+    }
+    setLaunchProfile() {
+        throw new Error("Launch profiles are only supported by Codex sessions");
+    }
+    setFastMode() {
+        throw new Error("Fast mode is only supported by Codex sessions");
+    }
+    getSelectedLaunchProfile() {
+        return this.launchProfile;
+    }
+    syncFromCodexState() {
+        return {
+            threadId: this.currentThreadId,
+            changed: false,
+            reattached: false,
+            changedFields: [],
+            info: this.getInfo(),
+        };
+    }
+    handback() {
+        const threadId = this.currentThreadId;
+        const workspace = this.currentWorkspace;
+        const sessionPath = this.currentSessionPath;
+        this.rpc?.stop();
+        this.rpc = null;
+        this.currentThreadId = null;
+        this.currentSessionPath = undefined;
+        return {
+            threadId,
+            workspace,
+            command: sessionPath
+                ? `cd ${shellQuote(workspace)} && pi --session ${shellQuote(sessionPath)}`
+                : undefined,
+            label: "Pi CLI",
+        };
+    }
+    dispose() {
+        this.rpc?.stop();
+        this.rpc = null;
+        this.processing = false;
+    }
+    async ensureSessionStarted() {
+        if (this.currentThreadId || this.currentSessionPath) {
+            this.getRpc().ensureStarted();
+            await this.refreshState().catch(() => { });
+            return;
+        }
+        await this.newThread(this.currentWorkspace, this.currentModel);
+    }
+    restartRpc() {
+        this.rpc?.stop();
+        this.rpc = new PiRpcClient({
+            commandPath: this.cliPath,
+            cwd: this.currentWorkspace,
+            sessionDir: this.sessionDir,
+            sessionPath: this.currentSessionPath,
+            model: this.currentModel,
+            thinking: this.currentThinking,
+            env: { PI_CODING_AGENT_SESSION_DIR: this.sessionDir },
+        });
+    }
+    restartRpcIfIdle() {
+        if (this.processing) {
+            return;
+        }
+        this.restartRpc();
+    }
+    getRpc() {
+        if (!this.rpc) {
+            this.restartRpc();
+        }
+        return this.rpc;
+    }
+    ensureIdle(action) {
+        if (this.processing) {
+            throw new Error(`Cannot ${action} while a turn is in progress`);
+        }
+    }
+    async refreshState() {
+        const response = await this.getRpc().send({ type: "get_state" }, 30_000);
+        const data = objectValue(response.data);
+        if (!data) {
+            return;
+        }
+        const model = objectValue(data.model);
+        const provider = stringValue(model?.provider);
+        const modelId = stringValue(model?.id);
+        this.currentModel = provider && modelId ? `${provider}/${modelId}` : modelId ?? this.currentModel;
+        this.currentThinking = stringValue(data.thinkingLevel) ?? this.currentThinking;
+        this.currentSessionPath = stringValue(data.sessionFile) ?? this.currentSessionPath;
+        this.currentThreadId = stringValue(data.sessionId) ?? this.currentThreadId;
+        this.currentSessionName = stringValue(data.sessionName) ?? this.currentSessionName;
+        if (this.currentSessionPath) {
+            const record = readPiSessionRecord(this.currentSessionPath, path.basename(path.dirname(this.currentSessionPath)));
+            if (record?.cwd) {
+                this.currentWorkspace = record.cwd;
+            }
+        }
+    }
+    async refreshStats() {
+        const response = await this.getRpc().send({ type: "get_session_stats" }, 30_000);
+        const data = objectValue(response.data);
+        if (!data) {
+            return;
+        }
+        const tokens = objectValue(data.tokens);
+        const contextUsage = objectValue(data.contextUsage);
+        this.cachedStats = {
+            sessionUsage: tokens
+                ? {
+                    input: numberValue(tokens.input) ?? 0,
+                    output: numberValue(tokens.output) ?? 0,
+                    cacheRead: numberValue(tokens.cacheRead) ?? 0,
+                    cacheWrite: numberValue(tokens.cacheWrite) ?? 0,
+                    total: numberValue(tokens.total) ?? 0,
+                    cost: numberValue(data.cost) ?? undefined,
+                }
+                : undefined,
+            contextUsage: contextUsage
+                ? {
+                    tokens: numberValue(contextUsage.tokens),
+                    contextWindow: numberValue(contextUsage.contextWindow),
+                    percent: numberValue(contextUsage.percent),
+                }
+                : undefined,
+        };
+    }
+    async buildPromptPayload(input) {
+        if (typeof input === "string") {
+            return { message: input };
+        }
+        const textParts = [input.stagedFileInstructions, input.text].filter((part) => Boolean(part?.trim()));
+        const images = await Promise.all((input.imagePaths ?? []).map(async (imagePath) => ({
+            type: "image",
+            data: (await readFile(imagePath)).toString("base64"),
+            mimeType: mimeTypeForImage(imagePath),
+        })));
+        return {
+            message: textParts.join("\n\n") || "Please inspect the attached file(s).",
+            ...(images.length > 0 ? { images } : {}),
+        };
+    }
+    refreshFromTurnEnd(event, callbacks) {
+        const message = objectValue(event.message);
+        const usage = objectValue(message?.usage);
+        if (!usage) {
+            return;
+        }
+        callbacks.onTurnComplete?.({
+            inputTokens: numberValue(usage.input) ?? 0,
+            cachedInputTokens: numberValue(usage.cacheRead) ?? 0,
+            outputTokens: numberValue(usage.output) ?? 0,
+        });
+    }
+    async waitForAgentEnd(done) {
+        if (done()) {
+            return;
+        }
+        await new Promise((resolve, reject) => {
+            const startedAt = Date.now();
+            const timer = setInterval(() => {
+                if (done()) {
+                    clearInterval(timer);
+                    resolve();
+                    return;
+                }
+                if (!this.processing) {
+                    clearInterval(timer);
+                    resolve();
+                    return;
+                }
+                if (Date.now() - startedAt > 24 * 60 * 60 * 1000) {
+                    clearInterval(timer);
+                    reject(new Error("Timed out waiting for Pi agent to finish"));
+                }
+            }, 250);
+            timer.unref?.();
+        });
+    }
+}
+function handleMessageUpdate(event, callbacks) {
+    const update = objectValue(event.assistantMessageEvent);
+    const updateType = stringValue(update?.type);
+    if (updateType === "text_delta") {
+        const delta = stringValue(update?.delta);
+        if (delta) {
+            callbacks.onTextDelta(delta);
+        }
+    }
+    else if (updateType === "toolcall_end") {
+        const toolCall = objectValue(update?.toolCall);
+        const id = stringValue(toolCall?.id);
+        const name = stringValue(toolCall?.name);
+        if (id && name) {
+            callbacks.onToolStart(name, id);
+        }
+    }
+    else if (updateType === "error") {
+        const toolId = "pi-message-error";
+        callbacks.onToolStart("message_error", toolId);
+        callbacks.onToolUpdate(toolId, stringValue(update?.reason) ?? "Message error");
+        callbacks.onToolEnd(toolId, true);
+    }
+}
+function objectValue(value) {
+    return typeof value === "object" && value !== null ? value : null;
+}
+function stringValue(value) {
+    return typeof value === "string" && value.trim() ? value : null;
+}
+function numberValue(value) {
+    return typeof value === "number" && Number.isFinite(value) ? value : null;
+}
+function extractContentText(container) {
+    const content = container?.content;
+    if (typeof content === "string") {
+        return content;
+    }
+    if (!Array.isArray(content)) {
+        return "";
+    }
+    return content
+        .map((entry) => {
+        const block = objectValue(entry);
+        return stringValue(block?.text) ?? "";
+    })
+        .join("");
+}
+function computeTextDelta(previousText, nextText) {
+    return nextText.startsWith(previousText) ? nextText.slice(previousText.length) : nextText;
+}
+function splitPiModelSlug(slug) {
+    const separatorIndex = slug.indexOf("/");
+    if (separatorIndex === -1) {
+        return { provider: "openai-codex", modelId: slug };
+    }
+    return {
+        provider: slug.slice(0, separatorIndex),
+        modelId: slug.slice(separatorIndex + 1),
+    };
+}
+function normalizePiThinking(value) {
+    if (PI_THINKING_LEVELS.includes(value)) {
+        return value;
+    }
+    throw new Error(`Unsupported Pi thinking level: ${value}`);
+}
+function parseCompactTokenCount(value) {
+    if (!value) {
+        return undefined;
+    }
+    const match = value.match(/^(\d+(?:\.\d+)?)([KMB])?$/i);
+    if (!match) {
+        return undefined;
+    }
+    const number = Number(match[1]);
+    const unit = match[2]?.toUpperCase();
+    const multiplier = unit === "M" ? 1_000_000 : unit === "K" ? 1_000 : unit === "B" ? 1_000_000_000 : 1;
+    return Math.round(number * multiplier);
+}
+function mimeTypeForImage(filePath) {
+    switch (path.extname(filePath).toLowerCase()) {
+        case ".png":
+            return "image/png";
+        case ".webp":
+            return "image/webp";
+        case ".gif":
+            return "image/gif";
+        default:
+            return "image/jpeg";
+    }
+}
+function shellQuote(value) {
+    return `'${value.replace(/'/g, `'\\''`)}'`;
+}