@averyyy/pi-coding-agent 0.99.1-piclient.1 → 1.0.0-piclient.2

package/dist/core/pi-server-client.js

@@ -1,6 +1,7 @@
 import { randomUUID } from "node:crypto";
 import { setTimeout as sleep } from "node:timers/promises";
-import { EventStream, parseStreamingJson, } from "@earendil-works/pi-ai";
+import { EventStream, getModelType, parseStreamingJson, } from "@earendil-works/pi-ai";
+import { raceWithAbortSignal } from "../utils/abort.js";
 import { buildPiServerTreePrefixHashes, hashPiServerSessionEntries, hashPiServerStaticContext, } from "./pi-server-protocol.js";
 import { ChunkRequest } from "./pi-server-request.js";
 class PiServerEventStream extends EventStream {
@@ -294,10 +295,6 @@ function getServerSentEventFieldValue(line, field) {
     const value = line.slice(prefix.length);
     return value.startsWith(" ") ? value.slice(1) : value;
 }
-function parseServerSentEvents(bodyText) {
-    const parser = new ServerSentEventParser();
-    return [...parser.feed(bodyText), ...parser.finish()];
-}
 function parseServerSentEventData(event, errorPrefix) {
     try {
         return JSON.parse(event.data);
@@ -310,17 +307,44 @@ function formatServerSentEventError(response, payload) {
     const message = isObject(payload) && typeof payload.error === "string" ? payload.error : JSON.stringify(payload);
     return `${getResponseStatus(response)}; content-type: ${getResponseContentType(response)}; body excerpt: ${getBodyExcerpt(message)}`;
 }
-async function readPiServerEventStreamJson(response, errorPrefix) {
-    const bodyText = await response.text();
-    for (const event of parseServerSentEvents(bodyText)) {
-        if (event.event === "error") {
-            throw new Error(`${errorPrefix} (${formatServerSentEventError(response, parseServerSentEventData(event, errorPrefix))})`);
+async function readPiServerEventStreamJson(response, errorPrefix, onCallback) {
+    if (!response.body)
+        throw new Error(`${errorPrefix} (missing event stream body)`);
+    const reader = response.body.getReader();
+    const decoder = new TextDecoder();
+    const parser = new ServerSentEventParser();
+    const processEvents = async (events) => {
+        for (const event of events) {
+            if (event.event === "provider_callback") {
+                if (!onCallback)
+                    throw new Error(`${errorPrefix} (unexpected provider callback)`);
+                await onCallback(parseServerSentEventData(event, errorPrefix));
+            }
+            else if (event.event === "error") {
+                throw new Error(`${errorPrefix} (${formatServerSentEventError(response, parseServerSentEventData(event, errorPrefix))})`);
+            }
+            else if (event.event === "result") {
+                return { value: parseServerSentEventData(event, errorPrefix) };
+            }
         }
-        if (event.event === "result") {
-            return parseServerSentEventData(event, errorPrefix);
+        return undefined;
+    };
+    try {
+        for (;;) {
+            const chunk = await reader.read();
+            const result = await processEvents(chunk.done
+                ? [...parser.feed(decoder.decode()), ...parser.finish()]
+                : parser.feed(decoder.decode(chunk.value, { stream: true })));
+            if (result)
+                return result.value;
+            if (chunk.done)
+                break;
         }
+        throw new Error(`${errorPrefix} (expected result event)`);
+    }
+    finally {
+        await reader.cancel().catch(() => { });
     }
-    throw new Error(`${errorPrefix} (${formatResponseDetails(response, bodyText)}; expected result event)`);
 }
 async function readPiServerCompactResponse(response, errorPrefix) {
     const contentType = getResponseContentType(response).toLowerCase().split(";")[0]?.trim();
@@ -607,6 +631,10 @@ function serializeOptions(options) {
     return {
         temperature: options?.temperature,
         maxTokens: options?.maxTokens,
+        samplingParams: options?.samplingParams,
+        toolChoice: options?.toolChoice,
+        deferred: options?.deferred,
+        env: options?.env,
         reasoning: options?.reasoning,
         cacheRetention: options?.cacheRetention,
         sessionId: options?.sessionId,
@@ -621,6 +649,35 @@ function serializeOptions(options) {
         maxRetryDelayMs: options?.maxRetryDelayMs,
     };
 }
+export function streamRawPiServer(model, context, options) {
+    if (options?.transformPreparedStreamOptions) {
+        throw new Error("Prepared stream option transforms cannot execute on pi-server");
+    }
+    return streamPiServer(model, context, {
+        signal: options?.signal,
+        sessionId: options?.sessionId,
+        onProviderStreamEvent: options?.onProviderStreamEvent,
+        onPayload: options?.onPayload,
+        onResponse: options?.onResponse,
+        fetch: options?.fetch,
+        requestKind: "api",
+        apiOptions: options ?? {},
+    });
+}
+export function fetchDeferredPiServer(model, handle, options) {
+    return streamPiServer(model, { messages: [] }, {
+        ...options,
+        requestKind: "deferred",
+        deferredHandle: handle,
+        apiOptions: options,
+    });
+}
+export async function cancelDeferredPiServer(model, handle, options) {
+    assertNoCustomPiServerFetch(options);
+    const result = await postPiServerModelOperation("/api/cancel-deferred", { model, handle, options: serializeApiOptions(options ?? {}) }, model, options);
+    if (result.cancelled !== true || result.stopReason !== "stop")
+        throw new Error("pi-server deferred cancellation did not complete");
+}
 export async function compactPiServer(model, context, options) {
     const sessionId = options?.sessionId ?? "default";
     const runId = randomUUID();
@@ -712,6 +769,8 @@ export async function compactPiServer(model, context, options) {
     }
 }
 export async function streamPiServer(model, context, options) {
+    assertNoCustomPiServerFetch(options);
+    assertNoCustomPiServerFetch(options?.apiOptions);
     const sessionId = options?.sessionId ?? randomUUID();
     const runId = randomUUID();
     const isEphemeralSession = options?.sessionId === undefined;
@@ -737,6 +796,28 @@ export async function streamPiServer(model, context, options) {
     (async () => {
         let phase = "session_init";
         let streamOpened = false;
+        let observerFailed = false;
+        let callbackFailed = false;
+        let cacheRetention;
+        const requireCacheRetention = options?.requireCacheRetention || options?.onCacheRetentionResolved !== undefined;
+        const applyCachePolicy = async (retention) => {
+            try {
+                if (retention !== "none" && retention !== "short" && retention !== "long") {
+                    throw new Error("pi-server cache policy has an invalid cacheRetention");
+                }
+                if (cacheRetention !== undefined) {
+                    if (cacheRetention !== retention)
+                        throw new Error("pi-server cache policy changed during the same run");
+                    return;
+                }
+                cacheRetention = retention;
+                await raceWithAbortSignal(Promise.resolve().then(() => options?.onCacheRetentionResolved?.(retention)), options?.signal);
+            }
+            catch (error) {
+                observerFailed = true;
+                throw error;
+            }
+        };
         try {
             const request = createPiServerRequest(options?.signal);
             await ensureSessionInit(sessionId, context, request);
@@ -744,7 +825,13 @@ export async function streamPiServer(model, context, options) {
                 sessionId,
                 runId,
                 model,
-                options: serializeOptions(options),
+                options: options?.requestKind === "api" || options?.requestKind === "deferred"
+                    ? serializeApiOptions(options.apiOptions ?? {})
+                    : serializeOptions(options),
+                requestKind: options?.requestKind,
+                deferredHandle: options?.deferredHandle,
+                observeProviderStreamEvents: options?.onProviderStreamEvent !== undefined,
+                callbacks: { onPayload: options?.onPayload !== undefined, onResponse: options?.onResponse !== undefined },
                 contextOverlay: options?.contextOverlay ?? [
                     ...context.messages,
                     ...(options?.ephemeralMessages ?? []),
@@ -775,9 +862,44 @@ export async function streamPiServer(model, context, options) {
             const decoder = new TextDecoder();
             const parser = new ServerSentEventParser();
             let terminalReceived = false;
-            const consumeEvents = (events) => {
+            const consumeEvents = async (events) => {
                 for (const serverEvent of events) {
+                    if (serverEvent.event === "cache_policy") {
+                        try {
+                            const policy = parseServerSentEventData(serverEvent, "pi-server cache policy");
+                            await applyCachePolicy(isObject(policy) ? policy.cacheRetention : undefined);
+                        }
+                        catch (error) {
+                            observerFailed = true;
+                            throw error;
+                        }
+                        continue;
+                    }
+                    if (serverEvent.event === "provider_callback") {
+                        try {
+                            await handleProviderCallback(parseServerSentEventData(serverEvent, "pi-server callback"), model, options, sessionId, runId, request);
+                        }
+                        catch (error) {
+                            callbackFailed = true;
+                            throw error;
+                        }
+                        continue;
+                    }
                     const proxyEvent = parseServerSentEventData(serverEvent, "pi-server stream");
+                    if (requireCacheRetention && cacheRetention === undefined && proxyEvent.type !== "error") {
+                        observerFailed = true;
+                        throw new Error("pi-server stream is missing its resolved cache policy");
+                    }
+                    if (proxyEvent.type === "provider_stream_event") {
+                        try {
+                            await options?.onProviderStreamEvent?.(proxyEvent.data, proxyEvent.model);
+                        }
+                        catch (error) {
+                            observerFailed = true;
+                            throw error;
+                        }
+                        continue;
+                    }
                     const event = processProxyEvent(proxyEvent, partial);
                     if (!event)
                         continue;
@@ -793,10 +915,10 @@ export async function streamPiServer(model, context, options) {
                 if (options?.signal?.aborted) {
                     throw new Error("Request aborted by user");
                 }
-                consumeEvents(parser.feed(decoder.decode(value, { stream: true })));
+                await consumeEvents(parser.feed(decoder.decode(value, { stream: true })));
             }
-            consumeEvents(parser.feed(decoder.decode()));
-            consumeEvents(parser.finish());
+            await consumeEvents(parser.feed(decoder.decode()));
+            await consumeEvents(parser.finish());
             if (options?.signal?.aborted) {
                 throw new Error("Request aborted by user");
             }
@@ -814,9 +936,17 @@ export async function streamPiServer(model, context, options) {
             // failures before the stream starts (502/504/fetch failed) never opened a run
             // we can poll; attempting recovery against a down proxy reclassifies the error
             // as history_reconcile and blocks session auto-retry.
-            if (!options?.signal?.aborted && phase === "provider_stream" && streamOpened) {
+            if (!options?.signal?.aborted &&
+                phase === "provider_stream" &&
+                streamOpened &&
+                !observerFailed &&
+                !callbackFailed) {
                 try {
                     const recoveredRun = await waitForPiServerRunCompletion(sessionId, runId, createPiServerRequest(options?.signal), options);
+                    if (recoveredRun?.cacheRetention !== undefined ||
+                        (requireCacheRetention && recoveredRun?.status === "completed")) {
+                        await applyCachePolicy(recoveredRun?.cacheRetention);
+                    }
                     if (recoveredRun?.status === "completed" && recoveredRun.message) {
                         switch (recoveredRun.message.stopReason) {
                             case "stop":
@@ -840,6 +970,11 @@ export async function streamPiServer(model, context, options) {
                         return;
                     }
                     if (recoveredRun?.status === "failed" && recoveredRun.errorMessage) {
+                        if (recoveredRun.message) {
+                            stream.push({ type: "error", reason: "error", error: recoveredRun.message });
+                            stream.end();
+                            return;
+                        }
                         errorMessage = recoveredRun.errorMessage;
                     }
                     if (recoveredRun?.status === "aborted") {
@@ -851,7 +986,9 @@ export async function streamPiServer(model, context, options) {
                     errorMessage = recoveryError instanceof Error ? recoveryError.message : String(recoveryError);
                 }
             }
-            if (options?.signal?.aborted && operation && !operation.cancelConfirmed) {
+            if ((options?.signal?.aborted || callbackFailed || observerFailed) &&
+                operation &&
+                !operation.cancelConfirmed) {
                 try {
                     await cancelPiServerOperation(operation);
                 }
@@ -871,7 +1008,7 @@ export async function streamPiServer(model, context, options) {
                     details: {
                         phase,
                         source: "pi-server",
-                        retryable: isRetryablePiServerFailure(phase, error),
+                        retryable: !callbackFailed && !observerFailed && isRetryablePiServerFailure(phase, error),
                     },
                 },
             ];
@@ -884,7 +1021,7 @@ export async function streamPiServer(model, context, options) {
         }
         finally {
             try {
-                await finishPiServerOperation(operation, options?.signal?.aborted === true);
+                await finishPiServerOperation(operation, options?.signal?.aborted === true || callbackFailed || observerFailed);
             }
             catch {
                 // The owning session's cancellation helper reports the failure and keeps
@@ -896,6 +1033,8 @@ export async function streamPiServer(model, context, options) {
 }
 function processProxyEvent(proxyEvent, partial) {
     switch (proxyEvent.type) {
+        case "provider_stream_event":
+            return undefined;
         case "start":
             return { type: "start", partial };
         case "text_start":
@@ -972,6 +1111,8 @@ function processProxyEvent(proxyEvent, partial) {
             partial.deferred = proxyEvent.deferred;
             return { type: "done", reason: proxyEvent.reason, message: partial };
         case "error":
+            if (proxyEvent.message)
+                return { type: "error", reason: proxyEvent.reason, error: proxyEvent.message };
             partial.stopReason = proxyEvent.reason;
             partial.errorMessage = proxyEvent.errorMessage;
             partial.usage = proxyEvent.usage;
@@ -984,4 +1125,149 @@ function processProxyEvent(proxyEvent, partial) {
     }
 }
 export { getServerUrl, getAuthToken };
+function serializeApiOptions(options) {
+    const { signal: _signal, fetch: _fetch, onPayload: _onPayload, onResponse: _onResponse, onProviderStreamEvent: _onProviderStreamEvent, telemetryContext: _telemetryContext, ...serialized } = options;
+    return serialized;
+}
+export async function fetchPiServerModels(options) {
+    const response = await createPiServerRequest(options?.signal).getJson("/api/models");
+    return readPiServerJson(response, "pi-server model catalog failed");
+}
+export async function generateImagesPiServer(model, context, options) {
+    assertNoCustomPiServerFetch(options);
+    return postPiServerModelOperation("/api/generate-images", {
+        model,
+        context,
+        options: {
+            apiKey: options?.apiKey,
+            env: options?.env,
+            headers: options?.headers,
+            metadata: options?.metadata,
+            timeoutMs: options?.timeoutMs,
+            maxRetries: options?.maxRetries,
+            maxRetryDelayMs: options?.maxRetryDelayMs,
+        },
+    }, model, options);
+}
+export async function classifyPiServer(model, context, options) {
+    assertNoCustomPiServerFetch(options);
+    return postPiServerModelOperation("/api/classify", {
+        model,
+        context,
+        options: {
+            apiKey: options?.apiKey,
+            env: options?.env,
+            headers: options?.headers,
+            temperature: options?.temperature,
+            timeoutMs: options?.timeoutMs,
+            maxRetries: options?.maxRetries,
+            maxRetryDelayMs: options?.maxRetryDelayMs,
+        },
+    }, model, options);
+}
+function assertNoCustomPiServerFetch(options) {
+    if (options?.fetch)
+        throw new Error("Custom fetch functions cannot execute on pi-server");
+}
+async function handleProviderCallback(callback, model, options, sessionId, runId, request) {
+    if (!callback ||
+        typeof callback.callbackId !== "string" ||
+        !callback.model ||
+        callback.model.provider !== model.provider ||
+        callback.model.id !== model.id ||
+        callback.model.api !== model.api ||
+        getModelType(callback.model) !== getModelType(model)) {
+        throw new Error("pi-server provider callback has an invalid model or callback identity");
+    }
+    const remoteModel = callback.model;
+    let result;
+    let failure;
+    let failed = false;
+    try {
+        if (callback.kind === "payload") {
+            if (!options?.onPayload)
+                throw new Error("Unexpected pi-server payload callback");
+            const replacement = await raceWithAbortSignal(Promise.resolve().then(() => options.onPayload(callback.payload, remoteModel)), options.signal);
+            result = replacement === undefined ? callback.payload : replacement;
+        }
+        else if (callback.kind === "response") {
+            if (!options?.onResponse)
+                throw new Error("Unexpected pi-server response callback");
+            await raceWithAbortSignal(Promise.resolve().then(() => options.onResponse(callback.response, remoteModel)), options.signal);
+        }
+        else {
+            throw new Error("Unknown pi-server provider callback kind");
+        }
+    }
+    catch (error) {
+        if (options?.signal?.aborted)
+            throw error;
+        failure = error;
+        failed = true;
+    }
+    const response = await request.postJson("/api/provider-callback", {
+        sessionId,
+        runId,
+        callbackId: callback.callbackId,
+        result,
+        error: !failed
+            ? undefined
+            : {
+                name: failure instanceof Error ? failure.name : "Error",
+                message: failure instanceof Error ? failure.message : String(failure),
+            },
+    });
+    const reply = await readPiServerJson(response, "pi-server callback reply failed");
+    if (!reply.accepted || reply.callbackId !== callback.callbackId || reply.runId !== runId) {
+        throw new Error("pi-server callback reply did not acknowledge the pending callback");
+    }
+    if (failed)
+        throw failure;
+}
+async function postPiServerModelOperation(endpoint, body, model, options) {
+    const signal = options?.signal;
+    const sessionId = randomUUID();
+    const runId = randomUUID();
+    const operation = registerPiServerOperation(sessionId, sessionId, runId, signal);
+    const request = createPiServerRequest(signal);
+    let callbackFailed = false;
+    try {
+        const response = await request.postJson(endpoint, {
+            ...body,
+            sessionId,
+            runId,
+            callbacks: { onPayload: options?.onPayload !== undefined, onResponse: options?.onResponse !== undefined },
+        });
+        if (!response.ok)
+            return await readPiServerJson(response, "pi-server model operation failed");
+        await ensurePiServerEventStream(response);
+        try {
+            return await readPiServerEventStreamJson(response, "pi-server model operation failed", async (callback) => {
+                try {
+                    await handleProviderCallback(callback, model, options, sessionId, runId, request);
+                }
+                catch (error) {
+                    callbackFailed = true;
+                    throw error;
+                }
+            });
+        }
+        catch (error) {
+            if (signal?.aborted || callbackFailed)
+                throw error;
+            const recovered = await waitForPiServerRunCompletion(sessionId, runId, request, { signal });
+            if (recovered?.operationResult)
+                return recovered.operationResult;
+            throw error;
+        }
+    }
+    catch (error) {
+        if ((signal?.aborted || callbackFailed) && operation && !operation.cancelConfirmed)
+            await cancelPiServerOperation(operation);
+        throw error;
+    }
+    finally {
+        await finishPiServerOperation(operation, signal?.aborted === true || callbackFailed);
+    }
+}
 //# sourceMappingURL=pi-server-client.js.map