@averyyy/pi-server 0.85.1-piclient.2 → 0.85.1-piclient.4

package/dist/server.js

@@ -15,6 +15,10 @@ const packageMetadata = createRequire(import.meta.url)("../package.json");
 const PI_SERVER_VERSION = packageMetadata.version;
 function createRequestModels(model, options) {
     const models = createModels();
+    const requestStream = (requestModel, context, streamOptions) => streamSimple(requestModel, context, {
+        ...streamOptions,
+        ...(options.timeoutMs !== undefined ? { timeoutMs: options.timeoutMs } : {}),
+    });
     models.setProvider(createProvider({
         id: model.provider,
         name: model.provider,
@@ -26,8 +30,8 @@ function createRequestModels(model, options) {
             },
         },
         api: {
-            stream: (requestModel, context, streamOptions) => streamSimple(requestModel, context, streamOptions),
-            streamSimple: (requestModel, context, streamOptions) => streamSimple(requestModel, context, streamOptions),
+            stream: requestStream,
+            streamSimple: requestStream,
         },
     }));
     return models;
@@ -91,27 +95,64 @@ function runKey(sessionId, runId) {
 function getStreamRun(sessionId, runId) {
     return streamRuns.get(runKey(sessionId, runId));
 }
-function startStreamRun(sessionId, runId) {
-    const existing = getStreamRun(sessionId, runId);
-    if (existing?.status === "completed")
-        return existing;
-    const now = Date.now();
-    const run = existing ?? {
+function createRunRecord(sessionId, runId, kind) {
+    let resolveSettled;
+    const settled = new Promise((resolve) => {
+        resolveSettled = resolve;
+    });
+    return {
         sessionId,
         runId,
+        kind,
         status: "running",
-        createdAt: now,
-        updatedAt: now,
+        createdAt: Date.now(),
+        updatedAt: Date.now(),
+        cancelRequested: false,
+        controller: new AbortController(),
+        settled,
+        resolveSettled,
     };
+}
+function settleRun(run) {
+    if (!run)
+        return;
+    const resolve = run.resolveSettled;
+    run.resolveSettled = undefined;
+    resolve?.();
+}
+function startStreamRun(sessionId, runId, kind = "stream") {
+    const existing = getStreamRun(sessionId, runId);
+    if (existing?.status === "completed" || existing?.status === "aborted")
+        return existing;
+    const run = existing ?? createRunRecord(sessionId, runId, kind);
     if (existing) {
         run.message = undefined;
         run.errorMessage = undefined;
+        run.cancelRequested = false;
+        run.controller = new AbortController();
+        let resolveSettled;
+        run.settled = new Promise((resolve) => {
+            resolveSettled = resolve;
+        });
+        run.resolveSettled = resolveSettled;
     }
+    run.kind = kind;
     run.status = "running";
-    run.updatedAt = now;
+    run.updatedAt = Date.now();
     streamRuns.set(runKey(sessionId, runId), run);
     return run;
 }
+function streamRunResponseBody(run) {
+    return {
+        sessionId: run.sessionId,
+        runId: run.runId,
+        status: run.status,
+        createdAt: run.createdAt,
+        updatedAt: run.updatedAt,
+        ...(run.message ? { message: run.message } : {}),
+        ...(run.errorMessage ? { errorMessage: run.errorMessage } : {}),
+    };
+}
 function replayStreamRun(run) {
     if (!run.message)
         return [];
@@ -149,6 +190,13 @@ function replayStreamRun(run) {
 function completeStreamRun(run, message) {
     if (!run)
         return;
+    if (run.cancelRequested) {
+        run.status = "aborted";
+        run.message = undefined;
+        run.errorMessage = "Stream run aborted";
+        run.updatedAt = Date.now();
+        return;
+    }
     run.status = "completed";
     run.message = message;
     run.errorMessage = undefined;
@@ -157,6 +205,12 @@ function completeStreamRun(run, message) {
 function failStreamRun(run, errorMessage) {
     if (!run)
         return;
+    if (run.cancelRequested) {
+        run.status = "aborted";
+        run.errorMessage = "Stream run aborted";
+        run.updatedAt = Date.now();
+        return;
+    }
     run.status = "failed";
     run.errorMessage = errorMessage;
     run.updatedAt = Date.now();
@@ -377,17 +431,51 @@ function prepareSessionCompact(body) {
     const options = body.options ?? {};
     return {
         session,
+        sessionId: session.sessionId,
+        revision: session.revision,
+        treeHash: session.treeHash,
+        leafId: session.leafId,
+        entryCount: session.entries.length,
         preparation: preparationResult.value,
         options,
     };
 }
-async function completeSessionCompact(config, body, prepared) {
-    const result = await compactLegacy(prepared.preparation, createRequestModels(body.model, prepared.options), body.model, body.customInstructions, undefined, prepared.options.reasoning);
+function sessionCompactConflict(body, prepared, currentSession) {
+    return {
+        status: 409,
+        body: {
+            error: "Session changed while compaction was running",
+            code: PiServerErrorCode.SESSION_STATE_CONFLICT,
+            details: {
+                sessionId: body.sessionId,
+                expectedRevision: prepared.revision,
+                expectedTreeHash: prepared.treeHash,
+                expectedLeafId: prepared.leafId,
+                actualRevision: currentSession?.revision ?? null,
+                actualTreeHash: currentSession?.treeHash ?? null,
+                actualLeafId: currentSession?.leafId ?? null,
+            },
+        },
+    };
+}
+async function completeSessionCompact(config, body, prepared, run) {
+    const result = await compactLegacy(prepared.preparation, createRequestModels(body.model, prepared.options), body.model, body.customInstructions, run?.controller.signal, prepared.options.reasoning);
     if (!result.ok) {
         return { status: 500, body: { error: result.error.message, code: PiServerErrorCode.INTERNAL_ERROR } };
     }
-    const baseTreeHash = prepared.session.treeHash;
-    const baseEntryCount = prepared.session.entries.length;
+    if (run?.cancelRequested || run?.controller.signal.aborted) {
+        return { status: 409, body: { error: "Compaction aborted", code: PiServerErrorCode.INVALID_REQUEST } };
+    }
+    const currentSession = getSession(prepared.sessionId);
+    if (!currentSession ||
+        currentSession !== prepared.session ||
+        currentSession.revision !== prepared.revision ||
+        currentSession.treeHash !== prepared.treeHash ||
+        currentSession.leafId !== prepared.leafId) {
+        return sessionCompactConflict(body, prepared, currentSession);
+    }
+    const baseTreeHash = prepared.treeHash;
+    const baseEntryCount = prepared.entryCount;
     const compaction = result.value;
     const { session: updatedSession, entry: compactionEntry } = appendCompactionEntry(body.sessionId, compaction);
     persistSession(config, updatedSession);
@@ -427,7 +515,7 @@ function writeServerSentEvent(res, event, body) {
     res.write(`event: ${event}\n`);
     res.write(`data: ${JSON.stringify(body)}\n\n`);
 }
-async function handleSessionCompactStream(config, body, prepared, res) {
+async function handleSessionCompactStream(config, body, prepared, res, run) {
     res.writeHead(200, {
         "Content-Type": "text/event-stream",
         "Cache-Control": "no-cache",
@@ -436,25 +524,33 @@ async function handleSessionCompactStream(config, body, prepared, res) {
     res.flushHeaders();
     res.write(STREAM_HEARTBEAT);
     const heartbeat = setInterval(() => {
+        touchStreamRun(run);
         if (!res.writableEnded) {
             res.write(STREAM_HEARTBEAT);
         }
     }, STREAM_HEARTBEAT_INTERVAL_MS);
     heartbeat.unref();
     try {
-        const result = await completeSessionCompact(config, body, prepared);
+        const result = await completeSessionCompact(config, body, prepared, run);
         writeServerSentEvent(res, result.status >= 400 ? "error" : "result", result.body);
+        finishCompactRun(run, result);
     }
     catch (err) {
         const message = err instanceof Error ? err.message : String(err);
+        if (run) {
+            run.status = run.cancelRequested ? "aborted" : "failed";
+            run.errorMessage = message;
+            run.updatedAt = Date.now();
+        }
         writeServerSentEvent(res, "error", { error: message });
     }
     finally {
         clearInterval(heartbeat);
+        settleRun(run);
         res.end();
     }
 }
-async function handleSessionCompactJsonStream(config, body, prepared, res) {
+async function handleSessionCompactJsonStream(config, body, prepared, res, run) {
     res.writeHead(200, {
         "Content-Type": "application/json",
         "Cache-Control": "no-cache",
@@ -463,35 +559,71 @@ async function handleSessionCompactJsonStream(config, body, prepared, res) {
     res.flushHeaders();
     res.write(JSON_HEARTBEAT);
     const heartbeat = setInterval(() => {
+        touchStreamRun(run);
         if (!res.writableEnded) {
             res.write(JSON_HEARTBEAT);
         }
     }, STREAM_HEARTBEAT_INTERVAL_MS);
     heartbeat.unref();
     try {
-        const result = await completeSessionCompact(config, body, prepared);
+        const result = await completeSessionCompact(config, body, prepared, run);
         res.write(JSON.stringify(result.body));
+        finishCompactRun(run, result);
     }
     catch (err) {
         const message = err instanceof Error ? err.message : String(err);
+        if (run) {
+            run.status = run.cancelRequested ? "aborted" : "failed";
+            run.errorMessage = message;
+            run.updatedAt = Date.now();
+        }
         res.write(JSON.stringify({ error: message }));
     }
     finally {
         clearInterval(heartbeat);
+        settleRun(run);
         res.end();
     }
 }
+function finishCompactRun(run, result) {
+    if (!run)
+        return;
+    if (run.cancelRequested || run.controller.signal.aborted) {
+        run.status = "aborted";
+        run.errorMessage = "Compaction aborted";
+    }
+    else {
+        run.status = result.status >= 400 ? "failed" : "completed";
+        if (result.status >= 400 && "error" in result.body) {
+            run.errorMessage = result.body.error;
+        }
+    }
+    run.updatedAt = Date.now();
+}
 async function handleSessionCompact(config, body, res) {
     const prepared = prepareSessionCompact(body);
     if ("status" in prepared) {
         sendJson(res, prepared.status, prepared.body);
         return;
     }
+    let run;
+    if (body.runId) {
+        const existingRun = getStreamRun(body.sessionId, body.runId);
+        if (existingRun?.status === "aborted") {
+            sendJson(res, 409, { error: "Compaction run was aborted", code: PiServerErrorCode.INVALID_REQUEST });
+            return;
+        }
+        if (existingRun?.status === "running") {
+            sendError(res, 409, "A run with this runId is already in progress", PiServerErrorCode.RUN_IN_PROGRESS);
+            return;
+        }
+        run = startStreamRun(body.sessionId, body.runId, "compact");
+    }
     if (body.streamResponse) {
-        await handleSessionCompactStream(config, body, prepared, res);
+        await handleSessionCompactStream(config, body, prepared, res, run);
         return;
     }
-    await handleSessionCompactJsonStream(config, body, prepared, res);
+    await handleSessionCompactJsonStream(config, body, prepared, res, run);
 }
 function handleSessionHistory(sessionId, from, entriesFrom, revision, baseTreeHash, res) {
     const session = getSession(sessionId);
@@ -515,7 +647,35 @@ function handleSessionRun(sessionId, runId, res) {
         sendError(res, 404, "run not found", PiServerErrorCode.RUN_NOT_FOUND);
         return;
     }
-    sendJson(res, 200, run);
+    sendJson(res, 200, streamRunResponseBody(run));
+}
+async function handleSessionRunAbort(sessionId, runId, res) {
+    if (!sessionId || !runId) {
+        sendError(res, 400, "sessionId and runId are required", PiServerErrorCode.REQUIRED_FIELD_MISSING);
+        return;
+    }
+    let run = getStreamRun(sessionId, runId);
+    if (!run) {
+        run = createRunRecord(sessionId, runId, "stream");
+        run.status = "aborted";
+        run.cancelRequested = true;
+        run.errorMessage = "Stream run aborted before it started";
+        run.updatedAt = Date.now();
+        settleRun(run);
+        streamRuns.set(runKey(sessionId, runId), run);
+        sendJson(res, 200, streamRunResponseBody(run));
+        return;
+    }
+    if (run.status !== "running") {
+        await run.settled;
+        sendJson(res, 200, streamRunResponseBody(run));
+        return;
+    }
+    run.cancelRequested = true;
+    run.updatedAt = Date.now();
+    run.controller.abort();
+    await run.settled;
+    sendJson(res, 200, streamRunResponseBody(run));
 }
 export function buildStreamContext(session, body) {
     const messages = body.contextOverlay ?? [...session.messages, ...(body.ephemeralMessages ?? [])];
@@ -555,7 +715,6 @@ function handleStream(config, body, res) {
     const context = buildStreamContext(session, body);
     const existingRun = body.runId ? getStreamRun(body.sessionId, body.runId) : undefined;
     const resolvedModel = body.model;
-    const streamOptions = { ...(body.options ?? {}) };
     if (existingRun?.status === "completed") {
         res.writeHead(200, {
             "Content-Type": "text/event-stream",
@@ -570,10 +729,19 @@ function handleStream(config, body, res) {
         endStreamResponse(res);
         return;
     }
+    if (existingRun?.status === "aborted") {
+        sendError(res, 409, "The stream run was aborted", PiServerErrorCode.INVALID_REQUEST);
+        return;
+    }
     if (existingRun?.status === "running") {
         sendError(res, 409, "A stream with this runId is already in progress", PiServerErrorCode.RUN_IN_PROGRESS);
         return;
     }
+    const run = body.runId ? startStreamRun(body.sessionId, body.runId) : undefined;
+    const streamOptions = {
+        ...(body.options ?? {}),
+        ...(run ? { signal: run.controller.signal } : {}),
+    };
     res.writeHead(200, {
         "Content-Type": "text/event-stream",
         "Cache-Control": "no-cache",
@@ -581,7 +749,6 @@ function handleStream(config, body, res) {
     });
     res.flushHeaders();
     res.write(STREAM_HEARTBEAT);
-    const run = body.runId ? startStreamRun(body.sessionId, body.runId) : undefined;
     const heartbeat = setInterval(() => {
         touchStreamRun(run);
         if (!res.writableEnded && !res.destroyed) {
@@ -597,11 +764,12 @@ function handleStream(config, body, res) {
         clearInterval(heartbeat);
         const message = err instanceof Error ? err.message : String(err);
         failStreamRun(run, message);
+        settleRun(run);
         writeStreamError(res, message);
         endStreamResponse(res);
         return;
     }
-    (async () => {
+    void (async () => {
         try {
             for await (const event of stream) {
                 const proxyEvent = toProxyEvent(event);
@@ -617,17 +785,20 @@ function handleStream(config, body, res) {
                 }
             }
         }
+        catch (err) {
+            const message = err instanceof Error ? err.message : String(err);
+            failStreamRun(run, message);
+            writeStreamError(res, message);
+        }
         finally {
             clearInterval(heartbeat);
+            if (run?.status === "running") {
+                failStreamRun(run, run.cancelRequested ? "Stream run aborted" : "Stream ended before a terminal event");
+            }
+            settleRun(run);
+            endStreamResponse(res);
         }
-        endStreamResponse(res);
-    })().catch((err) => {
-        clearInterval(heartbeat);
-        const message = err instanceof Error ? err.message : String(err);
-        failStreamRun(run, message);
-        writeStreamError(res, message);
-        endStreamResponse(res);
-    });
+    })();
 }
 async function handlePostRequest(config, pathname, body, res) {
     if (pathname === "/api/receive") {
@@ -683,7 +854,7 @@ export function createPiServer(configOverride) {
     const config = loadConfig(configOverride);
     loadPersistedSessions(config.sessionStoreDir);
     const server = createServer(async (req, res) => {
-        const url = new URL(req.url ?? "/", `http://${config.host}:${config.port}`);
+        const url = new URL(req.url ?? "/", "http://localhost");
         if (req.method === "GET" && url.pathname === "/health") {
             sendJson(res, 200, { status: "ok" });
             return;
@@ -704,6 +875,12 @@ export function createPiServer(configOverride) {
             handleSessionRun(decodeURIComponent(runMatch[1]), decodeURIComponent(runMatch[2]), res);
             return;
         }
+        const runAbortMatch = /^\/api\/session\/([^/]+)\/runs\/([^/]+)\/abort$/.exec(url.pathname);
+        if (req.method === "POST" && runAbortMatch) {
+            await readBody(req);
+            await handleSessionRunAbort(decodeURIComponent(runAbortMatch[1]), decodeURIComponent(runAbortMatch[2]), res);
+            return;
+        }
         if (req.method === "GET" && url.pathname.startsWith("/api/session/") && url.pathname.endsWith("/history")) {
             const encodedSessionId = url.pathname.slice("/api/session/".length, -"/history".length);
             const fromParam = url.searchParams.get("from");
@@ -835,6 +1012,7 @@ function toProxyEvent(event) {
                 reason: event.reason,
                 usage: event.message.usage,
                 deferred: event.message.deferred,
+                message: event.message,
             };
         case "error":
             return {
@@ -851,7 +1029,10 @@ export function startServer(configOverride) {
     const config = loadConfig(configOverride);
     const server = createPiServer(configOverride);
     server.listen(config.port, config.host, () => {
-        console.log(`pi-server listening on ${config.host}:${config.port}`);
+        const address = server.address();
+        const port = typeof address === "object" && address !== null ? address.port : config.port;
+        const host = config.host.includes(":") && !config.host.startsWith("[") ? `[${config.host}]` : config.host;
+        console.log(`pi-server listening on ${host}:${port}`);
     });
     return server;
 }